Allow adding multiple ContextStore fields to one key class, part 1
ADD
javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder.java
+/*
+ * 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);
+}

MODIFY
javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java
javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java
@@ -167,8 +167,7 @@ public abstract class InstrumentationModule implements Ordered {
* 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 abstract class InstrumentationModule implements Ordered {
* 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);
+ }
}

MODIFY
javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java
javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java
@@ -20,10 +20,10 @@ public class 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.

ADD
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/ContextStoreMappings.java
+/*
+ * 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;
+ }
+}

MODIFY
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java
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 class FieldBackedProvider implements InstrumentationContextProvider {
"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(

ADD
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/InstrumentationContextBuilderImpl.java
+/*
+ * 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);
+ }
+}

MODIFY
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java
@@ -17,7 +17,9 @@ import io.opentelemetry.javaagent.tooling.HelperInjector;
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 io.opentelemetry.javaagent.tooling.muzzle.ReferenceMatcher;
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 @@ public final class InstrumentationModuleInstaller {
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 @@ public final class InstrumentationModuleInstaller {
(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);
}
}

Allow adding multiple ContextStore fields to one key class, part 1 (#4067)