diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java index 896342058a2..7d334d5405b 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java @@ -30,7 +30,7 @@ public void put(final Object key, final Object context) { } @Override - public Object putIfAbsent(final Object key, final Object context) { + public Object getOrPut(final Object key, final Object context) { if (key instanceof FieldBackedContextAccessor) { final FieldBackedContextAccessor accessor = (FieldBackedContextAccessor) key; Object existingContext = accessor.$get$__datadogContext$(storeId); @@ -45,18 +45,12 @@ public Object putIfAbsent(final Object key, final Object context) { } return existingContext; } else { - return weakStore().putIfAbsent(key, context); + return weakStore().getOrPut(key, context); } } @Override - public Object putIfAbsent(final Object key, final Factory contextFactory) { - return computeIfAbsent(key, contextFactory); - } - - @Override - public Object computeIfAbsent( - Object key, KeyAwareFactory contextFactory) { + public Object getOrCompute(Object key, KeyAwareFactory contextFactory) { if (key instanceof FieldBackedContextAccessor) { final FieldBackedContextAccessor accessor = (FieldBackedContextAccessor) key; Object existingContext = accessor.$get$__datadogContext$(storeId); @@ -71,7 +65,7 @@ public Object computeIfAbsent( } return existingContext; } else { - return weakStore().computeIfAbsent(key, contextFactory); + return weakStore().getOrCompute(key, contextFactory); } } @@ -95,14 +89,14 @@ public Object remove(Object key) { } // only create WeakMap-based fall-back when we need it - private volatile WeakMapContextStore weakStore; + private volatile WeakMapPerStore weakStore; private final Object synchronizationInstance = new Object(); - WeakMapContextStore weakStore() { + WeakMapPerStore weakStore() { if (null == weakStore) { synchronized (synchronizationInstance) { if (null == weakStore) { - weakStore = new WeakMapContextStore<>(); + weakStore = new WeakMapPerStore<>(); } } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java index 529e475ab87..8ae3a3239c5 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java @@ -134,14 +134,4 @@ private static FieldBackedContextStore createStore(final int storeId) { } return store; } - - /** Injection helper that immediately delegates to the weak-map for the given context store. */ - public static Object weakGet(final Object key, final int storeId) { - return getContextStore(storeId).weakStore().get(key); - } - - /** Injection helper that immediately delegates to the weak-map for the given context store. */ - public static void weakPut(final Object key, final int storeId, final Object context) { - getContextStore(storeId).weakStore().put(key, context); - } } 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 765f0df6a6c..6f7e4f7b560 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 @@ -5,6 +5,7 @@ import java.util.HashMap; import java.util.Map; import java.util.function.Supplier; +import javax.annotation.Nullable; /** * An {@code InstanceStore} is a class global map for registering instances. This can be useful when @@ -36,8 +37,9 @@ private InstanceStore() {} * Gets the instance of {@code T} currently associated with the given key. * * @param key the instance key - * @return the associated instance + * @return the associated instance; {@code null} if there was none */ + @Nullable public T get(String key) { return store.get(key); } @@ -61,7 +63,7 @@ public void put(String key, T instance) { * @param instanceFactory the factory to create instances * @return final associated instance */ - public T putIfAbsent(String key, Supplier instanceFactory) { + public T getOrCreate(String key, Supplier instanceFactory) { return store.computeIfAbsent(key, k -> instanceFactory.get()); } @@ -71,6 +73,7 @@ public T putIfAbsent(String key, Supplier instanceFactory) { * @param key the instance key * @return the previously associated instance; {@code null} if there was none */ + @Nullable public T remove(String key) { return store.remove(key); } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java similarity index 56% rename from dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java rename to dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java index 1750b7f2b9c..36917246b1f 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java @@ -1,49 +1,53 @@ package datadog.trace.bootstrap; +import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStore; + import datadog.trace.api.internal.VisibleForTesting; +import datadog.trace.bootstrap.ContextStore.KeyAwareFactory; /** - * Weak {@link ContextStore} that acts as a fall-back when field-injection isn't possible. + * Weak "map-per-store" fall-back to track contexts when field-injection isn't possible. * *

This class should be created lazily because it uses weak maps with background cleanup. */ -final class WeakMapContextStore implements ContextStore { - private static final int DEFAULT_MAX_SIZE = 50_000; - - private final int maxSize; - private final WeakMap map = WeakMap.Supplier.newWeakMap(); +public final class WeakMapPerStore { - public WeakMapContextStore(int maxSize) { - this.maxSize = maxSize; + /** Injection helper that immediately delegates to the weak-map for the given context store. */ + public static Object get(final Object key, final int storeId) { + return getContextStore(storeId).weakStore().get(key); } - public WeakMapContextStore() { - this(DEFAULT_MAX_SIZE); + /** Injection helper that immediately delegates to the weak-map for the given context store. */ + public static void put(final Object key, final int storeId, final Object context) { + getContextStore(storeId).weakStore().put(key, context); } - @Override + private static final int MAX_SIZE = 50_000; + + private final WeakMap map = WeakMap.Supplier.newWeakMap(); + + WeakMapPerStore() {} + @SuppressWarnings("unchecked") - public V get(final K key) { + V get(final K key) { return (V) map.get(key); } - @Override - public void put(final K key, final V context) { - if (map.size() < maxSize) { + void put(final K key, final V context) { + if (map.size() < MAX_SIZE) { map.put(key, context); } } - @Override - public V putIfAbsent(final K key, final V context) { + V getOrPut(final K key, final V context) { V existingContext = get(key); if (null == existingContext) { // This whole part with using synchronized is only because // we want to avoid prematurely calling the factory if - // someone else is doing a putIfAbsent at the same time. + // someone else is doing a getOrPut at the same time. // There is still the possibility that there is a concurrent // call to put that will win, but that is indistinguishable - // from the put happening right after the putIfAbsent. + // from the put happening right after the getOrPut. synchronized (map) { existingContext = get(key); if (null == existingContext) { @@ -55,21 +59,15 @@ public V putIfAbsent(final K key, final V context) { return existingContext; } - @Override - public V putIfAbsent(final K key, final Factory contextFactory) { - return computeIfAbsent(key, contextFactory); - } - - @Override - public V computeIfAbsent(K key, KeyAwareFactory contextFactory) { + V getOrCompute(K key, KeyAwareFactory contextFactory) { V existingContext = get(key); if (null == existingContext) { // This whole part with using synchronized is only because // we want to avoid prematurely calling the factory if - // someone else is doing a putIfAbsent at the same time. + // someone else is doing a getOrCompute at the same time. // There is still the possibility that there is a concurrent // call to put that will win, but that is indistinguishable - // from the put happening right after the putIfAbsent. + // from the put happening right after the getOrCompute. synchronized (map) { existingContext = get(key); if (null == existingContext) { @@ -81,9 +79,8 @@ public V computeIfAbsent(K key, KeyAwareFactory contextFactory) { return existingContext; } - @Override @SuppressWarnings("unchecked") - public V remove(final K key) { + V remove(final K key) { return (V) map.remove(key); } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java index 0a2477a8b44..6767c55ac77 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java @@ -36,7 +36,7 @@ private ConcurrentState() {} public static ConcurrentState captureContinuation( ContextStore contextStore, K key, Context context) { if (shouldCapture(context)) { - final ConcurrentState state = contextStore.putIfAbsent(key, FACTORY); + final ConcurrentState state = contextStore.getOrCreate(key, FACTORY); if (!state.captureAndSetContinuation(context) && log.isDebugEnabled()) { log.debug( "continuation was already set for {} in context {}, no continuation captured.", diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java index 43ca4fdf7bf..cb2d682aa75 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java @@ -41,7 +41,7 @@ public static boolean shouldAttachStateToTask(final Object task, final Context c */ public static State setupState( final ContextStore contextStore, final T task, final Context context) { - final State state = contextStore.putIfAbsent(task, State.FACTORY); + final State state = contextStore.getOrCreate(task, State.FACTORY); if (!state.captureAndSetContinuation(context)) { log.debug( "continuation was already set for {} in context {}, no continuation captured.", 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 eb2dbdbf1db..6b9a3fee93f 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 @@ -53,13 +53,13 @@ class InstanceStoreTest extends DDSpecification { InstanceStore.of(Some).put(key, some1) when: - def current = InstanceStore.of(Some).putIfAbsent(key, Some::new) + def current = InstanceStore.of(Some).getOrCreate(key, Some::new) then: current == some1 when: - current = InstanceStore.of(Some).putIfAbsent(key, Some::new) + current = InstanceStore.of(Some).getOrCreate(key, Some::new) then: current == some1 @@ -71,13 +71,13 @@ class InstanceStoreTest extends DDSpecification { def key = nextKey() when: - def current = someStore.putIfAbsent(key, () -> some1) + def current = someStore.getOrCreate(key, () -> some1) then: current == some1 when: - current = someStore.putIfAbsent(key, Some::new) + current = someStore.getOrCreate(key, Some::new) then: current == some1 @@ -91,7 +91,7 @@ class InstanceStoreTest extends DDSpecification { someStore.put(key, some1) when: - def current = someStore.putIfAbsent(key, new Creator(invocations)) + def current = someStore.getOrCreate(key, new Creator(invocations)) then: current == some1 @@ -105,14 +105,14 @@ class InstanceStoreTest extends DDSpecification { def key = nextKey() when: - def current = someStore.putIfAbsent(key, new Creator(invocations, some1)) + def current = someStore.getOrCreate(key, new Creator(invocations, some1)) then: current == some1 invocations.get() == 1 when: - current = someStore.putIfAbsent(key, new Creator(invocations)) + current = someStore.getOrCreate(key, new Creator(invocations)) then: current == some1 diff --git a/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy b/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy index 2618d70d520..58a7aa1d13a 100644 --- a/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy +++ b/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy @@ -43,7 +43,7 @@ import datadog.trace.civisibility.source.SourcePathResolver import datadog.trace.civisibility.source.index.RepoIndexBuilder import datadog.trace.civisibility.telemetry.CiVisibilityMetricCollectorImpl import datadog.trace.civisibility.test.ExecutionStrategy -import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore +import datadog.trace.civisibility.utils.StrongMapContextStore import datadog.trace.civisibility.writer.ddintake.CiTestCovMapperV2 import datadog.trace.civisibility.writer.ddintake.CiTestCycleMapperV1 import datadog.trace.common.writer.ListWriter @@ -269,8 +269,8 @@ abstract class CiVisibilityInstrumentationTest extends InstrumentationSpecificat { testFrameworkSessionFactory.startSession(moduleName, component, null, capabilities) }, moduleName, false, - suiteStore != null ? suiteStore : new ConcurrentHashMapContextStore<>(), - testStore != null ? testStore : new ConcurrentHashMapContextStore<>()) + suiteStore != null ? suiteStore : new StrongMapContextStore<>(), + testStore != null ? testStore : new StrongMapContextStore<>()) } } diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java index d094cdf838a..af7f653e921 100644 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java +++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java @@ -38,7 +38,7 @@ import datadog.trace.civisibility.source.index.RepoIndex; import datadog.trace.civisibility.telemetry.CiVisibilityMetricCollectorImpl; import datadog.trace.civisibility.test.ExecutionStrategy; -import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore; +import datadog.trace.civisibility.utils.StrongMapContextStore; import datadog.trace.util.throwable.FatalAgentMisconfigurationError; import java.lang.instrument.Instrumentation; import java.nio.file.Path; @@ -198,8 +198,8 @@ public TestEventsHandler create( repoServices.moduleName, component, null, capabilities), repoServices.moduleName, eagerSessionStart, - suiteStore != null ? suiteStore : new ConcurrentHashMapContextStore<>(), - testStore != null ? testStore : new ConcurrentHashMapContextStore<>()); + suiteStore != null ? suiteStore : new StrongMapContextStore<>(), + testStore != null ? testStore : new StrongMapContextStore<>()); handlers.add(handler); return handler; } diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java deleted file mode 100644 index 19a702c959f..00000000000 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java +++ /dev/null @@ -1,40 +0,0 @@ -package datadog.trace.civisibility.utils; - -import datadog.trace.bootstrap.ContextStore; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; - -public class ConcurrentHashMapContextStore implements ContextStore { - - private final ConcurrentMap m = new ConcurrentHashMap<>(); - - @Override - public C get(K key) { - return m.get(key); - } - - @Override - public void put(K key, C context) { - m.put(key, context); - } - - @Override - public C putIfAbsent(K key, C context) { - return m.computeIfAbsent(key, k -> context); - } - - @Override - public C putIfAbsent(K key, Factory contextFactory) { - return m.computeIfAbsent(key, k -> contextFactory.create()); - } - - @Override - public C computeIfAbsent(K key, KeyAwareFactory contextFactory) { - return m.computeIfAbsent(key, contextFactory::create); - } - - @Override - public C remove(K key) { - return m.remove(key); - } -} diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java new file mode 100644 index 00000000000..2ca2586ce41 --- /dev/null +++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java @@ -0,0 +1,36 @@ +package datadog.trace.civisibility.utils; + +import datadog.trace.bootstrap.ContextStore; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; + +/** Substitute {@link ContextStore} that uses strong-references to track contexts. */ +public class StrongMapContextStore implements ContextStore { + private final ConcurrentMap map = new ConcurrentHashMap<>(); + + @Override + public C get(K key) { + return map.get(key); + } + + @Override + public void put(K key, C context) { + map.put(key, context); + } + + @Override + public C getOrPut(K key, C context) { + C existing = map.putIfAbsent(key, context); + return existing != null ? existing : context; + } + + @Override + public C getOrCompute(K key, KeyAwareFactory contextFactory) { + return map.computeIfAbsent(key, contextFactory::create); + } + + @Override + public C remove(K key) { + return map.remove(key); + } +} diff --git a/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java b/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java index dc62a3261c7..4c20bba1d2b 100644 --- a/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java +++ b/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java @@ -14,7 +14,7 @@ import datadog.trace.bootstrap.ContextStore; import datadog.trace.civisibility.domain.TestFrameworkModule; import datadog.trace.civisibility.domain.TestFrameworkSession; -import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore; +import datadog.trace.civisibility.utils.StrongMapContextStore; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; import org.junit.jupiter.api.Test; @@ -89,8 +89,8 @@ private static TestEventsHandlerImpl handler( private static TestEventsHandlerImpl handler( Supplier testSessionSupplier, boolean eagerSessionStart) { - ContextStore suiteStore = new ConcurrentHashMapContextStore<>(); - ContextStore testStore = new ConcurrentHashMapContextStore<>(); + ContextStore suiteStore = new StrongMapContextStore<>(); + ContextStore testStore = new StrongMapContextStore<>(); return new TestEventsHandlerImpl<>( NoOpMetricCollector.INSTANCE, testSessionSupplier, diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java index 1392a08fbfa..ee8ae19ca09 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java @@ -8,7 +8,7 @@ import datadog.trace.api.Pair; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.FieldBackedContextAccessor; -import datadog.trace.bootstrap.FieldBackedContextStores; +import datadog.trace.bootstrap.WeakMapPerStore; import java.io.Serializable; import java.util.Arrays; import java.util.BitSet; @@ -35,9 +35,6 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper { private static final Logger log = LoggerFactory.getLogger(FieldBackedContextInjector.class); - static final String FIELD_BACKED_CONTEXT_STORES_CLASS = - getInternalName(FieldBackedContextStores.class.getName()); - static final String FIELD_BACKED_CONTEXT_ACCESSOR_CLASS = getInternalName(FieldBackedContextAccessor.class.getName()); @@ -51,12 +48,12 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper { static final String PUTTER_METHOD_DESCRIPTOR = Type.getMethodDescriptor(Type.VOID_TYPE, Type.INT_TYPE, Type.getType(Object.class)); - static final String WEAK_GET_METHOD = "weakGet"; + static final String WEAK_REDIRECT_CLASS = getInternalName(WeakMapPerStore.class.getName()); + static final String WEAK_GET_METHOD_DESCRIPTOR = Type.getMethodDescriptor( Type.getType(Object.class), Type.getType(Object.class), Type.INT_TYPE); - static final String WEAK_PUT_METHOD = "weakPut"; static final String WEAK_PUT_METHOD_DESCRIPTOR = Type.getMethodDescriptor( Type.VOID_TYPE, Type.getType(Object.class), Type.INT_TYPE, Type.getType(Object.class)); @@ -444,11 +441,7 @@ private void invokeWeakGet(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ALOAD, 0); mv.visitIntInsn(Opcodes.ILOAD, 1); mv.visitMethodInsn( - Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, - WEAK_GET_METHOD, - WEAK_GET_METHOD_DESCRIPTOR, - false); + Opcodes.INVOKESTATIC, WEAK_REDIRECT_CLASS, "get", WEAK_GET_METHOD_DESCRIPTOR, false); mv.visitInsn(Opcodes.ARETURN); } @@ -457,11 +450,7 @@ private void invokeWeakPut(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ILOAD, 1); mv.visitIntInsn(Opcodes.ALOAD, 2); mv.visitMethodInsn( - Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, - WEAK_PUT_METHOD, - WEAK_PUT_METHOD_DESCRIPTOR, - false); + Opcodes.INVOKESTATIC, WEAK_REDIRECT_CLASS, "put", WEAK_PUT_METHOD_DESCRIPTOR, false); mv.visitInsn(Opcodes.RETURN); } diff --git a/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java b/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java index 192ac18e374..2a8dcb46fe3 100644 --- a/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java +++ b/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java @@ -82,7 +82,7 @@ public static void methodExit( @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) { final ContextStore contextStore = InstrumentationContext.get(KeyClass.class, Context.class); - final Context context = contextStore.putIfAbsent(thiz, new Context()); + final Context context = contextStore.getOrPut(thiz, new Context()); contextCount = ++context.count; } } @@ -93,7 +93,7 @@ public static void methodExit( @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) { final ContextStore contextStore = InstrumentationContext.get(KeyClass.class, Context.class); - final Context context = contextStore.putIfAbsent(thiz, Context::new); + final Context context = contextStore.getOrCreate(thiz, Context::new); contextCount = ++context.count; } } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java index 3468915c7f1..a192f3d38a1 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java @@ -23,7 +23,7 @@ public class DynamoDbInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws"); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java index 7cb54b62230..3672b49b88d 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java @@ -30,7 +30,7 @@ public class EventBridgeInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final String START_TIME_KEY = "x-datadog-start-time"; private static final String RESOURCE_NAME_KEY = "x-datadog-resource-name"; diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java index e03b6bb2427..4ab9a591572 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java @@ -21,7 +21,7 @@ public class S3Interceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws"); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java index f84a2dfa7e1..7c1a1a85b90 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java @@ -94,7 +94,7 @@ public class AwsSdkClientDecorator extends HttpClientDecorator KINESIS_STREAM_ARN_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("KinesisStreamArn", () -> new ExecutionAttribute<>("KinesisStreamArn")); + .getOrCreate("KinesisStreamArn", () -> new ExecutionAttribute<>("KinesisStreamArn")); // not static because this object would be ClassLoader specific if multiple SDK instances were // loaded by different loaders diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index 980f08ca118..366c67675e3 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -39,7 +39,7 @@ public class TracingExecutionInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final Logger log = LoggerFactory.getLogger(TracingExecutionInterceptor.class); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java index 5ab99ecae96..e78d3b2d188 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java @@ -18,7 +18,7 @@ public class SfnInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); public SfnInterceptor() {} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index 122cc55f399..6f84107b7dd 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -30,7 +30,7 @@ public class SnsInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private SdkBytes getMessageAttributeValueToInject( ExecutionAttributes executionAttributes, String snsTopicName) { diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java index 6be07206130..ed59931810a 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java @@ -32,7 +32,7 @@ public class SqsInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); public SqsInterceptor() {} diff --git a/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java b/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java index a90d65c2472..d32db6b803f 100644 --- a/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java +++ b/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java @@ -148,7 +148,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[ + "datadog.trace.bootstrap.instrumentation.jfr.exceptions.ExceptionSampleEvent:build_time," + "datadog.trace.bootstrap.instrumentation.jfr.backpressure.BackpressureSampleEvent:build_time," + "datadog.trace.bootstrap.instrumentation.jfr.directallocation.DirectAllocationTotalEvent:build_time," - + "datadog.trace.bootstrap.WeakMapContextStore:build_time," + + "datadog.trace.bootstrap.WeakMapPerStore:build_time," + "datadog.trace.config.inversion.ConfigHelper:rerun," + "datadog.trace.config.inversion.ConfigHelper$StrictnessPolicy:rerun," + "datadog.trace.config.inversion.GeneratedSupportedConfigurations:build_time," diff --git a/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java b/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java index ddcee1bcae5..3f7fdd6bdff 100644 --- a/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java +++ b/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java @@ -100,6 +100,6 @@ public static void attachSpanFromStore( return; } - targetContextStore.putIfAbsent(target, state); + targetContextStore.getOrPut(target, state); } } diff --git a/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java b/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java index ade78a85eed..be7f635672c 100644 --- a/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java +++ b/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java @@ -64,11 +64,11 @@ public static void openSession(@Advice.Return final Object session) { if (session instanceof Session) { final ContextStore contextStore = InstrumentationContext.get(Session.class, SessionState.class); - contextStore.putIfAbsent((Session) session, new SessionState(span)); + contextStore.getOrPut((Session) session, new SessionState(span)); } else if (session instanceof StatelessSession) { final ContextStore contextStore = InstrumentationContext.get(StatelessSession.class, SessionState.class); - contextStore.putIfAbsent((StatelessSession) session, new SessionState(span)); + contextStore.getOrPut((StatelessSession) session, new SessionState(span)); } } diff --git a/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java b/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java index 66b2910d592..00482ed4d67 100644 --- a/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java +++ b/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java @@ -57,7 +57,7 @@ public static void openSession(@Advice.Return final SharedSessionContract sessio final ContextStore contextStore = InstrumentationContext.get(SharedSessionContract.class, SessionState.class); - contextStore.putIfAbsent(session, new SessionState(span)); + contextStore.getOrPut(session, new SessionState(span)); } /** diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java b/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java index fb59f988a66..eb78b2cf27e 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java @@ -74,7 +74,7 @@ public static HttpUrlState methodEnter( final ContextStore contextStore = InstrumentationContext.get(HttpURLConnection.class, HttpUrlState.class); - final HttpUrlState state = contextStore.putIfAbsent(thiz, HttpUrlState.FACTORY); + final HttpUrlState state = contextStore.getOrCreate(thiz, HttpUrlState.FACTORY); synchronized (state) { final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpURLConnection.class); diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java index 8c3ffa48231..e3e657fb5d0 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java @@ -118,8 +118,7 @@ public static void bindProducerState( ackMode = Session.AUTO_ACKNOWLEDGE; } sessionState = - sessionStateStore.putIfAbsent( - session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); + sessionStateStore.getOrPut(session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); } boolean isQueue = PRODUCER_DECORATE.isQueue(destination); @@ -159,8 +158,7 @@ public static void bindConsumerState( ackMode = Session.AUTO_ACKNOWLEDGE; } sessionState = - sessionStateStore.putIfAbsent( - session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); + sessionStateStore.getOrPut(session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); } boolean isQueue = CONSUMER_DECORATE.isQueue(destination); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java index a2fa481491c..e428a73d94b 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java @@ -161,7 +161,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, @@ -214,7 +214,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java index 227d8872648..f0309fbcb84 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java @@ -267,7 +267,7 @@ public static void captureConfiguration( @Advice.Argument(0) ProducerConfig producerConfig) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); if (Config.get().isDataStreamsEnabled()) { KafkaConfigHelper.storePendingProducerConfig( state, KafkaConfigHelper.extractProducerConfig(producerConfig)); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java index d6acfe30369..94063de4412 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java @@ -90,7 +90,7 @@ public static void onEnter( String clusterId = newCluster.clusterResource().clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } @@ -111,7 +111,7 @@ public static void onEnter( String clusterId = response.clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java index f704eecade4..085fb7ba866 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java @@ -59,7 +59,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, KafkaConfigHelper.extractConsumerConfig(consumerConfig)); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java index 6557c1b6821..35ebb9f8564 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java @@ -52,7 +52,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, KafkaConfigHelper.extractConsumerConfig(consumerConfig)); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java index 307440b06fb..af2a8ce62ee 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java @@ -16,7 +16,7 @@ public static void onEnter( String clusterId = response.clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java index a1abe7b1959..f5eab61835c 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java @@ -16,7 +16,7 @@ public static void onEnter( String clusterId = newCluster.clusterResource().clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java index a97caca83b2..ddff8190bcb 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java @@ -16,7 +16,7 @@ public static void captureConfiguration( @Advice.Argument(0) ProducerConfig producerConfig) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); if (Config.get().isDataStreamsEnabled()) { KafkaConfigHelper.storePendingProducerConfig( state, KafkaConfigHelper.extractProducerConfig(producerConfig)); diff --git a/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java b/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java index 1a1b030fac1..84f3feaff63 100644 --- a/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java +++ b/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java @@ -82,7 +82,7 @@ public static void beforeExecute(@Advice.This ScenarioRuntime scenarioRuntime) { ExecutionContext executionContext = InstrumentationContext.get(Scenario.class, ExecutionContext.class) - .computeIfAbsent(scenarioRuntime.scenario, ExecutionContext::create); + .getOrCompute(scenarioRuntime.scenario, ExecutionContext::create); // Indicate beforehand if the failures should be suppressed. This aligns the ordering with the // rest of the frameworks diff --git a/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java b/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java index 1be2b8ee422..f950a077085 100644 --- a/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java +++ b/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java @@ -23,7 +23,7 @@ public static void beforeExecute(@Advice.This ScenarioRuntime scenarioRuntime) { ExecutionContext executionContext = InstrumentationContext.get(Scenario.class, ExecutionContext.class) - .computeIfAbsent(scenarioRuntime.getScenario(), ExecutionContext::create); + .getOrCompute(scenarioRuntime.getScenario(), ExecutionContext::create); executionContext.setTestStarted(false); // Indicate beforehand whether failures should be suppressed. This aligns the ordering with diff --git a/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java b/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java index d37cdd69359..c5e7f366a13 100644 --- a/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java +++ b/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java @@ -121,7 +121,7 @@ public static MongoCommandListener tryRegister( @Override public void commandStarted(final CommandStartedEvent event) { if (listenerAccessor != null) { - listenerAccessor.putIfAbsent(event.getConnectionDescription(), this); + listenerAccessor.getOrPut(event.getConnectionDescription(), this); } // If DBM comment injection is enabled, the span is created on the connection instrumentation diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java index 7148cffd723..cc59edc6757 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java @@ -99,7 +99,7 @@ public static ContextScope activateScope(@Advice.Argument(0) final ChannelFuture final ContextContinuation continuation = contextStore - .putIfAbsent(future.getChannel(), ChannelTraceContext.Factory.INSTANCE) + .getOrCreate(future.getChannel(), ChannelTraceContext.Factory.INSTANCE) .getConnectionContinuation(); contextStore.get(future.getChannel()).setConnectionContinuation(null); if (continuation == null) { diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java index 5daa04a82f5..59605630ec6 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java @@ -72,7 +72,7 @@ public static void addConnectContinuation(@Advice.This final Channel channel) { InstrumentationContext.get(Channel.class, ChannelTraceContext.class); if (contextStore - .putIfAbsent(channel, ChannelTraceContext.Factory.INSTANCE) + .getOrCreate(channel, ChannelTraceContext.Factory.INSTANCE) .getConnectionContinuation() != null) { continuation.release(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java index 974ecbe5181..07fc9be8b0f 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java @@ -41,7 +41,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m } final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); ContextScope parentScope = null; final ContextContinuation continuation = channelTraceContext.getConnectionContinuation(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java index e25a261bb56..e340be8ee77 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java @@ -28,7 +28,7 @@ public HttpClientResponseTracingHandler( public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent msg) throws Exception { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); AgentSpan parent = channelTraceContext.getClientParentSpan(); if (parent == null) { @@ -56,7 +56,7 @@ public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent @Override public void exceptionCaught(ChannelHandlerContext ctx, ExceptionEvent e) throws Exception { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); AgentSpan parent = channelTraceContext.getClientParentSpan(); if (parent == null) { diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java index ae8cfd62eb5..3a2aeef4b3f 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java @@ -27,7 +27,7 @@ public HttpServerRequestTracingHandler( @Override public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent msg) { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); if (!(msg.getMessage() instanceof HttpRequest)) { final Context storedContext = channelTraceContext.getServerContext(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java index cede6cd43c9..dffde997846 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java @@ -27,7 +27,7 @@ public HttpServerResponseTracingHandler( @Override public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent msg) { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); final AgentSpan span = channelTraceContext.getServerSpan(); if (span == null || !(msg.getMessage() instanceof HttpResponse)) { diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java index d7451780a5b..8cd6d9d3748 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java @@ -36,7 +36,7 @@ public MaybeBlockResponseHandler(final ContextStore contextFactory) { - final C existing = map.get(key); - if (existing != null) { - return existing; - } - final C created = contextFactory.create(); - map.put(key, created); - return created; - } - - @Override - public C computeIfAbsent(final K key, final KeyAwareFactory contextFactory) { + public C getOrCompute(final K key, final KeyAwareFactory contextFactory) { final C existing = map.get(key); if (existing != null) { return existing; diff --git a/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java b/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java index df2f9a91a0c..1ddbd88fb7e 100644 --- a/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java +++ b/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java @@ -90,7 +90,7 @@ public static void transferToOptimizedSubscriber( context = subscriberContexts.get(source); } if (context != null) { - subscriberContexts.putIfAbsent(target, context); + subscriberContexts.getOrPut(target, context); } } diff --git a/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java b/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java index bfb2bb7cdbe..4a42aeef8fa 100644 --- a/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java +++ b/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java @@ -32,7 +32,7 @@ public static BiConsumer, AgentSpan> putInto( public static BiConsumer, AgentSpan> putIfAbsentInto( final ContextStore store) { - return (publisher, span) -> store.putIfAbsent(publisher, HandoffContext.anyThread(span)); + return (publisher, span) -> store.getOrPut(publisher, HandoffContext.anyThread(span)); } public static Function, Publisher> wrapFunction( diff --git a/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java b/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java index 75a6e3d4764..22149f9a2a4 100644 --- a/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java +++ b/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java @@ -76,7 +76,7 @@ public static void apply( int runStamp = tracker.nextOrdinal().runStamp(); RunContext context = RunContext.getOrCreate(runStamp); RunContext existingContext = - InstrumentationContext.get(Filter.class, RunContext.class).putIfAbsent(filter, context); + InstrumentationContext.get(Filter.class, RunContext.class).getOrPut(filter, context); if (existingContext != context) { // This shouldn't happen. // If it does, instrumentation isn't working as expected, or Scalatest internals changed. diff --git a/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java b/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java index e69d7902d73..a9b03588791 100644 --- a/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java +++ b/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java @@ -43,7 +43,7 @@ public static void after( DECORATE.onProcessStart(span, workflowName); Map map = InstrumentationContext.get(ProcessContext.class, Map.class) - .putIfAbsent(processContext, HashMap::new); + .getOrCreate(processContext, HashMap::new); map.put(wId, span); } } diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java index 41de9e9f2d5..dd310511fdf 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java @@ -49,8 +49,7 @@ public static void captureHandshakeSpan(@Advice.This final WsHandshakeRequest se // apply jee configuration overrides if any since the servlet instrumentation won't kick in // for this span. ClassloaderConfigurationOverrides.maybeEnrichSpan(span); - InstrumentationContext.get(WsHandshakeRequest.class, AgentSpan.class) - .putIfAbsent(self, span); + InstrumentationContext.get(WsHandshakeRequest.class, AgentSpan.class).getOrPut(self, span); } } } diff --git a/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java b/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java index c34602884bd..e4876b8c670 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java +++ b/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java @@ -28,12 +28,12 @@ public class UndertowDecorator @SuppressWarnings("unchecked") public static final AttachmentKey DATADOG_UNDERTOW_CONTINUATION = - attachmentStore.putIfAbsent( + attachmentStore.getOrCreate( "DD_UNDERTOW_CONTINUATION", () -> AttachmentKey.create(ContextContinuation.class)); @SuppressWarnings("unchecked") public static final AttachmentKey PARENT_CONTEXT_KEY = - attachmentStore.putIfAbsent( + attachmentStore.getOrCreate( "DD_UNDERTOW_PARENT_CONTEXT", () -> AttachmentKey.create(Context.class)); public static final UndertowDecorator DECORATE = new UndertowDecorator(); diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java index 5f475eb82cc..ef50904e530 100644 --- a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java @@ -68,7 +68,7 @@ public static void onEnter(@Advice.Argument(0) final Session session) { current.forceSamplingDecision(); } InstrumentationContext.get(Session.class, HandlerContext.Sender.class) - .putIfAbsent( + .getOrPut( session, new HandlerContext.Sender(current.getLocalRootSpan(), session.getId())); } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java index 982330b6805..7219b6f0ce9 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java @@ -96,7 +96,7 @@ public static void onOpen( if (Config.get().isWebsocketMessagesInheritSampling()) { current.forceSamplingDecision(); } - sessionStore.putIfAbsent( + sessionStore.getOrPut( session, new HandlerContext.Sender(current.getLocalRootSpan(), session.getId())); } } catch (Throwable t) { diff --git a/internal-api/build.gradle.kts b/internal-api/build.gradle.kts index cd0a0e59723..42670eb6bcb 100644 --- a/internal-api/build.gradle.kts +++ b/internal-api/build.gradle.kts @@ -77,6 +77,7 @@ extra["excludedClassesCoverage"] = listOf( "datadog.trace.api.debugger.DebuggerConfigUpdate", // Bootstrap API "datadog.trace.bootstrap.ActiveSubsystems", + "datadog.trace.bootstrap.ContextStore", "datadog.trace.bootstrap.ContextStore.Factory", "datadog.trace.bootstrap.instrumentation.api.java.lang.ProcessImplInstrumentationHelpers", "datadog.trace.bootstrap.instrumentation.api.Tags", diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java b/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java index f0528cd976e..6d22bbe92ad 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java @@ -1,5 +1,7 @@ package datadog.trace.bootstrap; +import javax.annotation.Nullable; + /** * Interface to represent context storage for instrumentations. * @@ -43,55 +45,60 @@ interface KeyAwareFactory { } /** - * Get context given the key + * Get context instance for the given key. * - * @param key the key to lookup - * @return context object + * @param key the context key + * @return context instance; {@code null} if the key had no context */ + @Nullable C get(K key); /** - * Put new context instance for given key + * Unconditionally put new context instance for the given key. * - * @param key key to use + * @param key the context key * @param context context instance to save */ void put(K key, C context); /** - * Put new context instance if key is absent + * Gets the context instance for the given key. If no context exists then associate it with the + * new context. * - * @param key key to use - * @param context new context instance to put - * @return old instance if it was present, or new instance + * @param key the context key + * @param context new context instance + * @return existing context instance if present; otherwise new instance */ - C putIfAbsent(K key, C context); + C getOrPut(K key, C context); /** - * Put new context instance if key is absent. Uses context factory to avoid creating objects if - * not needed. + * Gets the context instance for the given key. If no context exists then create one using the + * given factory and associate it with the key. * - * @param key key to use - * @param contextFactory factory instance to produce new context object - * @return old instance if it was present, or new instance + * @param key the context key + * @param contextFactory factory instance to produce new context instances + * @return existing context instance if present; otherwise new instance */ - C putIfAbsent(K key, Factory contextFactory); + default C getOrCreate(K key, Factory contextFactory) { + return getOrCompute(key, contextFactory); + } /** - * Put new context instance if key is absent. Uses context factory to avoid creating objects if - * not needed. + * Gets the context instance for the given key. If no context exists then create one using the + * given factory and associate it with the key. * - * @param key key to use - * @param contextFactory factory instance to produce new context object - * @return old instance if it was present, or new instance + * @param key the context key + * @param contextFactory factory instance to produce new context instances + * @return existing context instance if present; otherwise new instance */ - C computeIfAbsent(K key, KeyAwareFactory contextFactory); + C getOrCompute(K key, KeyAwareFactory contextFactory); /** - * Removes the existing value for key and return it. + * Removes the context instance for the given key. * - * @param key the key remove - * @return removed context object + * @param key the context key + * @return removed context instance; {@code null} if the key had no context */ + @Nullable C remove(K key); }