From 61a3d403b4c4e61248d857e46b9935d5d053ea3b Mon Sep 17 00:00:00 2001
From: Stuart McCulloch <stuart.mcculloch@datadoghq.com>
Date: Sat, 8 Mar 2025 22:11:44 +0000
Subject: [PATCH] Migrate zio instrumentation to new Context API

---
 .../zio/v2_0/FiberContext.java                | 53 ++++++++-----------
 1 file changed, 23 insertions(+), 30 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 07a6302b218..075b4262335 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
@@ -2,63 +2,56 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
-import datadog.trace.bootstrap.instrumentation.api.ScopeState;
 
 public class FiberContext {
-  private final ScopeState state;
+  private final Context context;
   private AgentScope.Continuation continuation;
   private AgentScope scope;
-  private ScopeState oldState;
+  private Context oldContext;
 
-  private FiberContext(ScopeState state) {
-    this.state = state;
-    this.scope = null;
-    this.oldState = null;
+  private FiberContext(Context context) {
+    this.context = context;
     this.continuation = captureActiveSpan();
+    this.scope = null;
+    this.oldContext = null;
   }
 
   public static FiberContext create() {
-    final ScopeState state = AgentTracer.get().newScopeState();
-    return new FiberContext(state);
+    return new FiberContext(Context.root());
   }
 
   public void onEnd() {
-    if (this.scope != null) {
-      this.scope.close();
-      this.scope = null;
+    if (scope != null) {
+      scope.close();
+      scope = null;
     }
     if (continuation != null) {
       continuation.cancel();
       continuation = null;
     }
-
-    if (this.oldState != null) {
-      this.oldState.activate();
-      this.oldState = null;
+    if (oldContext != null) {
+      oldContext.swap();
+      oldContext = null;
     }
   }
 
   public void onSuspend() {
-    if (this.scope != null && continuation != null) {
-      this.scope.close();
-      this.scope = null;
+    if (scope != null && continuation != null) {
+      scope.close();
+      scope = null;
     }
-    if (this.oldState != null) {
-      this.oldState.activate();
-      this.oldState = null;
+    if (oldContext != null) {
+      oldContext.swap();
+      oldContext = null;
     }
   }
 
   public void onResume() {
-    this.oldState = AgentTracer.get().newScopeState();
-    this.oldState.fetchFromActive();
-
-    this.state.activate();
-
-    if (this.continuation != null) {
-      this.scope = continuation.activate();
+    oldContext = context.swap();
+    if (continuation != null) {
+      scope = continuation.activate();
       continuation = null;
     }
   }