From 74c1f55e35fdeccf88b58cd0f160dc93f242401c Mon Sep 17 00:00:00 2001
From: Mateusz Rzeszutek <mrzeszutek@splunk.com>
Date: Tue, 14 Sep 2021 08:08:36 +0200
Subject: [PATCH] Allow adding multiple ContextStore fields to one key class,
 part 3 (#4110)

---
 gradle-plugins/build.gradle.kts               |  8 ++--
 .../generation/MuzzleCodeGenerator.java       | 45 ++++++++-----------
 2 files changed, 22 insertions(+), 31 deletions(-)

diff --git a/gradle-plugins/build.gradle.kts b/gradle-plugins/build.gradle.kts
index 02b4032ce125..06e20f6a990f 100644
--- 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/gradle-plugins/src/main/java/io/opentelemetry/javaagent/muzzle/generation/MuzzleCodeGenerator.java b/gradle-plugins/src/main/java/io/opentelemetry/javaagent/muzzle/generation/MuzzleCodeGenerator.java
index 265c6823fec2..ad078cbc09cf 100644
--- a/gradle-plugins/src/main/java/io/opentelemetry/javaagent/muzzle/generation/MuzzleCodeGenerator.java
+++ b/gradle-plugins/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();