Propagate context to cassandra4 callbacks (#3371)

This commit is contained in:
Lauri Tulmin 2021-06-21 23:12:57 +03:00 committed by GitHub
parent d31a2c7f17
commit 3967c42cda
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 42 additions and 16 deletions

View File

@ -26,7 +26,9 @@ import com.datastax.oss.driver.api.core.type.reflect.GenericType;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.Supplier;
import org.checkerframework.checker.nullness.qual.Nullable;
public class TracingCqlSession implements CqlSession {
@ -188,29 +190,50 @@ public class TracingCqlSession implements CqlSession {
public CompletionStage<AsyncResultSet> executeAsync(Statement<?> statement) {
String query = getQuery(statement);
CassandraRequest request = CassandraRequest.create(session, query);
Context context = instrumenter().start(Context.current(), request);
try (Scope ignored = context.makeCurrent()) {
CompletionStage<AsyncResultSet> stage = session.executeAsync(statement);
return stage.whenComplete(
(asyncResultSet, throwable) ->
instrumenter()
.end(context, request, getExecutionInfo(asyncResultSet, throwable), throwable));
}
return executeAsync(request, () -> session.executeAsync(statement));
}
@Override
public CompletionStage<AsyncResultSet> executeAsync(String query) {
CassandraRequest request = CassandraRequest.create(session, query);
Context context = instrumenter().start(Context.current(), request);
return executeAsync(request, () -> session.executeAsync(query));
}
private static CompletionStage<AsyncResultSet> executeAsync(
CassandraRequest request, Supplier<CompletionStage<AsyncResultSet>> query) {
Context parentContext = Context.current();
Context context = instrumenter().start(parentContext, request);
try (Scope ignored = context.makeCurrent()) {
CompletionStage<AsyncResultSet> stage = session.executeAsync(query);
return stage.whenComplete(
(asyncResultSet, throwable) ->
instrumenter()
.end(context, request, getExecutionInfo(asyncResultSet, throwable), throwable));
CompletionStage<AsyncResultSet> stage = query.get();
return wrap(
stage.whenComplete(
(asyncResultSet, throwable) ->
instrumenter()
.end(
context,
request,
getExecutionInfo(asyncResultSet, throwable),
throwable)),
parentContext);
}
}
static <T> CompletableFuture<T> wrap(CompletionStage<T> future, Context context) {
CompletableFuture<T> result = new CompletableFuture<>();
future.whenComplete(
(T value, Throwable throwable) -> {
try (Scope ignored = context.makeCurrent()) {
if (throwable != null) {
result.completeExceptionally(throwable);
} else {
result.complete(value);
}
}
});
return result;
}
private static String getQuery(Statement<?> statement) {
String query = null;
if (statement instanceof SimpleStatement) {

View File

@ -74,14 +74,17 @@ class CassandraClientTest extends AgentInstrumentationSpecification {
CqlSession session = getSession(keyspace)
runUnderTrace("parent") {
session.executeAsync(statement).toCompletableFuture().get()
session.executeAsync(statement).toCompletableFuture().whenComplete({result, throwable ->
runUnderTrace("child") {}
}) .get()
}
expect:
assertTraces(1) {
trace(0, 2) {
trace(0, 3) {
basicSpan(it, 0, "parent")
cassandraSpan(it, 1, spanName, expectedStatement, operation, keyspace, table, span(0))
basicSpan(it, 2, "child", span(0))
}
}