From 8bc35cefaab1a6debde0d06f44c9994e18c2ba97 Mon Sep 17 00:00:00 2001
From: Stuart McCulloch <stuart.mcculloch@datadoghq.com>
Date: Sun, 11 May 2025 21:43:23 +0100
Subject: [PATCH] Cleanup Zio fiber instrumentation to avoid repeated
 activation of continuation

---
 .../zio/v2_0/FiberContext.java                | 60 ++++++-------------
 .../zio/v2_0/TracingSupervisor.java           |  2 +-
 2 files changed, 19 insertions(+), 43 deletions(-)

diff --git a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java
index 2a37e7c553e..83ccbc21607 100644
--- a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java
+++ b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java
@@ -7,58 +7,34 @@
 import datadog.trace.bootstrap.instrumentation.api.ScopeState;
 
 public class FiberContext {
-  private final ScopeState state;
-  private AgentScope.Continuation continuation;
-  private AgentScope scope;
-  private ScopeState oldState;
+  private final ScopeState scopeState;
+  private final AgentScope.Continuation continuation;
 
-  private FiberContext(ScopeState state) {
-    this.state = state;
-    this.scope = null;
-    this.oldState = null;
-    this.continuation = captureActiveSpan();
-  }
+  private ScopeState oldScopeState;
 
-  public static FiberContext create() {
-    final ScopeState state = AgentTracer.get().newScopeState();
-    return new FiberContext(state);
+  public FiberContext() {
+    // copy scope stack to use for this fiber
+    this.scopeState = AgentTracer.get().oldScopeState().copy();
+    // stop enclosing trace from finishing early
+    this.continuation = captureActiveSpan();
   }
 
-  public void onEnd() {
-    if (this.scope != null) {
-      this.scope.close();
-      this.scope = null;
-    }
-    if (continuation != null) {
-      continuation.cancel();
-      continuation = null;
-    }
-
-    if (this.oldState != null) {
-      this.oldState.activate();
-      this.oldState = null;
-    }
+  public void onResume() {
+    oldScopeState = AgentTracer.get().oldScopeState();
+    scopeState.activate(); // swap in the fiber's scope stack
   }
 
   public void onSuspend() {
-    if (this.scope != null && continuation != null) {
-      this.scope.close();
-      this.scope = null;
-    }
-    if (this.oldState != null) {
-      this.oldState.activate();
-      this.oldState = null;
+    if (oldScopeState != null) {
+      oldScopeState.activate(); // swap bock the original scope stack
+      oldScopeState = null;
     }
   }
 
-  public void onResume() {
-    this.oldState = AgentTracer.get().oldScopeState();
-
-    this.state.activate();
-
-    if (this.continuation != null) {
-      this.scope = continuation.activate();
-      continuation = null;
+  public void onEnd() {
+    if (continuation != null) {
+      // release enclosing trace now the fiber has ended
+      continuation.cancel();
     }
   }
 }
diff --git a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/TracingSupervisor.java b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/TracingSupervisor.java
index a2a8c819d0c..3439fdc0a4b 100644
--- a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/TracingSupervisor.java
+++ b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/TracingSupervisor.java
@@ -34,7 +34,7 @@ public <R, E, A_> void onStart(
       Option<Fiber.Runtime<Object, Object>> parent,
       Fiber.Runtime<E, A_> fiber,
       Unsafe unsafe) {
-    FiberContext context = FiberContext.create();
+    FiberContext context = new FiberContext();
     contextStore.put(fiber, context);
   }