From 930977367599be36731d8b64d28bf8a9fcf18765 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 14 Oct 2025 12:23:18 +0100 Subject: [PATCH 1/2] Replace weak-map per context store with a single global weak-map --- .../bootstrap/FieldBackedContextStore.java | 27 +-- .../bootstrap/FieldBackedContextStores.java | 10 - .../bootstrap/GlobalWeakContextStore.java | 172 ++++++++++++++++++ .../java/datadog/trace/bootstrap/WeakMap.java | 35 ---- .../trace/bootstrap/WeakMapContextStore.java | 92 ---------- .../trace/agent/tooling/AgentInstaller.java | 2 - .../datadog/trace/agent/tooling/WeakMaps.java | 106 ----------- .../context/FieldBackedContextInjector.java | 12 +- .../trace/agent/tooling/WeakMapTest.groovy | 83 --------- ...veImageGeneratorRunnerInstrumentation.java | 5 +- 10 files changed, 185 insertions(+), 359 deletions(-) create mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/GlobalWeakContextStore.java delete mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java delete mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java delete mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMaps.java delete mode 100644 dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakMapTest.groovy 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..eca5efd0800 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 @@ -2,7 +2,7 @@ /** * {@link ContextStore} that attempts to store context in its keys by using bytecode-injected - * fields. Delegates to a lazy {@link WeakMap} for keys that don't have a field for this store. + * fields. Delegates to the global weak map for keys that don't have a field for this store. */ public final class FieldBackedContextStore implements ContextStore { final int storeId; @@ -16,7 +16,7 @@ public Object get(final Object key) { if (key instanceof FieldBackedContextAccessor) { return ((FieldBackedContextAccessor) key).$get$__datadogContext$(storeId); } else { - return weakStore().get(key); + return GlobalWeakContextStore.weakGet(key, storeId); } } @@ -25,7 +25,7 @@ public void put(final Object key, final Object context) { if (key instanceof FieldBackedContextAccessor) { ((FieldBackedContextAccessor) key).$put$__datadogContext$(storeId, context); } else { - weakStore().put(key, context); + GlobalWeakContextStore.weakPut(key, storeId, context); } } @@ -45,7 +45,7 @@ public Object putIfAbsent(final Object key, final Object context) { } return existingContext; } else { - return weakStore().putIfAbsent(key, context); + return GlobalWeakContextStore.weakPutIfAbsent(key, storeId, context); } } @@ -71,7 +71,7 @@ public Object computeIfAbsent( } return existingContext; } else { - return weakStore().computeIfAbsent(key, contextFactory); + return GlobalWeakContextStore.weakComputeIfAbsent(key, storeId, contextFactory); } } @@ -90,22 +90,7 @@ public Object remove(Object key) { } return existingContext; } else { - return weakStore().remove(key); + return GlobalWeakContextStore.weakRemove(key, storeId); } } - - // only create WeakMap-based fall-back when we need it - private volatile WeakMapContextStore weakStore; - private final Object synchronizationInstance = new Object(); - - WeakMapContextStore weakStore() { - if (null == weakStore) { - synchronized (synchronizationInstance) { - if (null == weakStore) { - weakStore = new WeakMapContextStore<>(); - } - } - } - return weakStore; - } } 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/GlobalWeakContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/GlobalWeakContextStore.java new file mode 100644 index 00000000000..ffe29789a31 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/GlobalWeakContextStore.java @@ -0,0 +1,172 @@ +package datadog.trace.bootstrap; + +import datadog.trace.api.Platform; +import datadog.trace.bootstrap.ContextStore.KeyAwareFactory; +import datadog.trace.util.AgentTaskScheduler; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.lang.ref.Reference; +import java.lang.ref.ReferenceQueue; +import java.lang.ref.WeakReference; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; + +/** + * Global weak {@link ContextStore} that acts as a fall-back when field-injection isn't possible. + */ +public final class GlobalWeakContextStore { + + // global map of weak (key + store-id) wrappers mapped to context values + private static final Map globalMap = new ConcurrentHashMap<>(); + + // stale key wrappers that are now eligible for collection + private static final ReferenceQueue staleKeys = new ReferenceQueue<>(); + + private static final long CLEAN_FREQUENCY_SECONDS = 1; + + private static final int MAX_KEYS_CLEANED_PER_CYCLE = 1_000; + + static { + if (!Platform.isNativeImageBuilder()) { + AgentTaskScheduler.get() + .scheduleAtFixedRate( + GlobalWeakContextStore::cleanStaleKeys, + CLEAN_FREQUENCY_SECONDS, + CLEAN_FREQUENCY_SECONDS, + TimeUnit.SECONDS); + } + } + + /** Checks for stale key wrappers and removes them from the global map. */ + static void cleanStaleKeys() { + int count = 0; + Reference ref; + while ((ref = staleKeys.poll()) != null) { + globalMap.remove(ref); + if (++count >= MAX_KEYS_CLEANED_PER_CYCLE) { + break; // limit work done per call + } + } + } + + private GlobalWeakContextStore() {} + + public static Object weakGet(Object key, int storeId) { + return globalMap.get(new LookupKey(key, storeId)); + } + + public static void weakPut(Object key, int storeId, Object context) { + if (context != null) { + globalMap.put(new StoreKey(key, storeId), context); + } else { + globalMap.remove(new LookupKey(key, storeId)); + } + } + + public static Object weakPutIfAbsent(Object key, int storeId, Object context) { + LookupKey lookupKey = new LookupKey(key, storeId); + Object existing; + if (null == (existing = globalMap.get(lookupKey))) { + // 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. + // 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. + synchronized (key) { + if (null == (existing = globalMap.get(lookupKey))) { + weakPut(key, storeId, existing = context); + } + } + } + return existing; + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + public static Object weakComputeIfAbsent( + Object key, int storeId, KeyAwareFactory contextFactory) { + LookupKey lookupKey = new LookupKey(key, storeId); + Object existing; + if (null == (existing = globalMap.get(lookupKey))) { + // 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. + // 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. + synchronized (key) { + if (null == (existing = globalMap.get(lookupKey))) { + weakPut(key, storeId, existing = contextFactory.create(key)); + } + } + } + return existing; + } + + public static Object weakRemove(Object key, int storeId) { + return globalMap.remove(new LookupKey(key, storeId)); + } + + /** Reference key used to weakly associate a key and store-id with a context value. */ + static final class StoreKey extends WeakReference { + final int hash; + final int storeId; + + StoreKey(Object key, int storeId) { + super(key, staleKeys); + this.hash = (31 * storeId) + System.identityHashCode(key); + this.storeId = storeId; + } + + @Override + public int hashCode() { + return hash; + } + + @Override + @SuppressFBWarnings("Eq") // symmetric because it mirrors LookupKey.equals + public boolean equals(Object o) { + if (o instanceof LookupKey) { + LookupKey lookupKey = (LookupKey) o; + return storeId == lookupKey.storeId && get() == lookupKey.key; + } else if (o instanceof StoreKey) { + StoreKey storeKey = (StoreKey) o; + return storeId == storeKey.storeId && get() == storeKey.get(); + } else { + return false; + } + } + } + + /** Temporary key used for lookup purposes without the reference tracking overhead. */ + static final class LookupKey { + final Object key; + final int hash; + final int storeId; + + LookupKey(Object key, int storeId) { + this.key = key; + this.hash = (31 * storeId) + System.identityHashCode(key); + this.storeId = storeId; + } + + @Override + public int hashCode() { + return hash; + } + + @Override + @SuppressFBWarnings("Eq") // symmetric because it mirrors StoreKey.equals + public boolean equals(Object o) { + if (o instanceof StoreKey) { + StoreKey storeKey = (StoreKey) o; + return storeId == storeKey.storeId && key == storeKey.get(); + } else if (o instanceof LookupKey) { + LookupKey lookupKey = (LookupKey) o; + return storeId == lookupKey.storeId && key == lookupKey.key; + } else { + return false; + } + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java deleted file mode 100644 index 1aa859ea3d8..00000000000 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java +++ /dev/null @@ -1,35 +0,0 @@ -package datadog.trace.bootstrap; - -import java.util.function.Function; - -public interface WeakMap { - int size(); - - boolean containsKey(K target); - - V get(K key); - - void put(K key, V value); - - void putIfAbsent(K key, V value); - - V computeIfAbsent(K key, Function supplier); - - V remove(K key); - - abstract class Supplier { - private static volatile Supplier SUPPLIER; - - protected abstract WeakMap get(); - - public static WeakMap newWeakMap() { - return SUPPLIER.get(); - } - - public static synchronized void registerIfAbsent(Supplier supplier) { - if (null == SUPPLIER) { - SUPPLIER = supplier; - } - } - } -} 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/WeakMapContextStore.java deleted file mode 100644 index a3926412889..00000000000 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java +++ /dev/null @@ -1,92 +0,0 @@ -package datadog.trace.bootstrap; - -/** - * Weak {@link ContextStore} that acts as a fall-back 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 WeakMapContextStore(int maxSize) { - this.maxSize = maxSize; - } - - public WeakMapContextStore() { - this(DEFAULT_MAX_SIZE); - } - - @Override - @SuppressWarnings("unchecked") - public V get(final K key) { - return (V) map.get(key); - } - - @Override - public void put(final K key, final V context) { - if (map.size() < maxSize) { - map.put(key, context); - } - } - - @Override - public V putIfAbsent(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. - // 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. - synchronized (map) { - existingContext = get(key); - if (null == existingContext) { - existingContext = context; - put(key, existingContext); - } - } - } - 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 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. - // 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. - synchronized (map) { - existingContext = get(key); - if (null == existingContext) { - existingContext = contextFactory.create(key); - put(key, existingContext); - } - } - } - return existingContext; - } - - @Override - @SuppressWarnings("unchecked") - public V remove(final K key) { - return (V) map.remove(key); - } - - // Package reachable for testing - int size() { - return map.size(); - } -} diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 5b664280f91..a7f653d4a2b 100644 --- a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java @@ -56,8 +56,6 @@ public class AgentInstaller { static { addByteBuddyRawSetting(); disableByteBuddyNexus(); - // register weak map supplier as early as possible - WeakMaps.registerAsSupplier(); circularityErrorWorkaround(); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMaps.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMaps.java deleted file mode 100644 index 90f39205b6a..00000000000 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMaps.java +++ /dev/null @@ -1,106 +0,0 @@ -package datadog.trace.agent.tooling; - -import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap; -import datadog.trace.api.Platform; -import datadog.trace.bootstrap.WeakMap; -import datadog.trace.util.AgentTaskScheduler; -import datadog.trace.util.AgentTaskScheduler.Task; -import java.util.concurrent.TimeUnit; -import java.util.function.Function; - -public class WeakMaps { - private static final long CLEAN_FREQUENCY_SECONDS = 1; - - public static WeakMap newWeakMap() { - final WeakConcurrentMap map = new WeakConcurrentMap<>(false, true); - if (!Platform.isNativeImageBuilder()) { - AgentTaskScheduler.get() - .weakScheduleAtFixedRate( - MapCleaningTask.INSTANCE, - map, - CLEAN_FREQUENCY_SECONDS, - CLEAN_FREQUENCY_SECONDS, - TimeUnit.SECONDS); - } - return new Adapter<>(map); - } - - private WeakMaps() {} - - public static void registerAsSupplier() { - WeakMap.Supplier.registerIfAbsent( - new WeakMap.Supplier() { - @Override - protected WeakMap get() { - return WeakMaps.newWeakMap(); - } - }); - } - - // Important to use explicit class to avoid implicit hard references to target - private static class MapCleaningTask implements Task> { - static final MapCleaningTask INSTANCE = new MapCleaningTask(); - - @Override - public void run(final WeakConcurrentMap target) { - target.expungeStaleEntries(); - } - } - - private static class Adapter implements WeakMap { - private final WeakConcurrentMap map; - - private Adapter(final WeakConcurrentMap map) { - this.map = map; - } - - @Override - public int size() { - return map.approximateSize(); - } - - @Override - public boolean containsKey(final K key) { - return map.containsKey(key); - } - - @Override - public V get(final K key) { - return map.get(key); - } - - @Override - public void put(final K key, final V value) { - if (null != value) { - map.put(key, value); - } else { - map.remove(key); // WeakConcurrentMap doesn't accept null values - } - } - - @Override - public void putIfAbsent(final K key, final V value) { - map.putIfAbsent(key, value); - } - - @Override - public V computeIfAbsent(final K key, final Function supplier) { - V value = map.get(key); - if (null == value) { - synchronized (this) { - value = map.get(key); - if (null == value) { - value = supplier.apply(key); - map.put(key, value); - } - } - } - return value; - } - - @Override - public V remove(K key) { - return map.remove(key); - } - } -} 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 276ac6daf1b..11683c3b598 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.GlobalWeakContextStore; 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,6 +48,9 @@ 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 GLOBAL_WEAK_CONTEXT_STORE_CLASS = + getInternalName(GlobalWeakContextStore.class.getName()); + static final String WEAK_GET_METHOD = "weakGet"; static final String WEAK_GET_METHOD_DESCRIPTOR = Type.getMethodDescriptor( @@ -445,7 +445,7 @@ private void invokeWeakGet(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ILOAD, 1); mv.visitMethodInsn( Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, + GLOBAL_WEAK_CONTEXT_STORE_CLASS, WEAK_GET_METHOD, WEAK_GET_METHOD_DESCRIPTOR, false); @@ -458,7 +458,7 @@ private void invokeWeakPut(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ALOAD, 2); mv.visitMethodInsn( Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, + GLOBAL_WEAK_CONTEXT_STORE_CLASS, WEAK_PUT_METHOD, WEAK_PUT_METHOD_DESCRIPTOR, false); diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakMapTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakMapTest.groovy deleted file mode 100644 index e501161c6fe..00000000000 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakMapTest.groovy +++ /dev/null @@ -1,83 +0,0 @@ -package datadog.trace.agent.tooling - -import datadog.trace.test.util.GCUtils -import datadog.trace.test.util.DDSpecification -import spock.lang.IgnoreIf - -import java.lang.ref.WeakReference -import java.util.concurrent.TimeUnit - -class WeakMapTest extends DDSpecification { - - def "WeakMap accepts null values"() { - setup: - def map = WeakMaps.newWeakMap() - - when: - map.put('key', null) - - then: - noExceptionThrown() - } - - def "Calling newWeakMap creates independent maps"() { - setup: - def key = new Object() - def map1 = WeakMaps.newWeakMap() - def map2 = WeakMaps.newWeakMap() - - when: - map1.put(key, "value1") - map2.put(key, "value2") - - then: - map1.get(key) == "value1" - map2.get(key) == "value2" - } - - //@Flaky("awaitGC usage is flaky") - @IgnoreIf(reason="Often fails in Semeru runtime", value = { System.getProperty("java.runtime.name").contains("Semeru") }) - def "Unreferenced map gets cleaned up"() { - setup: - def map = WeakMaps.newWeakMap() - def ref = new WeakReference(map) - - when: - def mapRef = new WeakReference(map) - map = null - GCUtils.awaitGC(mapRef) - - then: - ref.get() == null - } - - //@Flaky("awaitGC usage is flaky") - @IgnoreIf(reason="Often fails in Semeru runtime", value = { System.getProperty("java.runtime.name").contains("Semeru") }) - def "Unreferenced keys get cleaned up"() { - setup: - def key = new Object() - def map = WeakMaps.newWeakMap() - map.put(key, "value") - GCUtils.awaitGC() - - expect: - map.size() == 1 - - when: - def keyRef = new WeakReference(key) - key = null - GCUtils.awaitGC(keyRef) - - // Sleep enough time for cleanup thread to get scheduled. - // But on a very slow box (or high load) scheduling may not be exactly predictable - // so we try a few times. - int count = 0 - while (map.size() != 0 && count < 10) { - Thread.sleep(TimeUnit.SECONDS.toMillis(WeakMaps.CLEAN_FREQUENCY_SECONDS)) - count++ - } - - then: - map.size() == 0 - } -} diff --git a/dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java b/dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java index 7281ca41b36..bff848d0d6c 100644 --- a/dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java +++ b/dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java @@ -74,7 +74,6 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[ + "datadog.environment.JavaVirtualMachine:rerun," + "datadog.environment.OperatingSystem:rerun," + "datadog.environment.OperatingSystem$Architecture:rerun," - + "datadog.trace.agent.tooling.WeakMaps$Adapter:build_time," + "datadog.trace.api.Config:rerun," + "datadog.trace.api.Platform:rerun," + "datadog.trace.api.Platform$Captured:build_time," @@ -139,7 +138,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.GlobalWeakContextStore:build_time," + "datadog.trace.instrumentation.guava10.GuavaAsyncResultExtension:build_time," + "datadog.trace.instrumentation.reactivestreams.ReactiveStreamsAsyncResultExtension:build_time," + "datadog.trace.instrumentation.reactor.core.ReactorAsyncResultExtension:build_time," @@ -161,8 +160,6 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[ + "datadog.slf4j.helpers.SubstituteLoggerFactory:build_time," + "datadog.slf4j.impl.StaticLoggerBinder:build_time," + "datadog.slf4j.LoggerFactory:build_time," - + "com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap:build_time," - + "com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap$1:build_time," + "net.bytebuddy:build_time," + "com.sun.proxy:build_time," + "jnr.enxio.channels:run_time," From 789a9969a4643a227d1f4fc4bb29cbbbd997e3d7 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 15 Oct 2025 12:06:09 +0100 Subject: [PATCH 2/2] Benchmark old and new approaches at different levels of concurrency --- .../bootstrap/WeakContextStoreBenchmark.java | 94 ++++++++++++++++ .../trace/bootstrap/weakmap/WeakMap.java | 35 ++++++ .../weakmap/WeakMapContextStore.java | 94 ++++++++++++++++ .../trace/bootstrap/weakmap/WeakMaps.java | 105 ++++++++++++++++++ 4 files changed, 328 insertions(+) create mode 100644 dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakContextStoreBenchmark.java create mode 100644 dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMap.java create mode 100644 dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMapContextStore.java create mode 100644 dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMaps.java diff --git a/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakContextStoreBenchmark.java b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakContextStoreBenchmark.java new file mode 100644 index 00000000000..780f79c5584 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakContextStoreBenchmark.java @@ -0,0 +1,94 @@ +package datadog.trace.bootstrap; + +import static java.util.concurrent.TimeUnit.MICROSECONDS; +import static org.openjdk.jmh.annotations.Mode.AverageTime; + +import datadog.trace.bootstrap.weakmap.WeakMapContextStore; +import datadog.trace.bootstrap.weakmap.WeakMaps; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.infra.Blackhole; + +@State(Scope.Benchmark) +@BenchmarkMode(AverageTime) +@OutputTimeUnit(MICROSECONDS) +@SuppressWarnings({"unused", "rawtypes", "unchecked"}) +public class WeakContextStoreBenchmark { + private static final int NUM_STORES = 3; + + @Param({"false", "true"}) + public boolean global; + + private BiFunction[] stores; + + private Map kvs; + + private AtomicInteger threadNumber; + + @Setup(Level.Trial) + public void setup() { + WeakMaps.registerAsSupplier(); + + stores = new BiFunction[NUM_STORES]; + for (int i = 0; i < NUM_STORES; i++) { + stores[i] = global ? globalWeakStorePutIfAbsent(i) : weakMapStorePutIfAbsent(i); + } + + kvs = new HashMap<>(); + for (int i = 0; i < 1_000; i++) { + kvs.put("key_" + i, "value_" + i); + } + + threadNumber = new AtomicInteger(); + } + + @Benchmark + @Fork(value = 1) + @Threads(value = 1) + public void singleThreaded(Blackhole blackhole) { + test(blackhole); + } + + @Benchmark + @Fork(value = 1) + @Threads(value = 10) + public void multiThreaded10(Blackhole blackhole) { + test(blackhole); + } + + @Benchmark + @Fork(value = 1) + @Threads(value = 100) + public void multiThreaded100(Blackhole blackhole) { + test(blackhole); + } + + private void test(Blackhole blackhole) { + // assign each benchmark thread a single store to operate on during the benchmark + // the number of concurrent requests to a store goes up as more threads are added + BiFunction store = stores[threadNumber.getAndIncrement() % NUM_STORES]; + for (Map.Entry e : kvs.entrySet()) { + blackhole.consume(store.apply(e.getKey(), e.getValue())); + } + } + + private static BiFunction globalWeakStorePutIfAbsent(int storeId) { + return (k, v) -> GlobalWeakContextStore.weakPutIfAbsent(k, storeId, v); + } + + private static BiFunction weakMapStorePutIfAbsent(int storeId) { + return new WeakMapContextStore<>(storeId)::putIfAbsent; + } +} diff --git a/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMap.java b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMap.java new file mode 100644 index 00000000000..9a81fe3dc23 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMap.java @@ -0,0 +1,35 @@ +package datadog.trace.bootstrap.weakmap; + +import java.util.function.Function; + +public interface WeakMap { + int size(); + + boolean containsKey(K target); + + V get(K key); + + void put(K key, V value); + + void putIfAbsent(K key, V value); + + V computeIfAbsent(K key, Function supplier); + + V remove(K key); + + abstract class Supplier { + private static volatile Supplier SUPPLIER; + + protected abstract WeakMap get(); + + public static WeakMap newWeakMap() { + return SUPPLIER.get(); + } + + public static synchronized void registerIfAbsent(Supplier supplier) { + if (null == SUPPLIER) { + SUPPLIER = supplier; + } + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMapContextStore.java b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMapContextStore.java new file mode 100644 index 00000000000..9a6ac11c419 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMapContextStore.java @@ -0,0 +1,94 @@ +package datadog.trace.bootstrap.weakmap; + +import datadog.trace.bootstrap.ContextStore; + +/** + * Weak {@link ContextStore} that acts as a fall-back when field-injection isn't possible. + * + *

This class should be created lazily because it uses weak maps with background cleanup. + */ +public 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 WeakMapContextStore(int maxSize) { + this.maxSize = maxSize; + } + + public WeakMapContextStore() { + this(DEFAULT_MAX_SIZE); + } + + @Override + @SuppressWarnings("unchecked") + public V get(final K key) { + return (V) map.get(key); + } + + @Override + public void put(final K key, final V context) { + if (map.size() < maxSize) { + map.put(key, context); + } + } + + @Override + public V putIfAbsent(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. + // 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. + synchronized (map) { + existingContext = get(key); + if (null == existingContext) { + existingContext = context; + put(key, existingContext); + } + } + } + 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 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. + // 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. + synchronized (map) { + existingContext = get(key); + if (null == existingContext) { + existingContext = contextFactory.create(key); + put(key, existingContext); + } + } + } + return existingContext; + } + + @Override + @SuppressWarnings("unchecked") + public V remove(final K key) { + return (V) map.remove(key); + } + + // Package reachable for testing + int size() { + return map.size(); + } +} diff --git a/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMaps.java b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMaps.java new file mode 100644 index 00000000000..d460851b5b3 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/weakmap/WeakMaps.java @@ -0,0 +1,105 @@ +package datadog.trace.bootstrap.weakmap; + +import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap; +import datadog.trace.api.Platform; +import datadog.trace.util.AgentTaskScheduler; +import datadog.trace.util.AgentTaskScheduler.Task; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; + +public class WeakMaps { + private static final long CLEAN_FREQUENCY_SECONDS = 1; + + public static WeakMap newWeakMap() { + final WeakConcurrentMap map = new WeakConcurrentMap<>(false, true); + if (!Platform.isNativeImageBuilder()) { + AgentTaskScheduler.get() + .weakScheduleAtFixedRate( + MapCleaningTask.INSTANCE, + map, + CLEAN_FREQUENCY_SECONDS, + CLEAN_FREQUENCY_SECONDS, + TimeUnit.SECONDS); + } + return new Adapter<>(map); + } + + private WeakMaps() {} + + public static void registerAsSupplier() { + WeakMap.Supplier.registerIfAbsent( + new WeakMap.Supplier() { + @Override + protected WeakMap get() { + return WeakMaps.newWeakMap(); + } + }); + } + + // Important to use explicit class to avoid implicit hard references to target + private static class MapCleaningTask implements Task> { + static final MapCleaningTask INSTANCE = new MapCleaningTask(); + + @Override + public void run(final WeakConcurrentMap target) { + target.expungeStaleEntries(); + } + } + + private static class Adapter implements WeakMap { + private final WeakConcurrentMap map; + + private Adapter(final WeakConcurrentMap map) { + this.map = map; + } + + @Override + public int size() { + return map.approximateSize(); + } + + @Override + public boolean containsKey(final K key) { + return map.containsKey(key); + } + + @Override + public V get(final K key) { + return map.get(key); + } + + @Override + public void put(final K key, final V value) { + if (null != value) { + map.put(key, value); + } else { + map.remove(key); // WeakConcurrentMap doesn't accept null values + } + } + + @Override + public void putIfAbsent(final K key, final V value) { + map.putIfAbsent(key, value); + } + + @Override + public V computeIfAbsent(final K key, final Function supplier) { + V value = map.get(key); + if (null == value) { + synchronized (this) { + value = map.get(key); + if (null == value) { + value = supplier.apply(key); + map.put(key, value); + } + } + } + return value; + } + + @Override + public V remove(K key) { + return map.remove(key); + } + } +}