diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java
index 3db9546e051..fdf65823555 100644
--- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java
+++ b/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.
*
- *
The {@code InstanceStore} is backed by a {@code WeakMapContextStore} that has a max size of
- * 100 keys.
+ *
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 implements ContextStore {
private static final ClassValue super ContextStore> classInstanceStore =
- GenericClassValue.of(
- (Function, ContextStore>) input -> new WeakMapContextStore<>(100));
+ GenericClassValue.of(input -> new InstanceStore<>());
+ /** @return global store of instances with the same common type */
@SuppressWarnings("unchecked")
- public static ContextStore of(Class type) {
- return (ContextStore) classInstanceStore.get(type);
+ public static InstanceStore of(Class type) {
+ return (InstanceStore) classInstanceStore.get(type);
+ }
+
+ // simple approach; instance stores don't need highly concurrent access or weak keys
+ private final Map 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 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);
}
}
diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy
index d3cc5ae54ca..eaac5757973 100644
--- a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy
+++ b/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 someStore = InstanceStore.of(Some) as WeakMapContextStore
-
- expect:
- someStore.size() == 0
- }
-
def "test returns existing value"() {
setup:
def someStore = InstanceStore.of(Some)