Make sure span is opened and closed with scope in Cassandra
MODIFY
dd-java-agent/instrumentation/datastax-cassandra-3/src/main/java/datadog/trace/instrumentation/datastax/cassandra/TracingSession.java
dd-java-agent/instrumentation/datastax-cassandra-3/src/main/java/datadog/trace/instrumentation/datastax/cassandra/TracingSession.java
@@ -15,8 +15,10 @@ import com.datastax.driver.core.Statement;
import com.google.common.base.Function;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
+import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
@@ -58,89 +60,101 @@ public class TracingSession implements Session {
@Override
public ResultSet execute(final String query) {
- final Span span = buildSpan(query);
- ResultSet resultSet = null;
+ try (final Scope scope = startSpanWithScope(query, true)) {
try {
- resultSet = session.execute(query);
+ final ResultSet resultSet = session.execute(query);
+ beforeSpanFinish(scope.span(), resultSet);
return resultSet;
- } finally {
- finishSpan(span, resultSet);
+ } catch (final RuntimeException e) {
+ beforeSpanFinish(scope.span(), e);
+ throw e;
+ }
}
}
@Override
public ResultSet execute(final String query, final Object... values) {
- final Span span = buildSpan(query);
- ResultSet resultSet = null;
+ try (final Scope scope = startSpanWithScope(query, true)) {
try {
- resultSet = session.execute(query, values);
+ final ResultSet resultSet = session.execute(query, values);
+ beforeSpanFinish(scope.span(), resultSet);
return resultSet;
- } finally {
- finishSpan(span, resultSet);
+ } catch (final RuntimeException e) {
+ beforeSpanFinish(scope.span(), e);
+ throw e;
+ }
}
}
@Override
public ResultSet execute(final String query, final Map<String, Object> values) {
- final Span span = buildSpan(query);
- ResultSet resultSet = null;
+ try (final Scope scope = startSpanWithScope(query, true)) {
try {
- resultSet = session.execute(query, values);
+ final ResultSet resultSet = session.execute(query, values);
+ beforeSpanFinish(scope.span(), resultSet);
return resultSet;
- } finally {
- finishSpan(span, resultSet);
+ } catch (final RuntimeException e) {
+ beforeSpanFinish(scope.span(), e);
+ throw e;
+ }
}
}
@Override
public ResultSet execute(final Statement statement) {
final String query = getQuery(statement);
- final Span span = buildSpan(query);
- ResultSet resultSet = null;
+ try (final Scope scope = startSpanWithScope(query, true)) {
try {
- resultSet = session.execute(statement);
+ final ResultSet resultSet = session.execute(statement);
+ beforeSpanFinish(scope.span(), resultSet);
return resultSet;
- } finally {
- finishSpan(span, resultSet);
+ } catch (final RuntimeException e) {
+ beforeSpanFinish(scope.span(), e);
+ throw e;
+ }
}
}
@Override
public ResultSetFuture executeAsync(final String query) {
- final Span span = buildSpan(query);
+ try (final Scope scope = startSpanWithScope(query, false)) {
final ResultSetFuture future = session.executeAsync(query);
- future.addListener(createListener(span, future), executorService);
+ future.addListener(createListener(scope.span(), future), executorService);
return future;
}
+ }
@Override
public ResultSetFuture executeAsync(final String query, final Object... values) {
- final Span span = buildSpan(query);
+ try (final Scope scope = startSpanWithScope(query, false)) {
final ResultSetFuture future = session.executeAsync(query, values);
- future.addListener(createListener(span, future), executorService);
+ future.addListener(createListener(scope.span(), future), executorService);
return future;
}
+ }
@Override
public ResultSetFuture executeAsync(final String query, final Map<String, Object> values) {
- final Span span = buildSpan(query);
+ try (final Scope scope = startSpanWithScope(query, false)) {
final ResultSetFuture future = session.executeAsync(query, values);
- future.addListener(createListener(span, future), executorService);
+ future.addListener(createListener(scope.span(), future), executorService);
return future;
}
+ }
@Override
public ResultSetFuture executeAsync(final Statement statement) {
final String query = getQuery(statement);
- final Span span = buildSpan(query);
+ try (final Scope scope = startSpanWithScope(query, false)) {
final ResultSetFuture future = session.executeAsync(statement);
- future.addListener(createListener(span, future), executorService);
+ future.addListener(createListener(scope.span(), future), executorService);
return future;
}
+ }
@Override
public PreparedStatement prepare(final String query) {
@@ -202,32 +216,31 @@ public class TracingSession implements Session {
return new Runnable() {
@Override
public void run() {
- try {
- finishSpan(span, future.get());
+ try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) {
+ beforeSpanFinish(span, future.get());
} catch (final InterruptedException | ExecutionException e) {
- finishSpan(span, e);
+ beforeSpanFinish(span, e);
}
}
};
}
- private Span buildSpan(final String query) {
+ private Scope startSpanWithScope(final String query, final boolean finishOnClose) {
final Span span = tracer.buildSpan("cassandra.execute").start();
+ final Scope scope = tracer.scopeManager().activate(span, finishOnClose);
DECORATE.afterStart(span);
DECORATE.onConnection(span, session);
DECORATE.onStatement(span, query);
- return span;
+ return scope;
}
- private static void finishSpan(final Span span, final ResultSet resultSet) {
+ private static void beforeSpanFinish(final Span span, final ResultSet resultSet) {
DECORATE.onResponse(span, resultSet);
DECORATE.beforeFinish(span);
- span.finish();
}
- private static void finishSpan(final Span span, final Exception e) {
+ private static void beforeSpanFinish(final Span span, final Exception e) {
DECORATE.onError(span, e);
DECORATE.beforeFinish(span);
- span.finish();
}
}

Make sure span is opened and closed with scope in Cassandra