KDWSS
diff --git a/‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/ContextStore.java‎
Lines changed: 59 additions & 0 deletions b/‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/ContextStore.java‎
Lines changed: 59 additions & 0 deletions
diff --git a/‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java‎
Lines changed: 10 additions & 27 deletions b/‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java‎
Lines changed: 10 additions & 27 deletions
@@ -0,0 +1,59 @@
+package datadog.trace.bootstrap;
+
+/**
+ * Interface to represent context storage for instrumentations.
+ *
+ * <p>Context instances are weakly referenced and will be garbage collected when their corresponding
+ * key instance is collected.
+ *
+ * @param <K> key type to do context lookups
+ * @param <C> context type
+ */
+public interface ContextStore<K, C> {
+
+  /**
+   * Factory interface to create context instances
+   *
+   * @param <C> context type
+   */
+  interface Factory<C> {
+
+    /** @return new context instance */
+    C create();
+  }
+
+  /**
+   * Get context given the key
+   *
+   * @param key the key to looup
+   * @return context object
+   */
+  C get(K key);
+
+  /**
+   * Put new context instance for given key
+   *
+   * @param key key to use
+   * @param context context instance to save
+   */
+  void put(K key, C context);
+
+  /**
+   * Put new context instance if key is absent
+   *
+   * @param key key to use
+   * @param context new context instance to put
+   * @return old instance if it was present, or new instance
+   */
+  C putIfAbsent(K key, C context);
+
+  /**
+   * Put new context instance if key is absent. Uses context factory to avoid creating objects if
+   * not needed.
+   *
+   * @param key key to use
+   * @param contextFactory factory instance to produce new context object
+   * @return old instance if it was present, or new instance
+   */
+  C putIfAbsent(K key, Factory<C> contextFactory);
+}
@@ -5,39 +5,22 @@ public class InstrumentationContext {
   private InstrumentationContext() {}
 
   /**
-   * Fetch a context instance out of the context store.
+   * Find a {@link ContextStore} instance for given key class and context class.
    *
-   * <p>
-   *
-   * <p>Conceptually, this can be thought of as a two pass map look up.
-   *
-   * <p>For example: <em>RunnableState runnableState = get(runnableImpl, Runnable.class,
-   * RunnableState.class)</em> --> <em>RunnableState runnableState = (RunnableState)
-   * GlobalContextMap.get(Runnable.class).get(runnableImpl)</em>
-   *
-   * <p>
+   * <p>Conceptually this can be thought of as a map lookup to fetch a second level map given
+   * keyClass.
    *
    * <p>However, the implementation is actually provided by bytecode transformation for performance
    * reasons.
    *
-   * <p>
-   *
-   * <p>Context classes are weakly referenced and will be garbage collected when their corresponding
-   * user instance is collected.
-   *
-   * <p>
-   *
-   * <p>Instrumenters making this call must define the user-context class relationship in
-   * datadog.trace.agent.tooling.Instrumenter.Default#contextStore.
-   *
-   * @param userInstance The instance to store context on.
-   * @param userClass The user class context is attached to.
+   * @param keyClass The key class context is attached to.
    * @param contextClass The context class attached to the user class.
-   * @param <K> user class
-   * @param <V> context class
-   * @return The context instance attached to userInstance.
+   * @param <K> key class
+   * @param <C> context class
+   * @return The instance of context store for given arguments.
    */
-  public static <K, V> V get(K userInstance, Class<K> userClass, Class<V> contextClass) {
-    throw new RuntimeException("calls to this method will be rewritten");
+  public static <K, C> ContextStore<K, C> get(
+      final Class<K> keyClass, final Class<C> contextClass) {
+    throw new RuntimeException("Calls to this method will be rewritten by MapBackedProvider");
   }
 }