Refactor ContinuableScope

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java

@@ -217,13 +217,13 @@ public final class ExecutorInstrumentation extends Instrumenter.Configurable {
     protected final TraceScope.Continuation continuation;
 
     public DatadogWrapper(TraceScope scope) {
-      continuation = scope.capture(true);
+      continuation = scope.capture();
       log.debug("created continuation {} from scope {}", continuation, scope);
     }
 
     public void cancel() {
       if (null != continuation) {
-        continuation.activate().close();
+        continuation.close();
         log.debug("canceled continuation {}", continuation);
       }
     }

dd-trace-api/src/main/java/datadog/trace/context/TraceScope.java

@@ -8,7 +8,7 @@ public interface TraceScope {
    *
    * <p>Should be called on the parent thread.
    */
-  Continuation capture(boolean finishOnClose);
+  Continuation capture();
 
   /** Close the activated context and allow any underlying spans to finish. */
   void close();
@@ -21,5 +21,8 @@ public interface TraceScope {
      * <p>Should be called on the child thread.
      */
     TraceScope activate();
+
+    /** Cancel the continuation. */
+    void close();
   }
 }

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContextualScopeManager.java

@@ -1,5 +1,6 @@
 package datadog.opentracing.scopemanager;
 
+import datadog.opentracing.DDSpan;
 import io.opentracing.Scope;
 import io.opentracing.ScopeManager;
 import io.opentracing.Span;
@@ -17,7 +18,11 @@ public class ContextualScopeManager implements ScopeManager {
         return context.activate(span, finishOnClose);
       }
     }
-    return new ContinuableScope(this, span, finishOnClose);
+    if (span instanceof DDSpan) {
+      return new ContinuableScope(this, (DDSpan) span, finishOnClose);
+    } else {
+      return new SimpleScope(this, span, finishOnClose);
+    }
   }
 
   @Override

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContinuableScope.java

@@ -1,11 +1,10 @@
 package datadog.opentracing.scopemanager;
 
+import datadog.opentracing.DDSpan;
 import datadog.opentracing.DDSpanContext;
 import datadog.opentracing.PendingTrace;
 import datadog.trace.context.TraceScope;
 import io.opentracing.Scope;
-import io.opentracing.Span;
-import io.opentracing.noop.NoopScopeManager;
 import java.io.Closeable;
 import java.lang.ref.WeakReference;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -14,26 +13,38 @@ import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class ContinuableScope implements Scope, TraceScope {
-  final ContextualScopeManager scopeManager;
-  final AtomicInteger refCount;
-  private final Span wrapped;
+  /** ScopeManager holding the thread-local to this scope. */
+  private final ContextualScopeManager scopeManager;
+  /**
+   * Span contained by this scope. Async scopes will hold a reference to the parent scope's span.
+   */
+  private final DDSpan spanUnderScope;
+  /** If true, finish the span when openCount hits 0. */
   private final boolean finishOnClose;
+  /** Count of open scope and continuations */
+  private final AtomicInteger openCount;
+  /** Scope to placed in the thread local after close. May be null. */
   private final Scope toRestore;
+  /** Continuation that created this scope. May be null. */
+  private final Continuation continuation;
 
   ContinuableScope(
-      final ContextualScopeManager scopeManager, final Span wrapped, final boolean finishOnClose) {
-    this(scopeManager, new AtomicInteger(1), wrapped, finishOnClose);
+      final ContextualScopeManager scopeManager,
+      final DDSpan spanUnderScope,
+      final boolean finishOnClose) {
+    this(scopeManager, new AtomicInteger(1), null, spanUnderScope, finishOnClose);
   }
 
   private ContinuableScope(
       final ContextualScopeManager scopeManager,
-      final AtomicInteger refCount,
-      final Span wrapped,
+      final AtomicInteger openCount,
+      final Continuation continuation,
+      final DDSpan spanUnderScope,
       final boolean finishOnClose) {
-
     this.scopeManager = scopeManager;
-    this.refCount = refCount;
-    this.wrapped = wrapped;
+    this.openCount = openCount;
+    this.continuation = continuation;
+    this.spanUnderScope = spanUnderScope;
     this.finishOnClose = finishOnClose;
     this.toRestore = scopeManager.tlsScope.get();
     scopeManager.tlsScope.set(this);
@@ -41,20 +52,22 @@ public class ContinuableScope implements Scope, TraceScope {
 
   @Override
   public void close() {
-    if (scopeManager.tlsScope.get() != this) {
-      return;
+    if (null != continuation) {
+      spanUnderScope.context().getTrace().cancelContinuation(continuation);
     }
 
-    if (refCount.decrementAndGet() == 0 && finishOnClose) {
-      wrapped.finish();
+    if (openCount.decrementAndGet() == 0 && finishOnClose) {
+      spanUnderScope.finish();
     }
 
+    if (scopeManager.tlsScope.get() == this) {
       scopeManager.tlsScope.set(toRestore);
     }
+  }
 
   @Override
-  public Span span() {
-    return wrapped;
+  public DDSpan span() {
+    return spanUnderScope;
   }
 
   /**
@@ -63,8 +76,8 @@ public class ContinuableScope implements Scope, TraceScope {
    * @param finishOnClose
    * @return
    */
-  public Continuation capture(final boolean finishOnClose) {
-    return new Continuation(this.finishOnClose && finishOnClose);
+  public Continuation capture() {
+    return new Continuation();
   }
 
   public class Continuation implements Closeable, TraceScope.Continuation {
@@ -72,72 +85,32 @@ public class ContinuableScope implements Scope, TraceScope {
 
     private final AtomicBoolean used = new AtomicBoolean(false);
     private final PendingTrace trace;
-    private final boolean finishSpanOnClose;
 
-    private Continuation(final boolean finishOnClose) {
-      this.finishSpanOnClose = finishOnClose;
-      refCount.incrementAndGet();
-      if (wrapped.context() instanceof DDSpanContext) {
-        final DDSpanContext context = (DDSpanContext) wrapped.context();
+    private Continuation() {
+      openCount.incrementAndGet();
+      final DDSpanContext context = (DDSpanContext) spanUnderScope.context();
       trace = context.getTrace();
       trace.registerContinuation(this);
-      } else {
-        trace = null;
-      }
     }
 
-    public ClosingScope activate() {
+    public ContinuableScope activate() {
       if (used.compareAndSet(false, true)) {
-        for (final ScopeContext context : scopeManager.scopeContexts) {
-          if (context.inContext()) {
-            return new ClosingScope(context.activate(wrapped, finishSpanOnClose));
-          }
-        }
-        return new ClosingScope(
-            new ContinuableScope(scopeManager, refCount, wrapped, finishSpanOnClose));
-      } else {
-        log.debug("Reusing a continuation not allowed.  Returning no-op scope.");
-        return new ClosingScope(NoopScopeManager.NoopScope.INSTANCE);
-      }
-    }
-
-    @Override
-    public void close() {
-      used.getAndSet(true);
-      if (trace != null) {
-        trace.cancelContinuation(this);
-      }
-    }
-
-    private class ClosingScope implements Scope, TraceScope {
-      private final Scope wrappedScope;
-
-      private ClosingScope(final Scope wrappedScope) {
-        this.wrappedScope = wrappedScope;
-      }
-
-      @Override
-      public Continuation capture(boolean finishOnClose) {
-        if (wrappedScope instanceof TraceScope) {
-          return ((TraceScope) wrappedScope).capture(finishOnClose);
+        return new ContinuableScope(scopeManager, openCount, this, spanUnderScope, finishOnClose);
       } else {
         log.debug(
-              "{} Failed to capture. ClosingScope does not wrap a TraceScope: {}.",
-              this,
-              wrappedScope);
-          return null;
+            "Failed to activate continuation. Reusing a continuation not allowed.  Returning a new scope. Spans will not be linked.");
+        return new ContinuableScope(
+            scopeManager, new AtomicInteger(1), null, spanUnderScope, finishOnClose);
       }
     }
 
     @Override
     public void close() {
-        wrappedScope.close();
-        ContinuableScope.Continuation.this.close();
-      }
-
-      @Override
-      public Span span() {
-        return wrappedScope.span();
+      if (used.compareAndSet(false, true)) {
+        trace.cancelContinuation(this);
+        ContinuableScope.this.close();
+      } else {
+        log.debug("Failed to close continuation {}. Already used.", this);
       }
     }
   }

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/SimpleScope.java

@@ -0,0 +1,39 @@
+package datadog.opentracing.scopemanager;
+
+import io.opentracing.Scope;
+import io.opentracing.Span;
+
+/** Simple scope implementation which does not propagate across threads. */
+public class SimpleScope implements Scope {
+  private final ContextualScopeManager scopeManager;
+  private final Span spanUnderScope;
+  private final boolean finishOnClose;
+  private final Scope toRestore;
+
+  public SimpleScope(
+      final ContextualScopeManager scopeManager,
+      final Span spanUnderScope,
+      final boolean finishOnClose) {
+    this.scopeManager = scopeManager;
+    this.spanUnderScope = spanUnderScope;
+    this.finishOnClose = finishOnClose;
+    this.toRestore = scopeManager.tlsScope.get();
+    scopeManager.tlsScope.set(this);
+  }
+
+  @Override
+  public void close() {
+    if (finishOnClose) {
+      spanUnderScope.finish();
+    }
+
+    if (scopeManager.tlsScope.get() == this) {
+      scopeManager.tlsScope.set(toRestore);
+    }
+  }
+
+  @Override
+  public Span span() {
+    return spanUnderScope;
+  }
+}

dd-trace-ot/src/test/groovy/datadog/opentracing/scopemanager/ScopeManagerTest.groovy

@@ -92,11 +92,11 @@ class ScopeManagerTest extends Specification {
     finishSpan << [true, false]
   }
 
-  def "ref counting scope doesn't close if non-zero"() {
+  def "ContinuableScope doesn't close if non-zero"() {
     setup:
     def builder = tracer.buildSpan("test")
     def scope = (ContinuableScope) builder.startActive(true)
-    def continuation = scope.capture(true)
+    def continuation = scope.capture()
 
     expect:
     !spanFinished(scope.span())
@@ -149,7 +149,7 @@ class ScopeManagerTest extends Specification {
     def builder = tracer.buildSpan("test")
     def scope = (ContinuableScope) builder.startActive(false)
     def span = scope.span()
-    def continuation = scope.capture(true)
+    def continuation = scope.capture()
     scope.close()
     span.finish()
 
@@ -188,7 +188,7 @@ class ScopeManagerTest extends Specification {
     ContinuableScope childScope = (ContinuableScope) tracer.buildSpan("parent").startActive(true)
     def childSpan = childScope.span()
 
-    def continuation = childScope.capture(true)
+    def continuation = childScope.capture()
     childScope.close()
 
     expect:
@@ -209,11 +209,11 @@ class ScopeManagerTest extends Specification {
 
     when:
     def newScope = continuation.activate()
-    def newContinuation = newScope.capture(true)
+    def newContinuation = newScope.capture()
 
     then:
-    newScope instanceof ContinuableScope.Continuation.ClosingScope
-    scopeManager.active() == newScope.wrappedScope
+    newScope instanceof ContinuableScope
+    scopeManager.active() == newScope
     newScope != childScope && newScope != parentScope
     newScope.span() == childSpan
     !spanFinished(childSpan)
@@ -236,16 +236,16 @@ class ScopeManagerTest extends Specification {
     def builder = tracer.buildSpan("test")
     def scope = (ContinuableScope) builder.startActive(false)
     def span = scope.span()
-    def continuation = scope.capture(false)
+    def continuation = scope.capture()
     scope.close()
     span.finish()
 
     def newScope = continuation.activate()
 
     expect:
-    newScope instanceof ContinuableScope.Continuation.ClosingScope
+    newScope instanceof ContinuableScope
     newScope != scope
-    scopeManager.active() == newScope.wrappedScope
+    scopeManager.active() == newScope
     spanFinished(span)
     writer == []
 
@@ -259,24 +259,7 @@ class ScopeManagerTest extends Specification {
     scopeManager.active() == null
     spanFinished(childSpan)
     childSpan.context().parentId == span.context().spanId
-    writer == []
-
-    when:
-    if (closeScope) {
-      newScope.close()
-    }
-    if (closeContinuation) {
-      continuation.close()
-    }
-
-    then:
     writer == [[childSpan, span]]
-
-    where:
-    closeScope | closeContinuation
-    true       | false
-    false      | true
-    true       | true
   }
 
   @Unroll
@@ -327,40 +310,27 @@ class ScopeManagerTest extends Specification {
     [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)] | _
   }
 
-  @Unroll
-  def "threadlocal to context with capture (#active)"() {
+  def "ContinuableScope put in threadLocal after continuation activation"() {
     setup:
-    contexts.each {
-      scopeManager.addScopeContext(it)
-    }
     ContinuableScope scope = (ContinuableScope) tracer.buildSpan("parent").startActive(true)
 
     expect:
     scopeManager.tlsScope.get() == scope
 
     when:
-    def cont = scope.capture(true)
+    def cont = scope.capture()
     scope.close()
 
     then:
     scopeManager.tlsScope.get() == null
 
     when:
-    active.each {
-      ((AtomicBoolean) contexts[it].enabled).set(true)
-    }
-    cont.activate()
+    scopeManager.addScopeContext(new AtomicReferenceScope(true))
+    def newScope = cont.activate()
 
     then:
-    scopeManager.tlsScope.get() == null
-
-    where:
-    active | contexts
-    [0]    | [new AtomicReferenceScope(false)]
-    [0]    | [new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
-    [1]    | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
-    [2]    | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
-    [0, 2] | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
+    newScope != scope
+    scopeManager.tlsScope.get() == newScope
   }
 
   @Unroll