Allow adding multiple ContextStore fields to one key class, part 1

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation;
+
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+
+/**
+ * This interface allows registering {@link ContextStore} class pairs.
+ *
+ * <p>This interface should not be implemented by the javaagent extension developer - the javaagent
+ * will provide the implementation of all transformations described here.
+ */
+public interface InstrumentationContextBuilder {
+
+  /**
+   * Register the association between the {@code keyClassName} and the {@code contextClassName}.
+   * Class pairs registered using this method will be available as {@link ContextStore}s in the
+   * runtime; obtainable by calling {@link InstrumentationContext#get(Class, Class)}.
+   *
+   * @param keyClassName The name of the class that an instance of class named {@code
+   *     contextClassName} will be attached to.
+   * @param contextClassName The instrumentation context class name.
+   * @return {@code this}.
+   * @see InstrumentationContext
+   * @see ContextStore
+   */
+  InstrumentationContextBuilder register(String keyClassName, String contextClassName);
+}

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java

@@ -167,8 +167,7 @@ public Map<String, ClassRef> getMuzzleReferences() {
    * compilation. Those helpers will be injected into the application classloader.
    *
    * <p>The actual implementation of this method is generated automatically during compilation by
-   * the {@code io.opentelemetry.javaagent.muzzle.generation.collector.MuzzleCodeGenerationPlugin}
-   * ByteBuddy plugin.
+   * the {@code io.opentelemetry.instrumentation.javaagent-codegen} Gradle plugin.
    *
    * <p><b>This method is generated automatically</b>: if you override it, the muzzle compile plugin
    * will not generate a new implementation, it will leave the existing one.
@@ -182,13 +181,31 @@ public List<String> getMuzzleHelperClassNames() {
    * associated with a context class stored in the value.
    *
    * <p>The actual implementation of this method is generated automatically during compilation by
-   * the {@code io.opentelemetry.javaagent.muzzle.generation.collector.MuzzleCodeGenerationPlugin}
-   * ByteBuddy plugin.
+   * the {@code io.opentelemetry.instrumentation.javaagent-codegen} Gradle plugin.
    *
    * <p><b>This method is generated automatically</b>: if you override it, the muzzle compile plugin
    * will not generate a new implementation, it will leave the existing one.
+   *
+   * @deprecated Use {@link #registerMuzzleContextStoreClasses(InstrumentationContextBuilder)}
+   *     instead.
    */
+  @Deprecated
   public Map<String, String> getMuzzleContextStoreClasses() {
     return Collections.emptyMap();
   }
+
+  /**
+   * Builds the associations between instrumented library classes and instrumentation context
+   * classes. Keys (and their subclasses) will be associated with a context class stored in the
+   * value.
+   *
+   * <p>The actual implementation of this method is generated automatically during compilation by
+   * the {@code io.opentelemetry.instrumentation.javaagent-codegen} Gradle plugin.
+   *
+   * <p><b>This method is generated automatically</b>: if you override it, the muzzle compile plugin
+   * will not generate a new implementation, it will leave the existing one.
+   */
+  public void registerMuzzleContextStoreClasses(InstrumentationContextBuilder builder) {
+    getMuzzleContextStoreClasses().forEach(builder::register);
+  }
 }

javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java

@@ -20,10 +20,10 @@ private InstrumentationContext() {}
    *
    * <p>When this method is called from outside of an Advice class it can only access {@link
    * ContextStore} when it is already created. For this {@link ContextStore} either needs to be
-   * added to {@code
-   * io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule#getMuzzleContextStoreClasses()}
+   * registered in {@code
+   * io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule#registerMuzzleContextStoreClasses(InstrumentationContextBuilder)}
    * or be used in an Advice or Helper class which automatically adds it to {@code
-   * InstrumentationModule#getMuzzleContextStoreClasses()}
+   * InstrumentationModule#registerMuzzleContextStoreClasses(InstrumentationContextBuilder)}.
    *
    * @param keyClass The key class context is attached to.
    * @param contextClass The context class attached to the user class.

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/ContextStoreMappings.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.context;
+
+import java.util.AbstractMap;
+import java.util.Map;
+import java.util.Set;
+
+public class ContextStoreMappings {
+  private final Set<Map.Entry<String, String>> entrySet;
+
+  public ContextStoreMappings(Set<Map.Entry<String, String>> entrySet) {
+    this.entrySet = entrySet;
+  }
+
+  public int size() {
+    return entrySet.size();
+  }
+
+  public boolean isEmpty() {
+    return entrySet.isEmpty();
+  }
+
+  public boolean hasMapping(String keyClassName, String contextClassName) {
+    return entrySet.contains(
+        new AbstractMap.SimpleImmutableEntry<>(keyClassName, contextClassName));
+  }
+
+  public Set<Map.Entry<String, String>> entrySet() {
+    return entrySet;
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java

@@ -113,7 +113,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
 
   private final Class<?> instrumenterClass;
   private final ByteBuddy byteBuddy;
-  private final Map<String, String> contextStore;
+  private final ContextStoreMappings contextStore;
 
   // fields-accessor-interface-name -> fields-accessor-interface-dynamic-type
   private final Map<String, DynamicType.Unloaded<?>> fieldAccessorInterfaces;
@@ -127,7 +127,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
 
   private final Instrumentation instrumentation;
 
-  public FieldBackedProvider(Class<?> instrumenterClass, Map<String, String> contextStore) {
+  public FieldBackedProvider(Class<?> instrumenterClass, ContextStoreMappings contextStore) {
     this.instrumenterClass = instrumenterClass;
     this.contextStore = contextStore;
     // This class is used only when running with javaagent, thus this calls is safe
@@ -239,13 +239,11 @@ public void visitMethodInsn(
                               "Incorrect Context Api Usage detected. Cannot find map holder class for %s context %s. Was that class defined in contextStore for instrumentation %s?",
                               keyClassName, contextClassName, instrumenterClass.getName()));
                     }
-                    if (!contextClassName.equals(contextStore.get(keyClassName))) {
+                    if (!contextStore.hasMapping(keyClassName, contextClassName)) {
                       throw new IllegalStateException(
                           String.format(
-                              "Incorrect Context Api Usage detected. Incorrect context class %s, expected %s for instrumentation %s",
-                              contextClassName,
-                              contextStore.get(keyClassName),
-                              instrumenterClass.getName()));
+                              "Incorrect Context Api Usage detected. Incorrect context class %s for instrumentation %s",
+                              contextClassName, instrumenterClass.getName()));
                     }
                     // stack: contextClass | keyClass
                     mv.visitMethodInsn(

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/InstrumentationContextBuilderImpl.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.context;
+
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationContextBuilder;
+import java.util.AbstractMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public class InstrumentationContextBuilderImpl implements InstrumentationContextBuilder {
+  private final Set<Map.Entry<String, String>> entrySet = new HashSet<>();
+
+  @Override
+  public InstrumentationContextBuilder register(String keyClassName, String contextClassName) {
+    entrySet.add(new AbstractMap.SimpleImmutableEntry<>(keyClassName, contextClassName));
+    return this;
+  }
+
+  public ContextStoreMappings build() {
+    return new ContextStoreMappings(entrySet);
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java

@@ -17,7 +17,9 @@
 import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.bytebuddy.LoggingFailSafeMatcher;
+import io.opentelemetry.javaagent.tooling.context.ContextStoreMappings;
 import io.opentelemetry.javaagent.tooling.context.FieldBackedProvider;
+import io.opentelemetry.javaagent.tooling.context.InstrumentationContextBuilderImpl;
 import io.opentelemetry.javaagent.tooling.context.InstrumentationContextProvider;
 import io.opentelemetry.javaagent.tooling.context.NoopContextProvider;
 import io.opentelemetry.javaagent.tooling.muzzle.HelperResourceBuilderImpl;
@@ -26,7 +28,6 @@
 import java.lang.instrument.Instrumentation;
 import java.security.ProtectionDomain;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.annotation.AnnotationSource;
@@ -118,9 +119,11 @@ AgentBuilder install(
 
   private static InstrumentationContextProvider createInstrumentationContextProvider(
       InstrumentationModule instrumentationModule) {
-    Map<String, String> contextStore = instrumentationModule.getMuzzleContextStoreClasses();
-    if (!contextStore.isEmpty()) {
-      return FieldBackedProviderFactory.get(instrumentationModule.getClass(), contextStore);
+    InstrumentationContextBuilderImpl builder = new InstrumentationContextBuilderImpl();
+    instrumentationModule.registerMuzzleContextStoreClasses(builder);
+    ContextStoreMappings mappings = builder.build();
+    if (!mappings.isEmpty()) {
+      return FieldBackedProviderFactory.get(instrumentationModule.getClass(), mappings);
     } else {
       return NoopContextProvider.INSTANCE;
     }
@@ -132,8 +135,9 @@ private static class FieldBackedProviderFactory {
           (keyClass, contextClass) -> FieldBackedProvider.getContextStore(keyClass, contextClass));
     }
 
-    static FieldBackedProvider get(Class<?> instrumenterClass, Map<String, String> contextStore) {
-      return new FieldBackedProvider(instrumenterClass, contextStore);
+    static FieldBackedProvider get(
+        Class<?> instrumenterClass, ContextStoreMappings contextStoreMappings) {
+      return new FieldBackedProvider(instrumenterClass, contextStoreMappings);
     }
   }