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> 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 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)