Add support for checkpointActiveForRollback...rollbackActiveToCheckpoint to clean up leaked scopes

Author: mcculls

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -977,6 +977,16 @@ public AgentScope activeScope() {
     return scopeManager.active();
   }
 
+  @Override
+  public void checkpointActiveForRollback() {
+    this.scopeManager.checkpointActiveForRollback();
+  }
+
+  @Override
+  public void rollbackActiveToCheckpoint() {
+    this.scopeManager.rollbackActiveToCheckpoint();
+  }
+
   @Override
   public void closeActive() {
     AgentScope activeScope = this.scopeManager.active();

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

@@ -15,10 +15,15 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
 
   final AgentSpan span; // package-private so scopeManager can access it directly
 
-  /** Flag to propagate this scope across async boundaries. */
-  private boolean isAsyncPropagating;
+  /** Flag that this scope should be allowed to propagate across async boundaries. */
+  private static final byte ASYNC_PROPAGATING = 1;
 
-  private final byte flags;
+  /** Flag that we intend to roll back the scope stack to this scope in the future. */
+  private static final byte CHECKPOINTED = 2;
+
+  private byte flags;
+
+  private final byte source;
 
   private short referenceCount = 1;
 
@@ -36,8 +41,8 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
       final Stateful scopeState) {
     this.scopeManager = scopeManager;
     this.span = span;
-    this.flags = source;
-    this.isAsyncPropagating = isAsyncPropagating;
+    this.source = source;
+    this.flags = isAsyncPropagating ? ASYNC_PROPAGATING : 0;
     this.scopeState = scopeState;
   }
 
@@ -116,7 +121,7 @@ final boolean alive() {
 
   @Override
   public final boolean isAsyncPropagating() {
-    return isAsyncPropagating;
+    return (flags & ASYNC_PROPAGATING) != 0;
   }
 
   @Override
@@ -126,14 +131,31 @@ public final AgentSpan span() {
 
   @Override
   public final void setAsyncPropagation(final boolean value) {
-    isAsyncPropagating = value;
+    if (value) {
+      flags |= ASYNC_PROPAGATING;
+    } else {
+      flags &= ~ASYNC_PROPAGATING;
+    }
   }
 
   @Override
   public final String toString() {
     return super.toString() + "->" + span;
   }
 
+  public void checkpoint() {
+    flags |= CHECKPOINTED;
+  }
+
+  public boolean rollback() {
+    if ((flags & CHECKPOINTED) != 0) {
+      flags &= ~CHECKPOINTED;
+      return false;
+    } else {
+      return true;
+    }
+  }
+
   public final void beforeActivated() {
     try {
       scopeState.activate(span.context());
@@ -164,7 +186,7 @@ public final void afterActivated() {
 
   @Override
   public byte source() {
-    return (byte) (flags & 0x7F);
+    return (byte) (source & 0x7F);
   }
 
   @Override

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

@@ -218,6 +218,24 @@ public AgentScope active() {
     return scopeStack().active();
   }
 
+  public void checkpointActiveForRollback() {
+    ContinuableScope active = scopeStack().active();
+    if (active != null) {
+      active.checkpoint();
+    }
+  }
+
+  public void rollbackActiveToCheckpoint() {
+    ContinuableScope active;
+    while ((active = scopeStack().active()) != null) {
+      if (active.rollback()) {
+        active.close();
+      } else {
+        break; // stop at the most recent checkpointed scope
+      }
+    }
+  }
+
   public AgentSpan activeSpan() {
     final ContinuableScope active = scopeStack().active();
     return active == null ? null : active.span;

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -121,10 +121,35 @@ public static AgentScope.Continuation captureSpan(final AgentSpan span) {
     return get().captureSpan(span);
   }
 
+  /**
+   * Checkpoints the active scope. A subsequent call to {@link #rollbackActiveToCheckpoint()} closes
+   * outstanding scopes up to but not including the most recent checkpointed scope.
+   *
+   * @deprecated This should only be used when scopes might leak onto the scope stack which cannot
+   *     be cleaned up by other means.
+   */
+  @Deprecated
+  public static void checkpointActiveForRollback() {
+    get().checkpointActiveForRollback();
+  }
+
+  /**
+   * Closes outstanding scopes up to but not including the most recent scope checkpointed with
+   * {@link #checkpointActiveForRollback()}. Closes all scopes if none have been checkpointed.
+   *
+   * @deprecated This should only be used when scopes have leaked onto the scope stack that cannot
+   *     be cleaned up by other means.
+   */
+  @Deprecated
+  public static void rollbackActiveToCheckpoint() {
+    get().rollbackActiveToCheckpoint();
+  }
+
   /**
    * Closes the scope for the currently active span.
    *
-   * @deprecated Prefer closing the scope returned by {@link #activateSpan} when available.
+   * @deprecated This should only be used when an instrumentation does not have access to the
+   *     original scope returned by {@link #activateSpan}.
    */
   @Deprecated
   public static void closeActive() {
@@ -322,6 +347,10 @@ AgentSpan startSpan(
 
     AgentScope.Continuation captureSpan(AgentSpan span);
 
+    void checkpointActiveForRollback();
+
+    void rollbackActiveToCheckpoint();
+
     void closeActive();
 
     void closePrevious(boolean finishSpan);
@@ -477,6 +506,12 @@ public boolean isAsyncPropagationEnabled() {
     @Override
     public void setAsyncPropagationEnabled(boolean asyncPropagationEnabled) {}
 
+    @Override
+    public void checkpointActiveForRollback() {}
+
+    @Override
+    public void rollbackActiveToCheckpoint() {}
+
     @Override
     public void closeActive() {}