Allow adding multiple ContextStore fields to one key class, part 3 (#…

…4110)
open-telemetry · Sep 14, 2021 · 74c1f55 · 74c1f55
1 parent 473f16f
commit 74c1f55
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 31 deletions.
diff --git a/gradle-plugins/build.gradle.kts b/gradle-plugins/build.gradle.kts
@@ -22,9 +22,9 @@ dependencies {
   implementation("com.google.guava:guava:30.1.1-jre")
   implementation("net.bytebuddy:byte-buddy-gradle-plugin:1.11.2")
 
-  implementation("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api:1.5.0-alpha")
-  implementation("io.opentelemetry.javaagent:opentelemetry-muzzle:1.5.0-alpha")
-  implementation("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api:1.5.0-alpha")
+  implementation("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api:1.6.0-alpha-SNAPSHOT")
+  implementation("io.opentelemetry.javaagent:opentelemetry-muzzle:1.6.0-alpha-SNAPSHOT")
+  implementation("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api:1.6.0-alpha-SNAPSHOT")
 
   implementation("org.eclipse.aether:aether-connector-basic:1.1.0")
   implementation("org.eclipse.aether:aether-transport-http:1.1.0")
@@ -37,7 +37,7 @@ dependencies {
   testImplementation("org.junit.jupiter:junit-jupiter-params")
   testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine")
 
-  testImplementation("io.opentelemetry.javaagent:opentelemetry-javaagent-instrumentation-api:1.5.0-alpha")
+  testImplementation("io.opentelemetry.javaagent:opentelemetry-javaagent-instrumentation-api:1.6.0-alpha-SNAPSHOT")
 }
 
 tasks.withType<Test>().configureEach {

diff --git a/...ugins/src/main/java/io/opentelemetry/javaagent/muzzle/generation/MuzzleCodeGenerator.java b/...ugins/src/main/java/io/opentelemetry/javaagent/muzzle/generation/MuzzleCodeGenerator.java
@@ -13,6 +13,7 @@
 import io.opentelemetry.javaagent.extension.muzzle.Flag;
 import io.opentelemetry.javaagent.extension.muzzle.MethodRef;
 import io.opentelemetry.javaagent.extension.muzzle.Source;
+import io.opentelemetry.javaagent.tooling.muzzle.ContextStoreMappings;
 import io.opentelemetry.javaagent.tooling.muzzle.HelperResource;
 import io.opentelemetry.javaagent.tooling.muzzle.HelperResourceBuilderImpl;
 import io.opentelemetry.javaagent.tooling.muzzle.ReferenceCollector;
@@ -52,7 +53,7 @@ final class MuzzleCodeGenerator implements AsmVisitorWrapper {
   private static final String MUZZLE_REFERENCES_METHOD_NAME = "getMuzzleReferences";
   private static final String MUZZLE_HELPER_CLASSES_METHOD_NAME = "getMuzzleHelperClassNames";
   private static final String MUZZLE_CONTEXT_STORE_CLASSES_METHOD_NAME =
-      "getMuzzleContextStoreClasses";
+      "registerMuzzleContextStoreClasses";
   private final URLClassLoader classLoader;
 
   public MuzzleCodeGenerator(URLClassLoader classLoader) {
@@ -510,50 +511,40 @@ private void generateMuzzleHelperClassNamesMethod(ReferenceCollector collector)
 
     private void generateMuzzleContextStoreClassesMethod(ReferenceCollector collector) {
       /*
-       * public Map<String, String> getMuzzleContextStoreClasses() {
-       *   Map<String, String> contextStore = new HashMap<>(...);
-       *   contextStore.put(..., ...);
-       *   return contextStore;
+       * public void registerMuzzleContextStoreClasses(InstrumentationContextBuilder builder) {
+       *   builder.register(..., ...);
        * }
        */
       MethodVisitor mv =
           super.visitMethod(
               Opcodes.ACC_PUBLIC,
               MUZZLE_CONTEXT_STORE_CLASSES_METHOD_NAME,
-              "()Ljava/util/Map;",
+              "(Lio/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder;)V",
               null,
               null);
       mv.visitCode();
 
-      Map<String, String> contextStoreClasses = collector.getContextStoreClasses();
+      ContextStoreMappings contextStoreMappings = collector.getContextStoreMappings();
 
-      writeNewMap(mv, contextStoreClasses.size());
-      // stack: map
-      mv.visitVarInsn(Opcodes.ASTORE, 1);
-      // stack: <empty>
-
-      contextStoreClasses.forEach(
+      mv.visitVarInsn(Opcodes.ALOAD, 1);
+      // stack: builder
+      contextStoreMappings.forEach(
           (className, contextClassName) -> {
-            mv.visitVarInsn(Opcodes.ALOAD, 1);
-            // stack: map
             mv.visitLdcInsn(className);
-            // stack: map, className
+            // stack: builder, className
             mv.visitLdcInsn(contextClassName);
-            // stack: map, className, contextClassName
+            // stack: builder, className, contextClassName
             mv.visitMethodInsn(
                 Opcodes.INVOKEINTERFACE,
-                "java/util/Map",
-                "put",
-                "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;",
+                "io/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder",
+                "register",
+                "(Ljava/lang/String;Ljava/lang/String;)Lio/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder;",
                 /* isInterface= */ true);
-            // stack: previousValue
-            mv.visitInsn(Opcodes.POP);
-            // stack: <empty>
+            // stack: builder
           });
-
-      mv.visitVarInsn(Opcodes.ALOAD, 1);
-      // stack: map
-      mv.visitInsn(Opcodes.ARETURN);
+      mv.visitInsn(Opcodes.POP);
+      // stack: <empty>
+      mv.visitInsn(Opcodes.RETURN);
 
       mv.visitMaxs(0, 0);
       mv.visitEnd();