From 667b87bac7b01b4378ce139c7b3362dd313b4501 Mon Sep 17 00:00:00 2001 From: Lauri Tulmin Date: Wed, 18 Aug 2021 10:36:04 +0300 Subject: [PATCH] Support looking up a ContextStore from outside of Advice (#3827) * Support looking up a ContextStore from outside of Advice * Add exception message * Move setting ContextStoreSupplier * Improve comment --- .../api/InstrumentationContext.java | 31 +++++++++++++++++-- .../tooling/context/FieldBackedProvider.java | 20 ++++++++++-- .../InstrumentationModuleInstaller.java | 14 ++++++++- 3 files changed, 59 insertions(+), 6 deletions(-) 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 9f80bc0656..6ed42aed9c 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 @@ -18,7 +18,12 @@ public class InstrumentationContext { *

In reality, the calls to this method are re-written to something more performant * while injecting advice into a method. * - *

This method must only be called within an Advice class. + *

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()} + * or be used in an Advice or Helper class which automatically adds it to {@code + * InstrumentationModule#getMuzzleContextStoreClasses()} * * @param keyClass The key class context is attached to. * @param contextClass The context class attached to the user class. @@ -28,7 +33,27 @@ public class InstrumentationContext { */ public static ContextStore get( Class keyClass, Class contextClass) { - throw new IllegalStateException( - "Calls to this method will be rewritten by Instrumentation Context Provider (e.g. FieldBackedProvider)"); + if (contextStoreSupplier == null) { + throw new IllegalStateException("Context store supplier not set"); + } + return contextStoreSupplier.get(keyClass, contextClass); + } + + public interface ContextStoreSupplier { + ContextStore get(Class keyClass, Class contextClass); + } + + private static volatile ContextStoreSupplier contextStoreSupplier; + + /** + * Sets the {@link ContextStoreSupplier} to execute when instrumentation needs to access {@link + * ContextStore}. This is called from the agent startup, instrumentation must not call this. + */ + public static void internalSetContextStoreSupplier(ContextStoreSupplier contextStoreSupplier) { + if (InstrumentationContext.contextStoreSupplier != null) { + // Only possible by misuse of this API, just ignore. + return; + } + InstrumentationContext.contextStoreSupplier = contextStoreSupplier; } } 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 a96b292e67..a4aaa2d591 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 @@ -21,6 +21,7 @@ import io.opentelemetry.javaagent.tooling.TransformSafeLogger; import io.opentelemetry.javaagent.tooling.Utils; import io.opentelemetry.javaagent.tooling.instrumentation.InstrumentationModuleInstaller; import java.lang.instrument.Instrumentation; +import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.security.ProtectionDomain; import java.util.Arrays; @@ -139,6 +140,21 @@ public class FieldBackedProvider implements InstrumentationContextProvider { bootstrapHelperInjector(contextStoreImplementations.values()); } + public static ContextStore getContextStore( + Class keyClass, Class contextClass) { + try { + String contextStoreClassName = + getContextStoreImplementationClassName(keyClass.getName(), contextClass.getName()); + Class contextStoreClass = Class.forName(contextStoreClassName, false, null); + Method method = contextStoreClass.getMethod("getContextStore", Class.class, Class.class); + return (ContextStore) method.invoke(null, keyClass, contextClass); + } catch (ClassNotFoundException exception) { + throw new IllegalStateException("Context store not found", exception); + } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException exception) { + throw new IllegalStateException("Failed to get context store", exception); + } + } + @Override public AgentBuilder.Identified.Extendable instrumentationTransformer( AgentBuilder.Identified.Extendable builder) { @@ -989,10 +1005,10 @@ public class FieldBackedProvider implements InstrumentationContextProvider { return (builder, typeDescription, classLoader, module) -> builder.visit(visitor); } - private String getContextStoreImplementationClassName( + private static String getContextStoreImplementationClassName( String keyClassName, String contextClassName) { return DYNAMIC_CLASSES_PACKAGE - + getClass().getSimpleName() + + FieldBackedProvider.class.getSimpleName() + "$ContextStore$" + Utils.convertToInnerClassName(keyClassName) + "$" 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 097c403b4b..08461edcf5 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 @@ -12,6 +12,7 @@ import static net.bytebuddy.matcher.ElementMatchers.not; import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; +import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext; import io.opentelemetry.javaagent.tooling.HelperInjector; import io.opentelemetry.javaagent.tooling.TransformSafeLogger; import io.opentelemetry.javaagent.tooling.Utils; @@ -113,12 +114,23 @@ public final class InstrumentationModuleInstaller { InstrumentationModule instrumentationModule) { Map contextStore = instrumentationModule.getMuzzleContextStoreClasses(); if (!contextStore.isEmpty()) { - return new FieldBackedProvider(instrumentationModule.getClass(), contextStore); + return FieldBackedProviderFactory.get(instrumentationModule.getClass(), contextStore); } else { return NoopContextProvider.INSTANCE; } } + private static class FieldBackedProviderFactory { + static { + InstrumentationContext.internalSetContextStoreSupplier( + (keyClass, contextClass) -> FieldBackedProvider.getContextStore(keyClass, contextClass)); + } + + static FieldBackedProvider get(Class instrumenterClass, Map contextStore) { + return new FieldBackedProvider(instrumenterClass, contextStore); + } + } + /** * A ByteBuddy matcher that decides whether this instrumentation should be applied. Calls * generated {@link ReferenceMatcher}: if any mismatch with the passed {@code classLoader} is