diff --git a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder.java b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationContextBuilder.java new file mode 100644 index 000000000000..06403b90d6a1 --- /dev/null +++ b/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. + * + *

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); +} diff --git a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java index 4be9ab512eb3..c413cd597ea6 100644 --- a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java +++ b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java @@ -167,8 +167,7 @@ public Map getMuzzleReferences() { * compilation. Those helpers will be injected into the application classloader. * *

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. * *

This method is generated automatically: 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 getMuzzleHelperClassNames() { * associated with a context class stored in the value. * *

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. * *

This method is generated automatically: 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 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. + * + *

The actual implementation of this method is generated automatically during compilation by + * the {@code io.opentelemetry.instrumentation.javaagent-codegen} Gradle plugin. + * + *

This method is generated automatically: 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); + } } diff --git a/javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java b/javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java index 6ed42aed9c98..9b27f026f49f 100644 --- a/javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java +++ b/javaagent-instrumentation-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/InstrumentationContext.java @@ -20,10 +20,10 @@ private InstrumentationContext() {} * *

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. diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/ContextStoreMappings.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/ContextStoreMappings.java new file mode 100644 index 000000000000..92e42a079af1 --- /dev/null +++ b/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> entrySet; + + public ContextStoreMappings(Set> 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> entrySet() { + return entrySet; + } +} diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java index 525e7485a70b..c7e2e54d1c02 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java +++ b/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 contextStore; + private final ContextStoreMappings contextStore; // fields-accessor-interface-name -> fields-accessor-interface-dynamic-type private final Map> fieldAccessorInterfaces; @@ -127,7 +127,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider { private final Instrumentation instrumentation; - public FieldBackedProvider(Class instrumenterClass, Map 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( diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/InstrumentationContextBuilderImpl.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/InstrumentationContextBuilderImpl.java new file mode 100644 index 000000000000..723e0409baec --- /dev/null +++ b/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> 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); + } +} diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java index 1183ff28dc97..ca549af6f384 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java +++ b/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 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 contextStore) { - return new FieldBackedProvider(instrumenterClass, contextStore); + static FieldBackedProvider get( + Class instrumenterClass, ContextStoreMappings contextStoreMappings) { + return new FieldBackedProvider(instrumenterClass, contextStoreMappings); } }