Handle reentrant scope cleanup in Akka/Pekko actor instrumentations (#8723)

Akka/Pekko actor instrumentations use the scope checkpoint and rollback mechanism to clean up scopes after async operations. To properly handle reentrant scope checkpointing and rollback we need to use a count in ContinuableScope like referenceCount rather than just a flag.

This is necessary because ContinuableScopeManager re-uses scope instances when activating the same span twice.

(cherry picked from commit 8b77a18)

Author: mcculls

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

@@ -19,13 +19,9 @@ class ContinuableScope implements AgentScope {
 
   final Context context; // package-private so scopeManager can access it directly
 
-  /** Flag that this scope should be allowed to propagate across async boundaries. */
-  private static final byte ASYNC_PROPAGATING = 1;
+  private boolean asyncPropagating;
 
-  /** 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 short checkpointCount = 0;
 
   private final byte source;
 
@@ -37,12 +33,12 @@ class ContinuableScope implements AgentScope {
       final ContinuableScopeManager scopeManager,
       final Context context,
       final byte source,
-      final boolean isAsyncPropagating,
+      final boolean asyncPropagating,
       final Stateful scopeState) {
     this.scopeManager = scopeManager;
     this.context = context;
     this.source = source;
-    this.flags = isAsyncPropagating ? ASYNC_PROPAGATING : 0;
+    this.asyncPropagating = asyncPropagating;
     this.scopeState = scopeState;
   }
 
@@ -121,7 +117,7 @@ final boolean alive() {
 
   @Override
   public final boolean isAsyncPropagating() {
-    return (flags & ASYNC_PROPAGATING) != 0;
+    return asyncPropagating;
   }
 
   @Override
@@ -136,11 +132,7 @@ public Context context() {
 
   @Override
   public final void setAsyncPropagation(final boolean value) {
-    if (value) {
-      flags |= ASYNC_PROPAGATING;
-    } else {
-      flags &= ~ASYNC_PROPAGATING;
-    }
+    asyncPropagating = value;
   }
 
   @Override
@@ -149,13 +141,13 @@ public final String toString() {
   }
 
   public void checkpoint() {
-    flags |= CHECKPOINTED;
+    checkpointCount++;
   }
 
   public boolean rollback() {
-    if ((flags & CHECKPOINTED) != 0) {
-      flags &= ~CHECKPOINTED;
-      return false;
+    if (checkpointCount > 0) {
+      checkpointCount--;
+      return false; // stop rollback at checkpoint
     } else {
       return true;
     }

dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy

@@ -1157,6 +1157,63 @@ class ScopeManagerTest extends DDCoreSpecification {
     scopeManager.activeSpan() == null
   }
 
+  def "rollback stops at most recent checkpoint"() {
+    when:
+    def span1 = tracer.buildSpan("test1", "test1").start()
+    def span2 = tracer.buildSpan("test2", "test2").start()
+    def span3 = tracer.buildSpan("test3", "test3").start()
+    then:
+    scopeManager.activeSpan() == null
+
+    when:
+    tracer.checkpointActiveForRollback()
+    tracer.activateSpan(span1)
+    tracer.checkpointActiveForRollback()
+    tracer.activateSpan(span2)
+    tracer.checkpointActiveForRollback()
+    tracer.activateSpan(span1)
+    tracer.checkpointActiveForRollback()
+    tracer.activateSpan(span2)
+    tracer.checkpointActiveForRollback()
+    tracer.activateSpan(span2)
+    tracer.checkpointActiveForRollback()
+    tracer.activateSpan(span1)
+    tracer.activateSpan(span2)
+    tracer.activateSpan(span3)
+    then:
+    scopeManager.activeSpan() == span3
+
+    when:
+    tracer.rollbackActiveToCheckpoint()
+    then:
+    scopeManager.activeSpan() == span2
+
+    when:
+    tracer.rollbackActiveToCheckpoint()
+    then:
+    scopeManager.activeSpan() == span2
+
+    when:
+    tracer.rollbackActiveToCheckpoint()
+    then:
+    scopeManager.activeSpan() == span1
+
+    when:
+    tracer.rollbackActiveToCheckpoint()
+    then:
+    scopeManager.activeSpan() == span2
+
+    when:
+    tracer.rollbackActiveToCheckpoint()
+    then:
+    scopeManager.activeSpan() == span1
+
+    when:
+    tracer.rollbackActiveToCheckpoint()
+    then:
+    scopeManager.activeSpan() == null
+  }
+
   boolean spanFinished(AgentSpan span) {
     return ((DDSpan) span)?.isFinished()
   }