make javaagent-tooling and testing-common indy-ready

javaagent-tooling/src/testExceptionHandler/java/io/opentelemetry/javaagent/tooling/bytebuddy/BadAdvice.java

@@ -5,14 +5,16 @@
 
 package io.opentelemetry.javaagent.tooling.bytebuddy;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import net.bytebuddy.asm.Advice;
 
 @SuppressWarnings("unused")
 public class BadAdvice {
 
   @Advice.OnMethodExit(suppress = Throwable.class)
-  public static void throwAnException(@Advice.Return(readOnly = false) boolean returnVal) {
-    returnVal = true;
+  public static void throwAnException(@Advice.Argument(0) AtomicBoolean isInstrumented) {
+    // mark that the advice has been executed
+    isInstrumented.set(true);
     throw new IllegalStateException("Test Exception");
   }
 

javaagent-tooling/src/testExceptionHandler/java/io/opentelemetry/javaagent/tooling/bytebuddy/ExceptionHandlerTest.java

@@ -15,11 +15,13 @@
 import io.opentelemetry.javaagent.bootstrap.ExceptionLogger;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
@@ -34,20 +36,24 @@ class ExceptionHandlerTest {
 
   @BeforeAll
   static void setUp() {
+    Advice.WithCustomMapping customMapping =
+        Advice.withCustomMapping()
+            // required for AssignReturned annotation and throwable suppression
+            .with(new Advice.AssignReturned.Factory().withSuppressed(Throwable.class));
     AgentBuilder builder =
         new AgentBuilder.Default()
             .disableClassFormatChanges()
             .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
             .type(named(ExceptionHandlerTest.class.getName() + "$SomeClass"))
             .transform(
-                new AgentBuilder.Transformer.ForAdvice()
+                new AgentBuilder.Transformer.ForAdvice(customMapping)
                     .with(
                         new AgentBuilder.LocationStrategy.Simple(
                             ClassFileLocator.ForClassLoader.of(BadAdvice.class.getClassLoader())))
                     .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
                     .advice(isMethod().and(named("isInstrumented")), BadAdvice.class.getName()))
             .transform(
-                new AgentBuilder.Transformer.ForAdvice()
+                new AgentBuilder.Transformer.ForAdvice(customMapping)
                     .with(
                         new AgentBuilder.LocationStrategy.Simple(
                             ClassFileLocator.ForClassLoader.of(BadAdvice.class.getClassLoader())))
@@ -80,7 +86,10 @@ void exceptionHandlerInvoked() {
     int initLogEvents = testHandler.getRecords().size();
 
     // Triggers classload and instrumentation
-    assertThat(SomeClass.isInstrumented()).isTrue();
+    AtomicBoolean isInstrumented = new AtomicBoolean(false);
+    SomeClass.isInstrumented(isInstrumented);
+    assertThat(isInstrumented.get()).describedAs("method should have been instrumented").isTrue();
+
     assertThat(testHandler.getRecords())
         .hasSize(initLogEvents + 1)
         .last()
@@ -104,8 +113,10 @@ void exceptionOnNondelegatingClassloader() throws Exception {
 
     Class<?> someClazz = loader.loadClass(SomeClass.class.getName());
     assertThat(someClazz.getClassLoader()).isSameAs(loader);
-    someClazz.getMethod("isInstrumented").invoke(null);
+    AtomicBoolean instrumented = new AtomicBoolean(false);
+    someClazz.getMethod("isInstrumented", AtomicBoolean.class).invoke(null, instrumented);
     assertThat(testHandler.getRecords()).hasSize(initLogEvents);
+    assertThat(instrumented.get()).describedAs("method should have been instrumented").isTrue();
   }
 
   @Test
@@ -120,8 +131,8 @@ void exceptionHandlerSetsCorrectStackSize() {
 
   public static class SomeClass {
 
-    public static boolean isInstrumented() {
-      return false;
+    public static void isInstrumented(AtomicBoolean instrumented) {
+      instrumented.set(false);
     }
 
     public static void smallStack() {

javaagent-tooling/src/testMissingType/java/io/opentelemetry/javaagent/tooling/bytebuddy/MissingTypeTest.java

@@ -16,6 +16,7 @@
 import net.bytebuddy.agent.ByteBuddyAgent;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.scaffold.MethodGraph;
 import net.bytebuddy.utility.JavaModule;
@@ -33,6 +34,8 @@ class MissingTypeTest {
 
   @BeforeAll
   static void setUp() {
+    Advice.WithCustomMapping customMapping =
+        Advice.withCustomMapping().with(new Advice.AssignReturned.Factory());
     AgentBuilder builder =
         new AgentBuilder.Default(
                 new ByteBuddy().with(MethodGraph.Compiler.ForDeclaredMethods.INSTANCE))
@@ -54,7 +57,7 @@ public void onError(
                 })
             .type(named(MissingTypeTest.class.getName() + "$SomeClass"))
             .transform(
-                new AgentBuilder.Transformer.ForAdvice()
+                new AgentBuilder.Transformer.ForAdvice(customMapping)
                     .with(
                         new AgentBuilder.LocationStrategy.Simple(
                             ClassFileLocator.ForClassLoader.of(TestAdvice.class.getClassLoader())))

javaagent-tooling/src/testMissingType/java/io/opentelemetry/javaagent/tooling/bytebuddy/TestAdvice.java

@@ -6,12 +6,14 @@
 package io.opentelemetry.javaagent.tooling.bytebuddy;
 
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 
 @SuppressWarnings("unused")
 public class TestAdvice {
 
+  @AssignReturned.ToReturned
   @Advice.OnMethodExit
-  public static void returnTrue(@Advice.Return(readOnly = false) boolean returnVal) {
-    returnVal = true;
+  public static boolean returnTrue() {
+    return true;
   }
 }

testing-common/integration-tests/src/main/java/context/ContextTestInstrumentation.java

@@ -5,15 +5,17 @@
 
 package context;
 
+import static context.ContextTestSingletons.CONTEXT;
+import static context.ContextTestSingletons.INTEGER;
+import static context.ContextTestSingletons.INTERFACE_INTEGER;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 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 library.KeyClass;
-import library.KeyInterface;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -42,78 +44,71 @@ public void transform(TypeTransformer transformer) {
 
   @SuppressWarnings("unused")
   public static class MarkInstrumentedAdvice {
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit
-    public static void methodExit(@Advice.Return(readOnly = false) boolean isInstrumented) {
-      isInstrumented = true;
+    public static boolean methodExit() {
+      return true;
     }
   }
 
   @SuppressWarnings("unused")
   public static class StoreAndIncrementApiUsageAdvice {
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit
-    public static void methodExit(
-        @Advice.This KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) {
-
-      VirtualField<KeyClass, Context> virtualField =
-          VirtualField.find(KeyClass.class, Context.class);
-
-      Context context = virtualField.get(thiz);
+    public static int methodExit(@Advice.This KeyClass thiz) {
+      Context context = CONTEXT.get(thiz);
       if (context == null) {
         context = new Context();
-        virtualField.set(thiz, context);
+        CONTEXT.set(thiz, context);
       }
 
-      contextCount = ++context.count;
+      return ++context.count;
     }
   }
 
   @SuppressWarnings("unused")
   public static class GetApiUsageAdvice {
+
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static boolean methodEnter() {
+      // always skip original method body
+      return true;
+    }
+
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit
-    public static void methodExit(
-        @Advice.This KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) {
-      VirtualField<KeyClass, Context> virtualField =
-          VirtualField.find(KeyClass.class, Context.class);
-      Context context = virtualField.get(thiz);
-      contextCount = context == null ? 0 : context.count;
+    public static int methodExit(@Advice.This KeyClass thiz) {
+      Context context = CONTEXT.get(thiz);
+      return context == null ? 0 : context.count;
     }
   }
 
   @SuppressWarnings("unused")
   public static class PutApiUsageAdvice {
     @Advice.OnMethodExit
     public static void methodExit(@Advice.This KeyClass thiz, @Advice.Argument(0) int value) {
-      VirtualField<KeyClass, Context> virtualField =
-          VirtualField.find(KeyClass.class, Context.class);
       Context context = new Context();
       context.count = value;
-      virtualField.set(thiz, context);
+      CONTEXT.set(thiz, context);
     }
   }
 
   @SuppressWarnings("unused")
   public static class RemoveApiUsageAdvice {
     @Advice.OnMethodExit
     public static void methodExit(@Advice.This KeyClass thiz) {
-      VirtualField<KeyClass, Context> virtualField =
-          VirtualField.find(KeyClass.class, Context.class);
-      virtualField.set(thiz, null);
+      CONTEXT.set(thiz, null);
     }
   }
 
   @SuppressWarnings("unused")
   public static class UseMultipleFieldsAdvice {
     @Advice.OnMethodExit
     public static void methodExit(@Advice.This KeyClass thiz) {
-      VirtualField<KeyClass, Context> field1 = VirtualField.find(KeyClass.class, Context.class);
-      VirtualField<KeyClass, Integer> field2 = VirtualField.find(KeyClass.class, Integer.class);
-      VirtualField<KeyInterface, Integer> interfaceField =
-          VirtualField.find(KeyInterface.class, Integer.class);
-
-      Context context = field1.get(thiz);
+      Context context = CONTEXT.get(thiz);
       int count = context == null ? 0 : context.count;
-      field2.set(thiz, count);
-      interfaceField.set(thiz, count);
+      INTEGER.set(thiz, count);
+      INTERFACE_INTEGER.set(thiz, count);
     }
   }
 }

testing-common/integration-tests/src/main/java/context/ContextTestInstrumentationModule.java

@@ -10,21 +10,29 @@
 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 ContextTestInstrumentationModule extends InstrumentationModule {
+public class ContextTestInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public ContextTestInstrumentationModule() {
     super("context-test-instrumentation");
   }
 
   @Override
   public boolean isHelperClass(String className) {
-    return className.equals(getClass().getPackage().getName() + ".Context");
+    String pkg = getClass().getPackage().getName();
+    return className.equals(pkg + ".Context") || className.equals(pkg + ".ContextTestSingletons");
   }
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new ContextTestInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }

testing-common/integration-tests/src/main/java/context/ContextTestSingletons.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package context;
+
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import library.KeyClass;
+import library.KeyInterface;
+
+public class ContextTestSingletons {
+
+  public static final VirtualField<KeyClass, Context> CONTEXT =
+      VirtualField.find(KeyClass.class, Context.class);
+  public static final VirtualField<KeyClass, Integer> INTEGER =
+      VirtualField.find(KeyClass.class, Integer.class);
+  public static final VirtualField<KeyInterface, Integer> INTERFACE_INTEGER =
+      VirtualField.find(KeyInterface.class, Integer.class);
+
+  private ContextTestSingletons() {}
+}

testing-common/integration-tests/src/main/java/field/VirtualFieldTestInstrumentationModule.java

@@ -14,6 +14,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.util.List;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -50,10 +51,11 @@ public void transform(TypeTransformer transformer) {
 
   @SuppressWarnings("unused")
   public static class TestAdvice {
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit
-    public static void onExit(@Advice.Return(readOnly = false) boolean result) {
+    public static boolean onExit() {
       VirtualFieldTestHelper.test();
-      result = true;
+      return true;
     }
   }
 }

testing-common/integration-tests/src/main/java/helper/DuplicateHelperInstrumentation.java

@@ -10,6 +10,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -26,9 +27,10 @@ public void transform(TypeTransformer transformer) {
 
   @SuppressWarnings("unused")
   public static class TestAdvice {
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void addSuffix(@Advice.Return(readOnly = false) String string) {
-      string = DuplicateHelper.addSuffix(string, " foo");
+    public static String addSuffix(@Advice.Return String string) {
+      return DuplicateHelper.addSuffix(string, " foo");
     }
   }
 }

testing-common/integration-tests/src/main/java/io/opentelemetry/javaagent/IndyInstrumentationTestModule.java

@@ -176,12 +176,11 @@ public static void onMethodEnter() {
         throw new RuntimeException("This exception should be suppressed");
       }
 
-      @Advice.AssignReturned.ToReturned
       @Advice.OnMethodExit(
           suppress = Throwable.class,
           onThrowable = Throwable.class,
           inline = false)
-      public static void onMethodExit(@Advice.Thrown Throwable throwable) {
+      public static void onMethodExit() {
         throw new RuntimeException("This exception should be suppressed");
       }
     }

testing-common/integration-tests/src/test/java/context/FieldBackedImplementationTest.java

@@ -140,7 +140,7 @@ void instanceState(Class<KeyClass> keyClass) throws Exception {
     KeyClass instance2 = keyClass.getConstructor().newInstance();
 
     // correct api usage stores state in map
-    instance1.incrementContextCount();
+    assertThat(instance1.incrementContextCount()).isEqualTo(1);
 
     assertThat(instance1.incrementContextCount()).isEqualTo(2);
     assertThat(instance2.incrementContextCount()).isEqualTo(1);
@@ -159,7 +159,10 @@ void modifyInstanceState(Class<KeyClass> keyClass) throws Exception {
   @MethodSource("provideKeyClassParameters")
   void removeInstanceState(Class<KeyClass> keyClass) throws Exception {
     KeyClass instance1 = keyClass.getConstructor().newInstance();
+
     instance1.putContextCount(10);
+    assertThat(instance1.getContextCount()).isEqualTo(10);
+
     instance1.removeContextCount();
 
     assertThat(instance1.getContextCount()).isEqualTo(0);