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 fdf65823555..67612ac1eed 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 @@ -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,15 +13,16 @@ * *

Instance keys are expected to be string literals, defined as constants in the helper classes. */ -public final class InstanceStore implements ContextStore { +public final class InstanceStore { - private static final ClassValue> classInstanceStore = - GenericClassValue.of(input -> new InstanceStore<>()); + @SuppressWarnings("rawtypes") + private static final ClassValue classInstanceStore = + GenericClassValue.of(type -> new InstanceStore<>()); /** @return global store of instances with the same common type */ @SuppressWarnings("unchecked") public static InstanceStore of(Class type) { - return (InstanceStore) classInstanceStore.get(type); + return classInstanceStore.get(type); } // simple approach; instance stores don't need highly concurrent access or weak keys @@ -28,33 +30,45 @@ public static InstanceStore of(Class type) { 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 instanceFactory) { + return store.computeIfAbsent(key, k -> instanceFactory.get()); } - @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 + /** + * 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); } 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 eaac5757973..eb2dbdbf1db 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,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 { + static class Creator implements Supplier { private AtomicInteger invocations private Some some @@ -127,7 +135,7 @@ class InstanceStoreTest extends DDSpecification { } @Override - Some create() { + Some get() { invocations.incrementAndGet() return some } diff --git a/dd-java-agent/instrumentation/undertow/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java b/dd-java-agent/instrumentation/undertow/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java index dfa2bea25d0..ee98d2fb51a 100644 --- a/dd-java-agent/instrumentation/undertow/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java +++ b/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 attachmentStore = + private static final InstanceStore attachmentStore = InstanceStore.of(AttachmentKey.class); @SuppressWarnings("unchecked") public static final AttachmentKey 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 =