From 8b77a18ede5e0d63468619562c094d87ef27e5a4 Mon Sep 17 00:00:00 2001
From: Stuart McCulloch <stuart.mcculloch@datadoghq.com>
Date: Wed, 23 Apr 2025 09:56:49 +0100
Subject: [PATCH] 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.
---
 .../core/scopemanager/ContinuableScope.java   | 28 ++++-----
 .../core/scopemanager/ScopeManagerTest.groovy | 57 +++++++++++++++++++
 2 files changed, 67 insertions(+), 18 deletions(-)

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java
index eb38e2289ff..8451f914405 100644
--- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java
+++ b/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;
     }
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy
index 9b3246f389e..5bad02bedae 100644
--- a/dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy
+++ b/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()
   }