simplify

Author: amarziali

dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java

@@ -1,22 +1,22 @@
 package datadog.trace.instrumentation.zio.v2_0;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.capture;
+
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ScopeState;
 
 public class FiberContext {
-
   private final ScopeState state;
-  private AgentSpan span;
+  private AgentScope.Continuation continuation;
   private AgentScope scope;
   private ScopeState oldState;
 
   private FiberContext(ScopeState state) {
     this.state = state;
-    this.span = AgentTracer.activeSpan();
     this.scope = null;
     this.oldState = null;
+    this.continuation = capture();
   }
 
   public static FiberContext create() {
@@ -29,6 +29,10 @@ public void onEnd() {
       this.scope.close();
       this.scope = null;
     }
+    if (continuation != null) {
+      continuation.cancel();
+      continuation = null;
+    }
 
     if (this.oldState != null) {
       this.oldState.activate();
@@ -37,13 +41,14 @@ public void onEnd() {
   }
 
   public void onSuspend() {
+    if (this.scope != null) {
+      this.scope.close();
+      this.scope = null;
+    }
     if (this.oldState != null) {
       this.oldState.activate();
       this.oldState = null;
     }
-    if (this.scope != null) {
-      this.scope.close();
-    }
   }
 
   public void onResume() {
@@ -52,8 +57,9 @@ public void onResume() {
 
     this.state.activate();
 
-    if (this.span != null && this.scope == null) {
-      this.scope = AgentTracer.activateSpan(span, true);
+    if (this.continuation != null) {
+      this.scope = continuation.activate();
+      continuation = null;
     }
   }
 }

dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/TracingSupervisor.java

@@ -4,6 +4,7 @@
 import scala.Option;
 import zio.Exit;
 import zio.Fiber;
+import zio.FiberId;
 import zio.Supervisor;
 import zio.Unsafe;
 import zio.ZEnvironment;
@@ -12,12 +13,9 @@
 
 @SuppressWarnings("unchecked")
 public final class TracingSupervisor extends Supervisor<Object> {
+  private final ContextStore<FiberId.Runtime, FiberContext> contextStore;
 
-  @SuppressWarnings("rawtypes")
-  private final ContextStore<Fiber.Runtime, FiberContext> contextStore;
-
-  @SuppressWarnings("rawtypes")
-  public TracingSupervisor(ContextStore<Fiber.Runtime, FiberContext> contextStore) {
+  public TracingSupervisor(ContextStore<FiberId.Runtime, FiberContext> contextStore) {
     this.contextStore = contextStore;
   }
 
@@ -35,24 +33,24 @@ public <R, E, A_> void onStart(
       Fiber.Runtime<E, A_> fiber,
       Unsafe unsafe) {
     FiberContext context = FiberContext.create();
-    contextStore.put(fiber, context);
+    contextStore.put(fiber.id(), context);
   }
 
   @Override
   public <R, E, A_> void onEnd(Exit<E, A_> value, Fiber.Runtime<E, A_> fiber, Unsafe unsafe) {
-    FiberContext context = contextStore.get(fiber);
+    FiberContext context = contextStore.get(fiber.id());
     if (context != null) context.onEnd();
   }
 
   @Override
   public <E, A_> void onSuspend(Fiber.Runtime<E, A_> fiber, Unsafe unsafe) {
-    FiberContext context = contextStore.get(fiber);
+    FiberContext context = contextStore.get(fiber.id());
     if (context != null) context.onSuspend();
   }
 
   @Override
   public <E, A_> void onResume(Fiber.Runtime<E, A_> fiber, Unsafe unsafe) {
-    FiberContext context = contextStore.get(fiber);
+    FiberContext context = contextStore.get(fiber.id());
     if (context != null) context.onResume();
   }
 }

dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/ZioRuntimeInstrumentation.java

@@ -14,7 +14,7 @@
 import java.util.Collections;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import zio.Fiber;
+import zio.FiberId;
 import zio.Supervisor;
 
 @AutoService(Instrumenter.class)
@@ -48,7 +48,10 @@ public String[] helperClassNames() {
 
   @Override
   public Map<String, String> contextStore() {
-    return singletonMap("zio.Fiber$Runtime", packageName + ".FiberContext");
+    // it looks like zio.Fiber$Runtime will clash with Runnable causing our threadpool
+    // instrumentation to fetch this key
+    // but failing in casting to State. Hence, using another class as a key
+    return singletonMap("zio.FiberId$Runtime", packageName + ".FiberContext");
   }
 
   @Override
@@ -60,9 +63,8 @@ public Map<String, String> contextStore() {
   public static final class DefaultSupervisor {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(@Advice.Return(readOnly = false) Supervisor<?> supervisor) {
-      @SuppressWarnings("rawtypes")
-      ContextStore<Fiber.Runtime, FiberContext> contextStore =
-          InstrumentationContext.get(Fiber.Runtime.class, FiberContext.class);
+      ContextStore<FiberId.Runtime, FiberContext> contextStore =
+          InstrumentationContext.get(FiberId.Runtime.class, FiberContext.class);
       supervisor = supervisor.$plus$plus(new TracingSupervisor(contextStore));
     }
   }