Make it possible to use InstrumentationContext (now VirtualField) fro…

docs/contributing/writing-instrumentation-module.md

@@ -31,7 +31,8 @@ class MyLibraryInstrumentationModule extends InstrumentationModule {
 An `InstrumentationModule` needs to have at least one name. The user of the javaagent can
 [suppress a chosen instrumentation](../suppressing-instrumentation.md) by referring to it by one of
 its names. The instrumentation module names use kebab-case. The main instrumentation name (the first
-one) is supposed to be the same as the gradle module name (excluding the version suffix if it has one).
+one) is supposed to be the same as the gradle module name (excluding the version suffix if it has
+one).
 
 ```java
 public MyLibraryInstrumentationModule() {
@@ -104,8 +105,8 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
 }
 ```
 
-The above example will skip instrumenting the application code if it does not contain the class that was
-introduced in the version your instrumentation covers.
+The above example will skip instrumenting the application code if it does not contain the class that
+was introduced in the version your instrumentation covers.
 
 ### `typeInstrumentations()`
 
@@ -174,8 +175,8 @@ of available transformations that you can apply:
 
 * Calling `applyAdviceToMethod(ElementMatcher<? super MethodDescription>, String)` allows you to
   apply an advice class (the second parameter) to all matching methods (the first parameter). It is
-  suggested to make the method matchers as strict as possible - the type instrumentation should
-  only instrument the code that it's supposed to, not more.
+  suggested to make the method matchers as strict as possible - the type instrumentation should only
+  instrument the code that it's supposed to, not more.
 * `applyTransformer(AgentBuilder.Transformer)` allows you to inject an arbitrary ByteBuddy
   transformer. This is an advanced, low-level option that will not be subjected to muzzle safety
   checks and helper class detection - use it responsibly.
@@ -250,9 +251,9 @@ Exceptions thrown by the advice methods will get caught and handled by a special
 that OpenTelemetry javaagent defines. The handler makes sure to properly log all unexpected
 exceptions.
 
-The `OnMethodEnter` and `OnMethodExit` advice methods often need to share several pieces
-of information. We use local variables prefixed with `otel` to pass context, scope (and sometimes
-more) between those methods.
+The `OnMethodEnter` and `OnMethodExit` advice methods often need to share several pieces of
+information. We use local variables prefixed with `otel` to pass context, scope (and sometimes more)
+between those methods.
 
 ```java
 @Advice.OnMethodEnter(suppress = Throwable.class)
@@ -305,18 +306,17 @@ for accessing these default methods from advice.
 In fact, we suggest avoiding Java 8 language features in advice classes at all - sometimes you don't
 know what bytecode version is used by the instrumented class.
 
-Sometimes there is a need to associate some context class with an instrumented library class,
-and the library does not offer a way to do this. The OpenTelemetry javaagent provides the
-`ContextStore` for that purpose:
+Sometimes there is a need to associate some context class with an instrumented library class, and
+the library does not offer a way to do this. The OpenTelemetry javaagent provides the
+`VirtualField` for that purpose:
 
 ```java
-ContextStore<Runnable, Context> contextStore =
-    InstrumentationContext.get(Runnable.class, Context.class);
+VirtualField<Runnable, Context> virtualField =
+    VirtualField.get(Runnable.class, Context.class);
 ```
 
-A `ContextStore` is conceptually very similar to a map. It is not a simple map though:
-the javaagent uses a lot of bytecode modification magic to make this optimal.
-Because of this, retrieving a `ContextStore` instance is rather limited:
-the `InstrumentationContext#get()` method can only be called in advice classes, and it MUST receive
-class references as its parameters - it won't work with variables, method params etc.
+A `VirtualField` is conceptually very similar to a map. It is not a simple map though:
+the javaagent uses a lot of bytecode modification magic to make this optimal. Because of this,
+retrieving a `VirtualField` instance is rather limited: the `VirtualField#get()`
+MUST receive class references as its parameters - it won't work with variables, method params etc.
 Both the key class and the context class must be known at compile time for it to work.

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/field/VirtualField.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.field;
+
+import io.opentelemetry.instrumentation.api.internal.RuntimeVirtualFieldSupplier;
+import java.util.function.Supplier;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Represents a "virtual" field of type {@code F} that is added to type {@code T} in the runtime.
+ *
+ * <p>A virtual field has similar semantics to a weak-keys strong-values map: the value will be
+ * garbage collected when their owner instance is collected. It is discouraged to use a virtual
+ * field for keeping values that might reference their key, as it may cause memory leaks.
+ *
+ * @param <T> The type that will contain the new virtual field.
+ * @param <F> The field type that'll be added to {@code T}.
+ */
+// we're using an abstract class here so that we can call static find() in pre-jdk8 advice classes
+public abstract class VirtualField<T, F> {
+
+  /**
+   * Finds a {@link VirtualField} instance for given {@code type} and {@code fieldType}.
+   *
+   * <p>Conceptually this can be thought of as a map lookup to fetch a second level map given {@code
+   * type}.
+   *
+   * <p>In runtime, when using the javaagent, the <em>calls</em> to this method are rewritten to
+   * something more performant while injecting advice into a method.
+   *
+   * <p>When using this method outside of Advice method, the {@link VirtualField} should be looked
+   * up once and stored in a field to avoid repeatedly calling this method.
+   *
+   * @param type The type that will contain the new virtual field.
+   * @param fieldType The field type that'll be added to {@code type}.
+   */
+  public static <U extends T, T, F> VirtualField<U, F> find(Class<T> type, Class<F> fieldType) {
+    return RuntimeVirtualFieldSupplier.get().find(type, fieldType);
+  }
+
+  /** Gets the value of this virtual field. */
+  @Nullable
+  public abstract F get(T object);
+
+  /** Sets the new value of this virtual field. */
+  public abstract void set(T object, @Nullable F fieldValue);
+
+  /** Sets the new value of this virtual field if the current value is {@code null}. */
+  public abstract void setIfNull(T object, F fieldValue);
+
+  /**
+   * Sets the new value of this virtual field if the current value is {@code null}.
+   *
+   * @return The old field value if it was present, or the result of evaluating passed {@code
+   *     fieldValueSupplier}.
+   */
+  public abstract F computeIfNull(T object, Supplier<F> fieldValueSupplier);
+}

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/internal/RuntimeVirtualFieldSupplier.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.internal;
+
+import io.opentelemetry.instrumentation.api.caching.Cache;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public final class RuntimeVirtualFieldSupplier {
+
+  private static final Logger logger = LoggerFactory.getLogger(RuntimeVirtualFieldSupplier.class);
+
+  public interface VirtualFieldSupplier {
+    <U extends T, T, F> VirtualField<U, F> find(Class<T> type, Class<F> fieldType);
+  }
+
+  private static final VirtualFieldSupplier DEFAULT = new CacheBasedVirtualFieldSupplier();
+
+  private static volatile VirtualFieldSupplier instance = DEFAULT;
+
+  public static void set(VirtualFieldSupplier virtualFieldSupplier) {
+    // only overwrite the default, cache-based supplier
+    if (instance != DEFAULT) {
+      logger.warn(
+          "Runtime VirtualField supplier has already been set up, further set() calls are ignored");
+      return;
+    }
+    instance = virtualFieldSupplier;
+  }
+
+  public static VirtualFieldSupplier get() {
+    return instance;
+  }
+
+  private static final class CacheBasedVirtualFieldSupplier
+      extends ClassValue<Map<Class<?>, VirtualField<?, ?>>> implements VirtualFieldSupplier {
+
+    @Override
+    public <U extends T, T, F> VirtualField<U, F> find(Class<T> type, Class<F> fieldType) {
+      return (VirtualField<U, F>)
+          get(type).computeIfAbsent(fieldType, k -> new CacheBasedVirtualField<>());
+    }
+
+    @Override
+    protected Map<Class<?>, VirtualField<?, ?>> computeValue(Class<?> type) {
+      return new ConcurrentHashMap<>();
+    }
+  }
+
+  private static final class CacheBasedVirtualField<T, F> extends VirtualField<T, F> {
+    private final Cache<T, F> cache = Cache.newBuilder().setWeakKeys().build();
+
+    @Override
+    public @Nullable F get(T object) {
+      return cache.get(object);
+    }
+
+    @Override
+    public void set(T object, @Nullable F fieldValue) {
+      if (fieldValue == null) {
+        cache.remove(object);
+      } else {
+        cache.put(object, fieldValue);
+      }
+    }
+
+    @Override
+    public void setIfNull(T object, F fieldValue) {
+      cache.computeIfAbsent(object, k -> fieldValue);
+    }
+
+    @Override
+    public F computeIfNull(T object, Supplier<F> fieldValueSupplier) {
+      return cache.computeIfAbsent(object, k -> fieldValueSupplier.get());
+    }
+  }
+
+  private RuntimeVirtualFieldSupplier() {}
+}

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/servlet/MappingResolver.java

@@ -142,9 +142,8 @@ public String getMapping() {
 
   /**
    * Factory interface for creating {@link MappingResolver} instances. The main reason this class is
-   * here is that we need to ensure that the class used for {@code InstrumentationContext} lookup is
-   * always the same. If we would use an injected class it could be different in different class
-   * loaders.
+   * here is that we need to ensure that the class used for {@code VirtualField} lookup is always
+   * the same. If we would use an injected class it could be different in different class loaders.
    */
   public interface Factory {
 

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaActorCellInstrumentation.java

@@ -11,10 +11,9 @@
 import akka.dispatch.Envelope;
 import akka.dispatch.sysmsg.SystemMessage;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.TaskAdviceHelper;
 import net.bytebuddy.asm.Advice;
@@ -43,9 +42,9 @@ public static class InvokeAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Scope enter(@Advice.Argument(0) Envelope envelope) {
-      ContextStore<Envelope, PropagatedContext> contextStore =
-          InstrumentationContext.get(Envelope.class, PropagatedContext.class);
-      return TaskAdviceHelper.makePropagatedContextCurrent(contextStore, envelope);
+      VirtualField<Envelope, PropagatedContext> virtualField =
+          VirtualField.find(Envelope.class, PropagatedContext.class);
+      return TaskAdviceHelper.makePropagatedContextCurrent(virtualField, envelope);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -61,9 +60,9 @@ public static class SystemInvokeAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Scope enter(@Advice.Argument(0) SystemMessage systemMessage) {
-      ContextStore<SystemMessage, PropagatedContext> contextStore =
-          InstrumentationContext.get(SystemMessage.class, PropagatedContext.class);
-      return TaskAdviceHelper.makePropagatedContextCurrent(contextStore, systemMessage);
+      VirtualField<SystemMessage, PropagatedContext> virtualField =
+          VirtualField.find(SystemMessage.class, PropagatedContext.class);
+      return TaskAdviceHelper.makePropagatedContextCurrent(virtualField, systemMessage);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDefaultSystemMessageQueueInstrumentation.java

@@ -12,10 +12,9 @@
 
 import akka.dispatch.sysmsg.SystemMessage;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorAdviceHelper;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
@@ -50,9 +49,9 @@ public static class DispatchSystemAdvice {
     public static PropagatedContext enter(@Advice.Argument(1) SystemMessage systemMessage) {
       Context context = Java8BytecodeBridge.currentContext();
       if (ExecutorAdviceHelper.shouldPropagateContext(context, systemMessage)) {
-        ContextStore<SystemMessage, PropagatedContext> contextStore =
-            InstrumentationContext.get(SystemMessage.class, PropagatedContext.class);
-        return ExecutorAdviceHelper.attachContextToTask(context, contextStore, systemMessage);
+        VirtualField<SystemMessage, PropagatedContext> virtualField =
+            VirtualField.find(SystemMessage.class, PropagatedContext.class);
+        return ExecutorAdviceHelper.attachContextToTask(context, virtualField, systemMessage);
       }
       return null;
     }

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDispatcherInstrumentation.java

@@ -10,10 +10,9 @@
 
 import akka.dispatch.Envelope;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorAdviceHelper;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
@@ -44,9 +43,9 @@ public static class DispatchEnvelopeAdvice {
     public static PropagatedContext enterDispatch(@Advice.Argument(1) Envelope envelope) {
       Context context = Java8BytecodeBridge.currentContext();
       if (ExecutorAdviceHelper.shouldPropagateContext(context, envelope)) {
-        ContextStore<Envelope, PropagatedContext> contextStore =
-            InstrumentationContext.get(Envelope.class, PropagatedContext.class);
-        return ExecutorAdviceHelper.attachContextToTask(context, contextStore, envelope);
+        VirtualField<Envelope, PropagatedContext> virtualField =
+            VirtualField.find(Envelope.class, PropagatedContext.class);
+        return ExecutorAdviceHelper.attachContextToTask(context, virtualField, envelope);
       }
       return null;
     }

instrumentation/akka-actor-fork-join-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaForkJoinPoolInstrumentation.java

@@ -11,10 +11,9 @@
 
 import akka.dispatch.forkjoin.ForkJoinTask;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorAdviceHelper;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
@@ -54,9 +53,9 @@ public static PropagatedContext enterJobSubmit(
         @Advice.Argument(value = 0, readOnly = false) ForkJoinTask<?> task) {
       Context context = Java8BytecodeBridge.currentContext();
       if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
-        ContextStore<ForkJoinTask<?>, PropagatedContext> contextStore =
-            InstrumentationContext.get(ForkJoinTask.class, PropagatedContext.class);
-        return ExecutorAdviceHelper.attachContextToTask(context, contextStore, task);
+        VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
+            VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
+        return ExecutorAdviceHelper.attachContextToTask(context, virtualField, task);
       }
       return null;
     }

instrumentation/akka-actor-fork-join-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaForkJoinTaskInstrumentation.java

@@ -15,10 +15,9 @@
 import akka.dispatch.forkjoin.ForkJoinPool;
 import akka.dispatch.forkjoin.ForkJoinTask;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.TaskAdviceHelper;
 import java.util.concurrent.Callable;
@@ -63,14 +62,14 @@ public static class ForkJoinTaskAdvice {
      */
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Scope enter(@Advice.This ForkJoinTask<?> thiz) {
-      ContextStore<ForkJoinTask<?>, PropagatedContext> contextStore =
-          InstrumentationContext.get(ForkJoinTask.class, PropagatedContext.class);
-      Scope scope = TaskAdviceHelper.makePropagatedContextCurrent(contextStore, thiz);
+      VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
+          VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
+      Scope scope = TaskAdviceHelper.makePropagatedContextCurrent(virtualField, thiz);
       if (thiz instanceof Runnable) {
-        ContextStore<Runnable, PropagatedContext> runnableContextStore =
-            InstrumentationContext.get(Runnable.class, PropagatedContext.class);
+        VirtualField<Runnable, PropagatedContext> runnableVirtualField =
+            VirtualField.find(Runnable.class, PropagatedContext.class);
         Scope newScope =
-            TaskAdviceHelper.makePropagatedContextCurrent(runnableContextStore, (Runnable) thiz);
+            TaskAdviceHelper.makePropagatedContextCurrent(runnableVirtualField, (Runnable) thiz);
         if (null != newScope) {
           if (null != scope) {
             newScope.close();
@@ -80,10 +79,10 @@ public static Scope enter(@Advice.This ForkJoinTask<?> thiz) {
         }
       }
       if (thiz instanceof Callable) {
-        ContextStore<Callable<?>, PropagatedContext> callableContextStore =
-            InstrumentationContext.get(Callable.class, PropagatedContext.class);
+        VirtualField<Callable<?>, PropagatedContext> callableVirtualField =
+            VirtualField.find(Callable.class, PropagatedContext.class);
         Scope newScope =
-            TaskAdviceHelper.makePropagatedContextCurrent(callableContextStore, (Callable<?>) thiz);
+            TaskAdviceHelper.makePropagatedContextCurrent(callableVirtualField, (Callable<?>) thiz);
         if (null != newScope) {
           if (null != scope) {
             newScope.close();