Simplify InstanceStore API

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java

@@ -4,6 +4,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * An {@code InstanceStore} is a class global map for registering instances. This can be useful when
@@ -12,49 +13,62 @@
  *
  * <p>Instance keys are expected to be string literals, defined as constants in the helper classes.
  */
-public final class InstanceStore<T> implements ContextStore<String, T> {
+public final class InstanceStore<T> {
 
-  private static final ClassValue<? super ContextStore<String, ?>> classInstanceStore =
-      GenericClassValue.of(input -> new InstanceStore<>());
+  @SuppressWarnings("rawtypes")
+  private static final ClassValue<InstanceStore> classInstanceStore =
+      GenericClassValue.of(type -> new InstanceStore<>());
 
   /** @return global store of instances with the same common type */
   @SuppressWarnings("unchecked")
   public static <T> InstanceStore<T> of(Class<T> type) {
-    return (InstanceStore<T>) classInstanceStore.get(type);
+    return classInstanceStore.get(type);
   }
 
   // simple approach; instance stores don't need highly concurrent access or weak keys
   private final Map<String, T> store = Collections.synchronizedMap(new HashMap<>());
 
   private InstanceStore() {}
 
-  @Override
+  /**
+   * Gets the instance of {@code T} currently associated with the given key.
+   *
+   * @param key the instance key
+   * @return the associated instance
+   */
   public T get(String key) {
     return store.get(key);
   }
 
-  @Override
+  /**
+   * Unconditionally associates an instance of {@code T} with the given key.
+   *
+   * @param key the instance key
+   * @param instance the instance
+   */
   public void put(String key, T instance) {
     store.put(key, instance);
   }
 
-  @Override
-  public T putIfAbsent(String key, T instance) {
-    T existing = store.putIfAbsent(key, instance);
-    return existing != null ? existing : instance;
+  /**
+   * If the given key is not already associated with an instance, create one using the factory and
+   * associate it. Unlike {@link java.util.Map#putIfAbsent} this always returns the final associated
+   * instance.
+   *
+   * @param key the instance key
+   * @param instanceFactory the factory to create instances
+   * @return final associated instance
+   */
+  public T putIfAbsent(String key, Supplier<T> instanceFactory) {
+    return store.computeIfAbsent(key, k -> instanceFactory.get());
   }
 
-  @Override
-  public T putIfAbsent(String key, Factory<T> instanceFactory) {
-    return store.computeIfAbsent(key, instanceFactory::create);
-  }
-
-  @Override
-  public T computeIfAbsent(String key, KeyAwareFactory<? super String, T> instanceFactory) {
-    return store.computeIfAbsent(key, instanceFactory::create);
-  }
-
-  @Override
+  /**
+   * Removes the instance of {@code T} currently associated with the given key.
+   *
+   * @param key the instance key
+   * @return the previously associated instance; {@code null} if there was none
+   */
   public T remove(String key) {
     return store.remove(key);
   }

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy

@@ -1,6 +1,7 @@
 package datadog.trace.bootstrap
 
 import datadog.trace.test.util.DDSpecification
+import java.util.function.Supplier
 import spock.lang.Shared
 
 import java.util.concurrent.atomic.AtomicInteger
@@ -15,27 +16,34 @@ class InstanceStoreTest extends DDSpecification {
     "key-${counter.incrementAndGet()}"
   }
 
-  def "test returns existing value"() {
+  def "test basic operation"() {
     setup:
     def someStore = InstanceStore.of(Some)
     def some1 = new Some()
     def some2 = new Some()
     def key = nextKey()
-    someStore.put(key, some1)
+    def current
 
     when:
-    def current = someStore.putIfAbsent(key, some2)
+    someStore.put(key, some1)
+    current = someStore.get(key)
 
     then:
     current == some1
-    current != some2
 
     when:
-    current = someStore.putIfAbsent(key, some2)
+    someStore.put(key, some2)
+    current = someStore.get(key)
 
     then:
-    current == some1
-    current != some2
+    current == some2
+
+    when:
+    someStore.remove(key)
+    current = someStore.get(key)
+
+    then:
+    current == null
   }
 
   def "test returns existing store"() {
@@ -45,13 +53,13 @@ class InstanceStoreTest extends DDSpecification {
     InstanceStore.of(Some).put(key, some1)
 
     when:
-    def current = InstanceStore.of(Some).putIfAbsent(key, new Some())
+    def current = InstanceStore.of(Some).putIfAbsent(key, Some::new)
 
     then:
     current == some1
 
     when:
-    current = InstanceStore.of(Some).putIfAbsent(key, new Some())
+    current = InstanceStore.of(Some).putIfAbsent(key, Some::new)
 
     then:
     current == some1
@@ -63,13 +71,13 @@ class InstanceStoreTest extends DDSpecification {
     def key = nextKey()
 
     when:
-    def current = someStore.putIfAbsent(key, some1)
+    def current = someStore.putIfAbsent(key, () -> some1)
 
     then:
     current == some1
 
     when:
-    current = someStore.putIfAbsent(key, new Some())
+    current = someStore.putIfAbsent(key, Some::new)
 
     then:
     current == some1
@@ -113,7 +121,7 @@ class InstanceStoreTest extends DDSpecification {
 
   static class Some {}
 
-  static class Creator implements ContextStore.Factory<Some> {
+  static class Creator implements Supplier<Some> {
     private AtomicInteger invocations
     private Some some
 
@@ -127,7 +135,7 @@ class InstanceStoreTest extends DDSpecification {
     }
 
     @Override
-    Some create() {
+    Some get() {
       invocations.incrementAndGet()
       return some
     }

dd-java-agent/instrumentation/undertow/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java

@@ -3,7 +3,6 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.gateway.BlockResponseFunction;
 import datadog.trace.api.naming.SpanNaming;
-import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstanceStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -23,13 +22,13 @@ public class UndertowDecorator
       UTF8BytesString.create("undertow-http-server");
 
   @SuppressWarnings("rawtypes")
-  private static final ContextStore<String, AttachmentKey> attachmentStore =
+  private static final InstanceStore<AttachmentKey> attachmentStore =
       InstanceStore.of(AttachmentKey.class);
 
   @SuppressWarnings("unchecked")
   public static final AttachmentKey<AgentScope.Continuation> DD_UNDERTOW_CONTINUATION =
       attachmentStore.putIfAbsent(
-          "DD_UNDERTOW_CONTINUATION", AttachmentKey.create(AgentScope.Continuation.class));
+          "DD_UNDERTOW_CONTINUATION", () -> AttachmentKey.create(AgentScope.Continuation.class));
 
   public static final UndertowDecorator DECORATE = new UndertowDecorator();
   public static final CharSequence UNDERTOW_REQUEST =