ZIO 2.0 instrumentation

dd-java-agent/instrumentation/zio/build.gradle

@@ -0,0 +1 @@
+apply from: "$rootDir/gradle/java.gradle"

dd-java-agent/instrumentation/zio/zio-2.0/build.gradle

@@ -0,0 +1,48 @@
+def zioVersion = '2.0.0'
+def scalaVersion = '2.12'
+
+muzzle {
+  pass {
+    group = 'dev.zio'
+    module = "zio_2.12"
+    versions = "[$zioVersion,)"
+    assertInverse = true
+  }
+  pass {
+    group = 'dev.zio'
+    module = "zio_2.13"
+    versions = "[$zioVersion,)"
+    assertInverse = true
+  }
+  pass {
+    group = 'dev.zio'
+    module = "zio_3"
+    versions = "[$zioVersion,)"
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply from: "$rootDir/gradle/test-with-scala.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+tasks.named("compileLatestDepTestGroovy").configure {
+  dependsOn "compileLatestDepTestScala"
+  classpath += files(compileLatestDepTestScala.destinationDirectory)
+}
+
+dependencies {
+  compileOnly group: 'dev.zio', name: "zio_$scalaVersion", version: zioVersion
+
+  testImplementation deps.scala
+  testImplementation group: 'dev.zio', name: "zio_$scalaVersion", version: zioVersion
+
+  latestDepTestImplementation group: 'dev.zio', name: "zio_$scalaVersion", version: '+'
+}

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

@@ -0,0 +1,70 @@
+package datadog.trace.instrumentation.zio.v2_0;
+
+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 AgentScope.Continuation continuation;
+  private AgentScope scope;
+  private ScopeState oldState;
+
+  private FiberContext(ScopeState state, AgentScope.Continuation continuation) {
+    this.state = state;
+    this.continuation = continuation;
+    this.scope = null;
+    this.oldState = null;
+  }
+
+  public static FiberContext create() {
+    final ScopeState state = AgentTracer.get().newScopeState();
+    final AgentScope scope = AgentTracer.get().activeScope();
+    final AgentScope.Continuation continuation;
+
+    if (scope != null && scope.isAsyncPropagating()) {
+      continuation = scope.capture();
+    } else {
+      continuation = null;
+    }
+
+    return new FiberContext(state, continuation);
+  }
+
+  public void onEnd() {
+    if (this.scope != null) {
+      this.scope.close();
+      this.scope = null;
+    }
+
+    if (this.continuation != null) {
+      this.continuation.cancel();
+      this.continuation = null;
+    }
+
+    if (this.oldState != null) {
+      this.oldState.activate();
+      this.oldState = null;
+    }
+  }
+
+  public void onSuspend() {
+    if (this.oldState != null) {
+      this.oldState.activate();
+      this.oldState = null;
+    }
+  }
+
+  public void onResume() {
+    this.oldState = AgentTracer.get().newScopeState();
+    this.oldState.fetchFromActive();
+
+    this.state.activate();
+
+    if (this.continuation != null && this.scope == null) {
+      this.scope = this.continuation.activate();
+      this.continuation = null;
+    }
+  }
+}

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

@@ -0,0 +1,58 @@
+package datadog.trace.instrumentation.zio.v2_0;
+
+import datadog.trace.bootstrap.ContextStore;
+import scala.Option;
+import zio.Exit;
+import zio.Fiber;
+import zio.Supervisor;
+import zio.Unsafe;
+import zio.ZEnvironment;
+import zio.ZIO;
+import zio.ZIO$;
+
+@SuppressWarnings("unchecked")
+public final class TracingSupervisor extends Supervisor<Object> {
+
+  @SuppressWarnings("rawtypes")
+  private final ContextStore<Fiber.Runtime, FiberContext> contextStore;
+
+  @SuppressWarnings("rawtypes")
+  public TracingSupervisor(ContextStore<Fiber.Runtime, FiberContext> contextStore) {
+    this.contextStore = contextStore;
+  }
+
+  @Override
+  @SuppressWarnings("rawtypes")
+  public ZIO value(Object trace) {
+    return ZIO$.MODULE$.unit();
+  }
+
+  @Override
+  public <R, E, A_> void onStart(
+      ZEnvironment<R> environment,
+      ZIO<R, E, A_> effect,
+      Option<Fiber.Runtime<Object, Object>> parent,
+      Fiber.Runtime<E, A_> fiber,
+      Unsafe unsafe) {
+    FiberContext context = FiberContext.create();
+    contextStore.put(fiber, context);
+  }
+
+  @Override
+  public <R, E, A_> void onEnd(Exit<E, A_> value, Fiber.Runtime<E, A_> fiber, Unsafe unsafe) {
+    FiberContext context = contextStore.get(fiber);
+    if (context != null) context.onEnd();
+  }
+
+  @Override
+  public <E, A_> void onSuspend(Fiber.Runtime<E, A_> fiber, Unsafe unsafe) {
+    FiberContext context = contextStore.get(fiber);
+    if (context != null) context.onSuspend();
+  }
+
+  @Override
+  public <E, A_> void onResume(Fiber.Runtime<E, A_> fiber, Unsafe unsafe) {
+    FiberContext context = contextStore.get(fiber);
+    if (context != null) context.onResume();
+  }
+}

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

@@ -0,0 +1,59 @@
+package datadog.trace.instrumentation.zio.v2_0;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import zio.Fiber;
+import zio.Supervisor;
+
+@AutoService(Instrumenter.class)
+public class ZioRuntimeInstrumentation extends Instrumenter.Tracing
+    implements Instrumenter.ForSingleType {
+
+  public ZioRuntimeInstrumentation() {
+    super("zio.experimental");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "zio.Runtime$";
+  }
+
+  @Override
+  protected final boolean defaultEnabled() {
+    return false;
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod().and(named("defaultSupervisor")), getClass().getName() + "$DefaultSupervisor");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".FiberContext", packageName + ".TracingSupervisor"};
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("zio.Fiber$Runtime", packageName + ".FiberContext");
+  }
+
+  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);
+      supervisor = supervisor.$plus$plus(new TracingSupervisor(contextStore));
+    }
+  }
+}