Combine continuation implementations into one which supports multiple activations

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

@@ -69,23 +69,23 @@ public static <K> void closeScope(
     if (throwable != null) {
       // This might lead to the continuation being consumed early, but it's better to be safe if we
       // threw an Exception on entry
-      state.closeContinuation();
+      state.cancelContinuation();
     }
   }
 
-  public static <K> void closeAndClearContinuation(
+  public static <K> void cancelAndClearContinuation(
       ContextStore<K, ConcurrentState> contextStore, K key) {
     final ConcurrentState state = contextStore.get(key);
     if (state == null) {
       return;
     }
-    state.closeAndClearContinuation();
+    state.cancelAndClearContinuation();
   }
 
   private boolean captureAndSetContinuation(final AgentScope scope) {
     if (CONTINUATION.compareAndSet(this, null, CLAIMED)) {
       // lazy write is guaranteed to be seen by getAndSet
-      CONTINUATION.lazySet(this, scope.captureConcurrent());
+      CONTINUATION.lazySet(this, scope.capture().hold());
       return true;
     }
     return false;
@@ -99,14 +99,14 @@ private AgentScope activateAndContinueContinuation() {
     return null;
   }
 
-  private void closeContinuation() {
+  private void cancelContinuation() {
     final AgentScope.Continuation continuation = CONTINUATION.get(this);
     if (continuation != null && continuation != CLAIMED) {
       continuation.cancel();
     }
   }
 
-  private void closeAndClearContinuation() {
+  private void cancelAndClearContinuation() {
     final AgentScope.Continuation continuation = CONTINUATION.get(this);
     if (continuation != null && continuation != CLAIMED) {
       // We should never be able to reuse this state

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

@@ -7,6 +7,11 @@ final class ContinuationClaim implements AgentScope.Continuation {
 
   public static final ContinuationClaim CLAIMED = new ContinuationClaim();
 
+  @Override
+  public AgentScope.Continuation hold() {
+    throw new IllegalStateException();
+  }
+
   @Override
   public AgentScope activate() {
     throw new IllegalStateException();

dd-java-agent/instrumentation/java-concurrent/java-completablefuture/src/main/java/java/util/concurrent/CompletableFutureAdvice.java

@@ -72,7 +72,7 @@ public static void exit(
       boolean claimed = !wasClaimed && !hadExecutor && zis.getForkJoinTaskTag() == 1;
       if (mode == ASYNC || (mode < ASYNC && claimed) || !zis.isLive()) {
         contextStore = InstrumentationContext.get(UniCompletion.class, ConcurrentState.class);
-        ConcurrentState.closeAndClearContinuation(contextStore, zis);
+        ConcurrentState.cancelAndClearContinuation(contextStore, zis);
       }
       if (scope != null || throwable != null) {
         if (contextStore == null) {

dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java

@@ -121,7 +121,7 @@ public void maybeInitialize() {
       if (!isInitialized) {
         final AgentScope activeScope = AgentTracer.get().activeScope();
         if (activeScope != null && activeScope.isAsyncPropagating()) {
-          continuation = activeScope.captureConcurrent();
+          continuation = activeScope.capture().hold();
         }
         isInitialized = true;
       }

dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/OtelScope.java

@@ -16,11 +16,6 @@ public Continuation capture() {
     return delegate.capture();
   }
 
-  @Override
-  public Continuation captureConcurrent() {
-    return delegate.captureConcurrent();
-  }
-
   @Override
   public void close() {
     delegate.close();

dd-java-agent/instrumentation/opentracing/api-0.31/src/main/java/datadog/trace/instrumentation/opentracing31/OTScopeManager.java

@@ -68,11 +68,6 @@ public Continuation capture() {
       return delegate.capture();
     }
 
-    @Override
-    public Continuation captureConcurrent() {
-      return delegate.captureConcurrent();
-    }
-
     public boolean isFinishSpanOnClose() {
       return finishSpanOnClose;
     }

dd-java-agent/instrumentation/opentracing/api-0.32/src/main/java/datadog/trace/instrumentation/opentracing32/OTScopeManager.java

@@ -78,11 +78,6 @@ public Continuation capture() {
       return delegate.capture();
     }
 
-    @Override
-    public Continuation captureConcurrent() {
-      return delegate.captureConcurrent();
-    }
-
     public boolean isFinishSpanOnClose() {
       return finishSpanOnClose;
     }

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

@@ -8,6 +8,11 @@ public static class NoopContinuation implements Continuation {
 
     private NoopContinuation() {}
 
+    @Override
+    public Continuation hold() {
+      return this;
+    }
+
     @Override
     public TraceScope activate() {
       return NoopTraceScope.INSTANCE;
@@ -24,11 +29,6 @@ public Continuation capture() {
     return NoopContinuation.INSTANCE;
   }
 
-  @Override
-  public Continuation captureConcurrent() {
-    return null;
-  }
-
   @Override
   public void close() {}
 }

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

@@ -14,18 +14,11 @@ public interface TraceScope extends Closeable {
    */
   Continuation capture();
 
-  /**
-   * Prevent the trace attached to this TraceScope from reporting until the returned Continuation is
-   * either activated (and the returned scope is closed), or canceled.
-   *
-   * <p>Should be called on the parent thread.
-   *
-   * <p>If the returned {@link Continuation} is activated, it needs to be canceled in addition to
-   * the returned {@link TraceScope} being closed. This is to allow multiple concurrent threads that
-   * activate the continuation to race in a safe way, and close the scopes without fear of closing
-   * the related {@code Span} prematurely.
-   */
-  Continuation captureConcurrent();
+  /** @deprecated Replaced by {@code capture().hold()}. */
+  @Deprecated
+  default Continuation captureConcurrent() {
+    return capture().hold();
+  }
 
   /** Close the activated context and allow any underlying spans to finish. */
   @Override
@@ -61,6 +54,15 @@ default void setAsyncPropagation(boolean value) {
    */
   interface Continuation {
 
+    /**
+     * Prevent the trace attached to this scope from reporting until the continuation is explicitly
+     * cancelled. You must call {@link #cancel()} at some point to avoid discarding traces.
+     *
+     * <p>Use this when you want to let multiple threads activate the continuation concurrently and
+     * close their scopes without fear of prematurely closing the related span.
+     */
+    Continuation hold();
+
     /**
      * Activate the continuation.
      *

dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java

@@ -135,21 +135,9 @@ public final void setAsyncPropagation(final boolean value) {
    * @return The new continuation, or null if this scope is not async propagating.
    */
   @Override
-  public final AbstractContinuation capture() {
+  public final ScopeContinuation capture() {
     return isAsyncPropagating
-        ? new SingleContinuation(scopeManager, span, source()).register()
-        : null;
-  }
-
-  /**
-   * The continuation returned must be closed or activated or the trace will not finish.
-   *
-   * @return The new continuation, or null if this scope is not async propagating.
-   */
-  @Override
-  public final AbstractContinuation captureConcurrent() {
-    return isAsyncPropagating
-        ? new ConcurrentContinuation(scopeManager, span, source()).register()
+        ? new ScopeContinuation(scopeManager, span, source()).register()
         : null;
   }
 

dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java

@@ -91,8 +91,8 @@ public AgentScope activate(
   }
 
   public AgentScope.Continuation captureSpan(final AgentSpan span) {
-    AbstractContinuation continuation =
-        new SingleContinuation(this, span, ScopeSource.INSTRUMENTATION.id());
+    ScopeContinuation continuation =
+        new ScopeContinuation(this, span, ScopeSource.INSTRUMENTATION.id());
     continuation.register();
     healthMetrics.onCaptureContinuation();
     return continuation;
@@ -136,12 +136,12 @@ private AgentScope activate(
   }
 
   /**
-   * Activates a scope for the given {@link AbstractContinuation}.
+   * Activates a scope for the given {@link ScopeContinuation}.
    *
    * @param continuation {@code null} if a continuation is re-used
    */
   ContinuableScope continueSpan(
-      final AbstractContinuation continuation, final AgentSpan span, final byte source) {
+      final ScopeContinuation continuation, final AgentSpan span, final byte source) {
     ScopeStack scopeStack = scopeStack();
 
     // optimization: if the top scope is already keeping the same span alive

dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuingScope.java

@@ -5,14 +5,14 @@
 
 final class ContinuingScope extends ContinuableScope {
   /** Continuation that created this scope. */
-  private final AbstractContinuation continuation;
+  private final ScopeContinuation continuation;
 
   ContinuingScope(
       final ContinuableScopeManager scopeManager,
       final AgentSpan span,
       final byte source,
       final boolean isAsyncPropagating,
-      final AbstractContinuation continuation,
+      final ScopeContinuation continuation,
       final Stateful scopeState) {
     super(scopeManager, span, source, isAsyncPropagating, scopeState);
     this.continuation = continuation;
@@ -21,7 +21,6 @@ final class ContinuingScope extends ContinuableScope {
   @Override
   void cleanup(final ScopeStack scopeStack) {
     super.cleanup(scopeStack);
-
     continuation.cancelFromContinuedScopeClose();
   }
 }