Making zio indy ready (#14758)

Author: LikeTheSalad

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/TracingSupervisor.java

@@ -18,6 +18,10 @@
 @SuppressWarnings("unchecked")
 public final class TracingSupervisor extends Supervisor<Object> {
 
+  @SuppressWarnings("rawtypes")
+  public static final VirtualField<Fiber.Runtime, FiberContext> RUNTIME_FIBER_CONTEXT =
+      VirtualField.find(Fiber.Runtime.class, FiberContext.class);
+
   @SuppressWarnings("rawtypes")
   private final VirtualField<Fiber.Runtime, FiberContext> virtualField;
 

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/ZioInstrumentationModule.java

@@ -10,10 +10,12 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class ZioInstrumentationModule extends InstrumentationModule {
+public class ZioInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
 
   public ZioInstrumentationModule() {
     super("zio", "zio-2.0");
@@ -23,4 +25,9 @@ public ZioInstrumentationModule() {
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new ZioRuntimeInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/ZioRuntimeInstrumentation.java

@@ -5,16 +5,15 @@
 
 package io.opentelemetry.javaagent.instrumentation.zio.v2_0;
 
+import static io.opentelemetry.javaagent.instrumentation.zio.v2_0.TracingSupervisor.RUNTIME_FIBER_CONTEXT;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
-import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import zio.Fiber;
 import zio.Supervisor;
 
 public class ZioRuntimeInstrumentation implements TypeInstrumentation {
@@ -36,11 +35,9 @@ public static final class DefaultSupervisor {
     private DefaultSupervisor() {}
 
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void onExit(@Advice.Return(readOnly = false) Supervisor<?> supervisor) {
-      @SuppressWarnings("rawtypes")
-      VirtualField<Fiber.Runtime, FiberContext> virtualField =
-          VirtualField.find(Fiber.Runtime.class, FiberContext.class);
-      supervisor = supervisor.$plus$plus(new TracingSupervisor(virtualField));
+    @Advice.AssignReturned.ToReturned
+    public static Object onExit(@Advice.Return Supervisor<?> supervisor) {
+      return supervisor.$plus$plus(new TracingSupervisor(RUNTIME_FIBER_CONTEXT));
     }
   }
 }