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

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

@@ -18,7 +18,12 @@ public class InstrumentationContext {
    * <p>In reality, the <em>calls</em> to this method are re-written to something more performant
    * while injecting advice into a method.
    *
-   * <p>This method must only be called within an Advice class.
+   * <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()}
+   * 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 <Q extends K, K, C> ContextStore<Q, C> get(
       Class<K> keyClass, Class<C> 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<Q extends K, K, C> {
+    ContextStore<Q, C> get(Class<K> keyClass, Class<C> 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;
   }
 }

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 <Q extends K, K, C> ContextStore<Q, C> getContextStore(
+      Class<K> keyClass, Class<C> 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<Q, C>) 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)
         + "$"

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<String, String> 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<String, String> 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