diff --git a/dd-java-agent/agent-bootstrap/build.gradle b/dd-java-agent/agent-bootstrap/build.gradle index 398c9ad1351..df898beb88d 100644 --- a/dd-java-agent/agent-bootstrap/build.gradle +++ b/dd-java-agent/agent-bootstrap/build.gradle @@ -33,6 +33,9 @@ dependencies { testImplementation group: 'com.google.guava', name: 'guava-testlib', version: '20.0' testImplementation libs.bundles.junit5 testImplementation libs.bundles.mockito + + // Old WeakMapContextStore baseline in WeakMapContextStoreBenchmark (#10479). + jmhImplementation group: 'com.blogspot.mydailyjava', name: 'weak-lock-free', version: '0.17' } // Must use Java 11 to build JFR enabled code - there is no JFR in OpenJDK 8 (revisit once JFR in Java 8 is available) diff --git a/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakMapContextStoreBenchmark.java b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakMapContextStoreBenchmark.java new file mode 100644 index 00000000000..b96765b54f9 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/WeakMapContextStoreBenchmark.java @@ -0,0 +1,117 @@ +package datadog.trace.bootstrap; + +import static java.util.concurrent.TimeUnit.MICROSECONDS; +import static org.openjdk.jmh.annotations.Mode.AverageTime; + +import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap; +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; + +/** + * Compares the current {@link WeakMapContextStore} (ConcurrentHashMap + inline ReferenceQueue + * expunge, no cap) against the previous implementation (capped {@link WeakConcurrentMap} with a + * periodic cleaner), at different levels of concurrency. This is the hot path for every + * context-store access when field injection is unavailable (issue #10479). + */ +@State(Scope.Benchmark) +@BenchmarkMode(AverageTime) +@OutputTimeUnit(MICROSECONDS) +@SuppressWarnings({"unused", "rawtypes", "unchecked"}) +public class WeakMapContextStoreBenchmark { + private static final int NUM_STORES = 3; + private static final int NUM_KEYS = 1_000; + + @Param({"old", "new"}) + public String impl; + + private BiFunction[] stores; + + private String[] keys; + private String[] values; + + private AtomicInteger threadNumber; + + @Setup(Level.Trial) + public void setup() { + stores = new BiFunction[NUM_STORES]; + for (int i = 0; i < NUM_STORES; i++) { + stores[i] = + "new".equals(impl) + ? new WeakMapContextStore<>()::putIfAbsent + : new CappedWeakConcurrentMapStore<>()::putIfAbsent; + } + + keys = new String[NUM_KEYS]; + values = new String[NUM_KEYS]; + for (int i = 0; i < NUM_KEYS; i++) { + keys[i] = "key_" + i; + values[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 (int i = 0; i < NUM_KEYS; i++) { + blackhole.consume(store.apply(keys[i], values[i])); + } + } + + /** The previous implementation: capped WeakConcurrentMap with synchronized putIfAbsent. */ + static final class CappedWeakConcurrentMapStore { + private static final int MAX_SIZE = 50_000; + + private final WeakConcurrentMap map = new WeakConcurrentMap<>(false, true); + + V putIfAbsent(final K key, final V context) { + V existingContext = map.get(key); + if (null == existingContext) { + synchronized (map) { + existingContext = map.get(key); + if (null == existingContext) { + existingContext = context; + if (map.approximateSize() < MAX_SIZE) { + map.put(key, existingContext); + } + } + } + } + return existingContext; + } + } +} 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 index 1750b7f2b9c..f51f592cbc1 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/WeakMapContextStore.java @@ -1,80 +1,105 @@ package datadog.trace.bootstrap; +import datadog.trace.api.Platform; import datadog.trace.api.internal.VisibleForTesting; +import datadog.trace.util.AgentTaskScheduler; +import java.lang.ref.Reference; +import java.lang.ref.ReferenceQueue; +import java.lang.ref.WeakReference; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; /** * 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. + *

Entries are keyed weakly by identity and reclaimed once their carrier is collected: collected + * keys are drained from a {@link ReferenceQueue} inline on every write — the only place the map can + * grow — and by a periodic background task, so dead entries and their contexts don't accumulate on + * stores that go idle or read-only. There is deliberately no size cap — growth is bounded by live + * carriers, exactly like the injected-field path. A previous 50k cap silently dropped live trace + * context under load when field injection was unavailable (issue #10479). + * + *

This class should be created lazily because it uses background cleanup. */ final class WeakMapContextStore implements ContextStore { - private static final int DEFAULT_MAX_SIZE = 50_000; + private static final long CLEAN_FREQUENCY_SECONDS = 1; - private final int maxSize; - private final WeakMap map = WeakMap.Supplier.newWeakMap(); + private static final ThreadLocal LOOKUP_KEY = ThreadLocal.withInitial(LookupKey::new); - public WeakMapContextStore(int maxSize) { - this.maxSize = maxSize; - } + private final ConcurrentHashMap map = new ConcurrentHashMap<>(); + private final ReferenceQueue queue = new ReferenceQueue<>(); - public WeakMapContextStore() { - this(DEFAULT_MAX_SIZE); + WeakMapContextStore() { + if (!Platform.isNativeImageBuilder()) { + AgentTaskScheduler.get() + .weakScheduleAtFixedRate( + ExpungeTask.INSTANCE, + this, + CLEAN_FREQUENCY_SECONDS, + CLEAN_FREQUENCY_SECONDS, + TimeUnit.SECONDS); + } } @Override - @SuppressWarnings("unchecked") public V get(final K key) { - return (V) map.get(key); + // no expunge: collected entries are unreachable via lookup anyway, reads must stay cheap + final LookupKey lookupKey = LOOKUP_KEY.get(); + try { + return map.get(lookupKey.withReferent(key)); + } finally { + lookupKey.clear(); + } } @Override public void put(final K key, final V context) { - if (map.size() < maxSize) { - map.put(key, context); + // replace-in-place first, so overwrites don't register a redundant WeakKey on the queue + final LookupKey lookupKey = LOOKUP_KEY.get(); + try { + if (null != map.replace(lookupKey.withReferent(key), context)) { + return; + } + } finally { + lookupKey.clear(); } + expunge(); + map.put(new WeakKey(key, queue), 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); - } - } + // check via get() first: the hit path must not allocate a WeakKey + final V existingContext = get(key); + if (null != existingContext) { + return existingContext; } - return existingContext; + expunge(); + final V raceContext = map.putIfAbsent(new WeakKey(key, queue), context); + return null != raceContext ? raceContext : context; } @Override public V putIfAbsent(final K key, final Factory contextFactory) { - return computeIfAbsent(key, contextFactory); + final V existingContext = get(key); + if (null != existingContext) { + return existingContext; + } + return computeIfAbsent(key, ignored -> contextFactory.create()); } @Override - public V computeIfAbsent(K key, KeyAwareFactory contextFactory) { + public V computeIfAbsent(final K key, final 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) { + // not the map's own computeIfAbsent: CHM forbids the mapping function from touching the + // map, and context factories may re-enter this store; a reentrant monitor allows that + synchronized (this) { existingContext = get(key); if (null == existingContext) { existingContext = contextFactory.create(key); - put(key, existingContext); + expunge(); + map.putIfAbsent(new WeakKey(key, queue), existingContext); } } } @@ -82,13 +107,101 @@ public V computeIfAbsent(K key, KeyAwareFactory contextFactory) { } @Override - @SuppressWarnings("unchecked") public V remove(final K key) { - return (V) map.remove(key); + expunge(); + final LookupKey lookupKey = LOOKUP_KEY.get(); + try { + return map.remove(lookupKey.withReferent(key)); + } finally { + lookupKey.clear(); + } } @VisibleForTesting int size() { + expunge(); return map.size(); } + + private void expunge() { + Reference ref; + while ((ref = queue.poll()) != null) { + map.remove(ref); + } + } + + // Explicit class to avoid an implicit hard reference to the store, which must stay collectible. + private static final class ExpungeTask + implements AgentTaskScheduler.Task> { + static final ExpungeTask INSTANCE = new ExpungeTask(); + + @Override + public void run(final WeakMapContextStore target) { + target.expunge(); + } + } + + /** Weak identity key; equality with the stored referent or a {@link LookupKey} for it. */ + private static final class WeakKey extends WeakReference { + private final int hash; + + WeakKey(final Object referent, final ReferenceQueue queue) { + super(referent, queue); + hash = System.identityHashCode(referent); + } + + @Override + public int hashCode() { + return hash; + } + + @Override + public boolean equals(final Object other) { + if (other == this) { + return true; + } + // A collected key can only equal itself; its map entry is removed via the reference queue. + final Object referent = get(); + if (null == referent) { + return false; + } + if (other instanceof WeakKey) { + return referent == ((WeakKey) other).get(); + } + return other instanceof LookupKey && referent == ((LookupKey) other).referent; + } + } + + /** + * Strong query key, reused per thread so lookups don't allocate. Never stored in the map, and + * must be cleared after use so it doesn't retain the carrier. + */ + private static final class LookupKey { + private Object referent; + + LookupKey withReferent(final Object referent) { + this.referent = referent; + return this; + } + + void clear() { + referent = null; + } + + @Override + public int hashCode() { + return System.identityHashCode(referent); + } + + @Override + public boolean equals(final Object other) { + if (other == this) { + return true; + } + if (other instanceof WeakKey) { + return referent == ((WeakKey) other).get(); + } + return other instanceof LookupKey && referent == ((LookupKey) other).referent; + } + } } diff --git a/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/WeakMapContextStoreTest.java b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/WeakMapContextStoreTest.java new file mode 100644 index 00000000000..d5d70fc82ba --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/WeakMapContextStoreTest.java @@ -0,0 +1,145 @@ +package datadog.trace.bootstrap; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +import java.lang.ref.WeakReference; +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; + +/** + * Regression tests for the {@link WeakMapContextStore} fall-back used when bytecode field injection + * is unavailable — e.g. under JDK 25 AOT class linking, where carrier classes come pre-linked from + * the cache and the context-store field cannot be added (issue #10479). + * + *

The store used to silently discard new entries past a 50k cap measured with a stale + * approximate size, dropping live trace context past ~50k context writes/sec. It must retain every + * live carrier's context and reclaim entries once carriers are collected. + */ +class WeakMapContextStoreTest { + + // The former production cap whose overflow used to be dropped silently. + private static final int FORMER_CAP = 50_000; + + @Test + void retainsLiveContextBeyondFormerCap() { + WeakMapContextStore store = new WeakMapContextStore<>(); + + List live = new ArrayList<>(); + for (int i = 0; i < FORMER_CAP + 1; i++) { + Object carrier = new Object(); + live.add(carrier); // nothing may be collected in this test + store.put(carrier, "ctx-" + i); + } + + assertEquals(FORMER_CAP + 1, store.size()); + assertNotNull( + store.get(live.get(FORMER_CAP)), + "fall-back store must not drop live context past the former cap (issue #10479)"); + } + + @Test + void retainsLiveContextAfterHighChurnOfDeadCarriers() throws InterruptedException { + WeakMapContextStore store = new WeakMapContextStore<>(); + + // Completed tasks: their carriers become garbage right after the context write. This is the + // realistic #10479 trigger — sustained churn, few carriers actually live at any moment. + WeakReference probe = null; + for (int i = 0; i < FORMER_CAP; i++) { + Object carrier = new Object(); + if (i == 0) { + probe = new WeakReference<>(carrier); + } + store.put(carrier, "dead-" + i); + } + + Assumptions.assumeTrue(awaitCollected(probe), "GC did not reclaim the dead carriers in time"); + + Object liveCarrier = new Object(); + store.put(liveCarrier, "live-ctx"); + + assertNotNull( + store.get(liveCarrier), + "fall-back store must not drop live context because of dead entries (issue #10479)"); + + // Reference enqueueing is asynchronous, so poll until the dead entries drain away. + long deadline = System.nanoTime() + 10_000_000_000L; + while (store.size() > 1 && System.nanoTime() < deadline) { + System.gc(); + Thread.sleep(50); + } + assertEquals(1, store.size(), "collected carriers must be expunged, not counted"); + } + + @Test + void removesEntryOnceCarrierIsCollected() throws InterruptedException { + WeakMapContextStore store = new WeakMapContextStore<>(); + + Object carrier = new Object(); + WeakReference probe = new WeakReference<>(carrier); + store.put(carrier, "ctx"); + assertEquals(1, store.size()); + + carrier = null; + Assumptions.assumeTrue(awaitCollected(probe), "GC did not reclaim the carrier in time"); + + assertEquals(0, store.size(), "entry must be reclaimed with its carrier"); + } + + @Test + void basicContextStoreContract() { + WeakMapContextStore store = new WeakMapContextStore<>(); + Object carrier = new Object(); + + assertNull(store.get(carrier)); + assertSame("first", store.putIfAbsent(carrier, "first")); + assertSame("first", store.putIfAbsent(carrier, "second")); + assertSame("first", store.computeIfAbsent(carrier, k -> "third")); + assertSame("first", store.get(carrier)); + + store.put(carrier, "replaced"); + assertSame("replaced", store.get(carrier)); + + assertSame("replaced", store.remove(carrier)); + assertNull(store.get(carrier)); + assertEquals(0, store.size()); + } + + @Test + void factoryMayReenterTheStore() { + // context factories run arbitrary instrumentation code that may touch the same store, which + // ConcurrentHashMap.computeIfAbsent forbids — so factories must run outside the map's locks + WeakMapContextStore store = new WeakMapContextStore<>(); + Object carrier = new Object(); + Object other = new Object(); + + Object context = + store.computeIfAbsent( + carrier, + k -> { + store.put(other, "nested-ctx"); + return "ctx"; + }); + + assertSame("ctx", context); + assertSame("ctx", store.get(carrier)); + assertSame("nested-ctx", store.get(other)); + } + + private static boolean awaitCollected(WeakReference ref) throws InterruptedException { + long deadline = System.nanoTime() + 10_000_000_000L; + while (System.nanoTime() < deadline) { + if (ref.get() == null) { + return true; + } + System.gc(); + Thread.sleep(50); + } + return ref.get() == null; + } +}