Fix async context loss when field injection is unavailable

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)

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<K, V> {
+    private static final int MAX_SIZE = 50_000;
+
+    private final WeakConcurrentMap<K, V> 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;
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java

@@ -1,94 +1,207 @@
 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.
  *
- * <p>This class should be created lazily because it uses weak maps with background cleanup.
+ * <p>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).
+ *
+ * <p>This class should be created lazily because it uses background cleanup.
  */
 final class WeakMapContextStore<K, V> implements ContextStore<K, V> {
-  private static final int DEFAULT_MAX_SIZE = 50_000;
+  private static final long CLEAN_FREQUENCY_SECONDS = 1;
 
-  private final int maxSize;
-  private final WeakMap<Object, Object> map = WeakMap.Supplier.newWeakMap();
+  private static final ThreadLocal<LookupKey> LOOKUP_KEY = ThreadLocal.withInitial(LookupKey::new);
 
-  public WeakMapContextStore(int maxSize) {
-    this.maxSize = maxSize;
-  }
+  private final ConcurrentHashMap<Object, V> map = new ConcurrentHashMap<>();
+  private final ReferenceQueue<Object> 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<V> 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<? super K, V> contextFactory) {
+  public V computeIfAbsent(final K key, final KeyAwareFactory<? super K, V> 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);
         }
       }
     }
     return existingContext;
   }
 
   @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<WeakMapContextStore<?, ?>> {
+    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<Object> {
+    private final int hash;
+
+    WeakKey(final Object referent, final ReferenceQueue<Object> 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;
+    }
+  }
 }