From f3e1eb4f5d085b1702f84cb549fe476dc8fd3e8e Mon Sep 17 00:00:00 2001 From: Nikolay Martynov Date: Thu, 4 Apr 2019 11:28:51 -0400 Subject: [PATCH] Make sure span is opened and closed with scope in Couchbase --- .../CouchbaseBucketInstrumentation.java | 31 +++++++++++-------- .../CouchbaseClusterInstrumentation.java | 29 ++++++++++------- 2 files changed, 35 insertions(+), 25 deletions(-) diff --git a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java index b7bb935dfe..742605e0b8 100644 --- a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java +++ b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java @@ -14,6 +14,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.api.DDTags; import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import io.opentracing.Scope; import io.opentracing.Span; import io.opentracing.noop.NoopSpan; import io.opentracing.util.GlobalTracer; @@ -113,14 +114,16 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default { declaringClass.getSimpleName().replace("CouchbaseAsync", "").replace("DefaultAsync", ""); final String resourceName = className + "." + method.getName(); - // just replace the no-op span. - spanRef.set( - DECORATE.afterStart( - GlobalTracer.get() - .buildSpan("couchbase.call") - .withTag(DDTags.RESOURCE_NAME, resourceName) - .withTag("bucket", bucket) - .start())); + final Span span = + GlobalTracer.get() + .buildSpan("couchbase.call") + .withTag(DDTags.RESOURCE_NAME, resourceName) + .withTag("bucket", bucket) + .start(); + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) { + // just replace the no-op span. + spanRef.set(DECORATE.afterStart(span)); + } } } @@ -136,8 +139,9 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default { final Span span = spanRef.getAndSet(null); if (span != null) { - DECORATE.beforeFinish(span); - span.finish(); + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) { + DECORATE.beforeFinish(span); + } } } } @@ -153,9 +157,10 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default { public void call(final Throwable throwable) { final Span span = spanRef.getAndSet(null); if (span != null) { - DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); - span.finish(); + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + } } } } diff --git a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java index 17fdfa7574..1235c84ed8 100644 --- a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java +++ b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java @@ -14,6 +14,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.api.DDTags; import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import io.opentracing.Scope; import io.opentracing.Span; import io.opentracing.noop.NoopSpan; import io.opentracing.util.GlobalTracer; @@ -108,13 +109,15 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default { declaringClass.getSimpleName().replace("CouchbaseAsync", "").replace("DefaultAsync", ""); final String resourceName = className + "." + method.getName(); - // just replace the no-op span. - spanRef.set( - DECORATE.afterStart( - GlobalTracer.get() - .buildSpan("couchbase.call") - .withTag(DDTags.RESOURCE_NAME, resourceName) - .start())); + final Span span = + GlobalTracer.get() + .buildSpan("couchbase.call") + .withTag(DDTags.RESOURCE_NAME, resourceName) + .start(); + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) { + // just replace the no-op span. + spanRef.set(DECORATE.afterStart(scope.span())); + } } } @@ -130,8 +133,9 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default { final Span span = spanRef.getAndSet(null); if (span != null) { - DECORATE.beforeFinish(span); - span.finish(); + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) { + DECORATE.beforeFinish(span); + } } } } @@ -147,9 +151,10 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default { public void call(final Throwable throwable) { final Span span = spanRef.getAndSet(null); if (span != null) { - DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); - span.finish(); + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + } } } }