com.datastax.oss.driver.internal.core.cql.QueryTraceFetcher Maven / Gradle / Ivy
The newest version!
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.datastax.oss.driver.internal.core.cql;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.cql.QueryTrace;
import com.datastax.oss.driver.api.core.cql.Row;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.cql.TraceEvent;
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
import com.datastax.oss.driver.shaded.guava.common.collect.Iterables;
import com.datastax.oss.driver.shaded.netty.util.concurrent.EventExecutor;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import net.jcip.annotations.ThreadSafe;
@ThreadSafe
class QueryTraceFetcher {
private final UUID tracingId;
private final CqlSession session;
private final DriverExecutionProfile config;
private final int maxAttempts;
private final long intervalNanos;
private final EventExecutor scheduler;
private final CompletableFuture resultFuture = new CompletableFuture<>();
QueryTraceFetcher(
UUID tracingId,
CqlSession session,
InternalDriverContext context,
DriverExecutionProfile config) {
this.tracingId = tracingId;
this.session = session;
String regularConsistency = config.getString(DefaultDriverOption.REQUEST_CONSISTENCY);
String traceConsistency = config.getString(DefaultDriverOption.REQUEST_TRACE_CONSISTENCY);
this.config =
traceConsistency.equals(regularConsistency)
? config
: config.withString(DefaultDriverOption.REQUEST_CONSISTENCY, traceConsistency);
this.maxAttempts = config.getInt(DefaultDriverOption.REQUEST_TRACE_ATTEMPTS);
this.intervalNanos = config.getDuration(DefaultDriverOption.REQUEST_TRACE_INTERVAL).toNanos();
this.scheduler = context.getNettyOptions().adminEventExecutorGroup().next();
querySession(maxAttempts);
}
CompletionStage fetch() {
return resultFuture;
}
private void querySession(int remainingAttempts) {
session
.executeAsync(
SimpleStatement.builder("SELECT * FROM system_traces.sessions WHERE session_id = ?")
.addPositionalValue(tracingId)
.setExecutionProfile(config)
.build())
.whenComplete(
(rs, error) -> {
if (error != null) {
resultFuture.completeExceptionally(error);
} else {
Row row = rs.one();
if (row == null || row.isNull("duration") || row.isNull("started_at")) {
// Trace is incomplete => fail if last try, or schedule retry
if (remainingAttempts == 1) {
resultFuture.completeExceptionally(
new IllegalStateException(
String.format(
"Trace %s still not complete after %d attempts",
tracingId, maxAttempts)));
} else {
scheduler.schedule(
() -> querySession(remainingAttempts - 1),
intervalNanos,
TimeUnit.NANOSECONDS);
}
} else {
queryEvents(row, new ArrayList<>(), null);
}
}
});
}
private void queryEvents(Row sessionRow, List events, ByteBuffer pagingState) {
session
.executeAsync(
SimpleStatement.builder("SELECT * FROM system_traces.events WHERE session_id = ?")
.addPositionalValue(tracingId)
.setPagingState(pagingState)
.setExecutionProfile(config)
.build())
.whenComplete(
(rs, error) -> {
if (error != null) {
resultFuture.completeExceptionally(error);
} else {
Iterables.addAll(events, rs.currentPage());
ByteBuffer nextPagingState = rs.getExecutionInfo().getPagingState();
if (nextPagingState == null) {
resultFuture.complete(buildTrace(sessionRow, events));
} else {
queryEvents(sessionRow, events, nextPagingState);
}
}
});
}
private QueryTrace buildTrace(Row sessionRow, Iterable eventRows) {
ImmutableList.Builder eventsBuilder = ImmutableList.builder();
for (Row eventRow : eventRows) {
UUID eventId = eventRow.getUuid("event_id");
int sourcePort = 0;
if (eventRow.getColumnDefinitions().contains("source_port")) {
sourcePort = eventRow.getInt("source_port");
}
eventsBuilder.add(
new DefaultTraceEvent(
eventRow.getString("activity"),
eventId == null ? -1 : eventId.timestamp(),
new InetSocketAddress(eventRow.getInetAddress("source"), sourcePort),
eventRow.getInt("source_elapsed"),
eventRow.getString("thread")));
}
Instant startedAt = sessionRow.getInstant("started_at");
int coordinatorPort = 0;
if (sessionRow.getColumnDefinitions().contains("coordinator_port")) {
coordinatorPort = sessionRow.getInt("coordinator_port");
}
return new DefaultQueryTrace(
tracingId,
sessionRow.getString("request"),
sessionRow.getInt("duration"),
new InetSocketAddress(sessionRow.getInetAddress("coordinator"), coordinatorPort),
sessionRow.getMap("parameters", String.class, String.class),
startedAt == null ? -1 : startedAt.toEpochMilli(),
eventsBuilder.build());
}
}