Simplify InstanceStore implementation to use simple, synchronized HashMaps

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

@@ -1,25 +1,61 @@
 package datadog.trace.bootstrap;
 
 import datadog.trace.api.GenericClassValue;
-import java.util.function.Function;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * An {@code InstanceStore} is a class global map for registering instances. This can be useful when
- * helper classes are injected into multiple class loaders but need to share unique keys of a type
+ * helper classes are injected into multiple class loaders and need to share instances of a type
  * from a common parent class loader.
  *
- * <p>The {@code InstanceStore} is backed by a {@code WeakMapContextStore} that has a max size of
- * 100 keys.
+ * <p>Instance keys are expected to be string literals, defined as constants in the helper classes.
  */
-public final class InstanceStore {
-  private InstanceStore() {}
+public final class InstanceStore<T> implements ContextStore<String, T> {
 
   private static final ClassValue<? super ContextStore<String, ?>> classInstanceStore =
-      GenericClassValue.of(
-          (Function<Class<?>, ContextStore<String, ?>>) input -> new WeakMapContextStore<>(100));
+      GenericClassValue.of(input -> new InstanceStore<>());
 
+  /** @return global store of instances with the same common type */
   @SuppressWarnings("unchecked")
-  public static <T> ContextStore<String, T> of(Class<T> type) {
-    return (ContextStore<String, T>) classInstanceStore.get(type);
+  public static <T> InstanceStore<T> of(Class<T> type) {
+    return (InstanceStore<T>) 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
+  public T get(String key) {
+    return store.get(key);
+  }
+
+  @Override
+  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;
+  }
+
+  @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
+  public T remove(String key) {
+    return store.remove(key);
   }
 }

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

@@ -1,15 +1,11 @@
 package datadog.trace.bootstrap
 
-import datadog.trace.agent.tooling.WeakMaps
 import datadog.trace.test.util.DDSpecification
 import spock.lang.Shared
 
 import java.util.concurrent.atomic.AtomicInteger
 
 class InstanceStoreTest extends DDSpecification {
-  static {
-    WeakMaps.registerAsSupplier()
-  }
 
   @Shared
   private AtomicInteger counter = new AtomicInteger()
@@ -19,14 +15,6 @@ class InstanceStoreTest extends DDSpecification {
     "key-${counter.incrementAndGet()}"
   }
 
-  def "test empty InstanceStore"() {
-    setup:
-    WeakMapContextStore<String, Some> someStore = InstanceStore.of(Some) as WeakMapContextStore<String, Some>
-
-    expect:
-    someStore.size() == 0
-  }
-
   def "test returns existing value"() {
     setup:
     def someStore = InstanceStore.of(Some)