Fix lookup-then-read race in GenerationalUtf8Cache.getUtf8

communication/src/main/java/datadog/communication/serialization/GenerationalUtf8Cache.java

@@ -2,6 +2,7 @@
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.nio.charset.StandardCharsets;
+import javax.annotation.concurrent.ThreadSafe;
 
 /**
  * 2-level generational cache of UTF8 values - primarily intended to be used for tag values
@@ -62,6 +63,7 @@
  * calling ValueUtf8Cache#reclibrate will adjust promotion thresholds to
  * provide better cache utilization.
  */
+@ThreadSafe
 @SuppressFBWarnings(
     value = "IS2_INCONSISTENT_SYNC",
     justification =
@@ -219,32 +221,42 @@ public final byte[] getUtf8(String value, long accessTimeMs) {
     CacheEntry[] tenuredEntries = this.tenuredEntries;
     int matchingTenuredIndex = lookupEntryIndex(tenuredEntries, MAX_TENURED_PROBES, adjHash, value);
     if (matchingTenuredIndex != -1) {
+      // The slot can be mutated concurrently between the lookup and this read: nulled (recalibrate
+      // purge / eviction) or reassigned to a *different* value. CacheEntry identity is immutable
+      // (adjHash/value/valueUtf8 are final), so re-validate the loaded reference against the
+      // request; anything but a match means the slot moved out from under us, so fall through and
+      // treat it as a miss rather than NPE'ing (null) or returning another value's bytes
+      // (reassigned).
       CacheEntry tenuredEntry = tenuredEntries[matchingTenuredIndex];
+      if (tenuredEntry != null && tenuredEntry.matches(adjHash, value)) {
+        tenuredEntry.hit(accessTimeMs);
 
-      tenuredEntry.hit(accessTimeMs);
-
-      this.tenuredHits += 1;
-      return tenuredEntry.utf8();
+        this.tenuredHits += 1;
+        return tenuredEntry.utf8();
+      }
     }
 
     CacheEntry[] edenEntries = this.edenEntries;
     int matchingEdenIndex = lookupEntryIndex(edenEntries, MAX_EDEN_PROBES, adjHash, value);
     if (matchingEdenIndex != -1) {
+      // Same lookup-then-read race as tenured, plus concurrent promotion nulls the slot (line
+      // below); re-validate the loaded reference and treat null-or-mismatch as a miss.
       CacheEntry edenEntry = edenEntries[matchingEdenIndex];
+      if (edenEntry != null && edenEntry.matches(adjHash, value)) {
+        double hits = edenEntry.hit(accessTimeMs);
+        if (hits > this.promotionThreshold) {
+          // mark promoted first - to avoid racy insertions
+          this.promotions += 1;
 
-      double hits = edenEntry.hit(accessTimeMs);
-      if (hits > this.promotionThreshold) {
-        // mark promoted first - to avoid racy insertions
-        this.promotions += 1;
+          boolean evicted = lruInsert(this.tenuredEntries, MAX_TENURED_PROBES, edenEntry);
+          if (evicted) this.tenuredEvictions += 1;
 
-        boolean evicted = lruInsert(this.tenuredEntries, MAX_TENURED_PROBES, edenEntry);
-        if (evicted) this.tenuredEvictions += 1;
+          edenEntries[matchingEdenIndex] = null;
+        }
 
-        edenEntries[matchingEdenIndex] = null;
+        this.edenHits += 1;
+        return edenEntry.utf8();
       }
-
-      this.edenHits += 1;
-      return edenEntry.utf8();
     }
 
     boolean wasMarked = Caching.mark(this.edenMarkers, adjHash);

communication/src/main/java/datadog/communication/serialization/SimpleUtf8Cache.java

@@ -1,6 +1,7 @@
 package datadog.communication.serialization;
 
 import java.nio.charset.StandardCharsets;
+import javax.annotation.concurrent.ThreadSafe;
 
 /**
  * A simple UTF8 cache - primarily intended for tag names
@@ -41,6 +42,7 @@
  * If there are no available slots in entries for a newly created CacheEntry,
  * a LFU: least frequently used eviction policy is used to free up a slot.
  */
+@ThreadSafe
 public final class SimpleUtf8Cache implements EncodingCache {
   static final int MAX_CAPACITY = 1024;
 

communication/src/test/java/datadog/communication/serialization/GenerationalUtf8CacheTest.java

@@ -7,8 +7,12 @@
 import static org.junit.jupiter.api.Assertions.assertSame;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.Random;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -175,6 +179,88 @@ public void fuzz() {
     assertNotEquals(0, promotedHits);
   }
 
+  @Test
+  public void concurrentAccess_neverThrowsOrReturnsWrongBytes() throws InterruptedException {
+    // Regression test for a lookup-then-read race in getUtf8(): a slot can be nulled (recalibrate
+    // purge / eviction) or reassigned to a *different* value between lookupEntryIndex() and the
+    // array read. Before the fix this either NPE'd (null slot) or, worse, silently returned another
+    // value's bytes (reassigned slot). Serialization is single-threaded today, but the cache is
+    // built to allow concurrent use, so this exercises that contract.
+    final GenerationalUtf8Cache cache = create();
+
+    // More distinct values than the cache can hold, so promotions/evictions churn slots hard.
+    final String[] values = new String[256];
+    for (int i = 0; i < values.length; ++i) {
+      values[i] = "value-" + i;
+    }
+
+    final int threadCount = 8;
+    final int iterationsPerThread = 200_000;
+    final CountDownLatch start = new CountDownLatch(1);
+    final AtomicReference<Throwable> failure = new AtomicReference<>();
+    final AtomicInteger readersRunning = new AtomicInteger(threadCount);
+
+    Thread[] readers = new Thread[threadCount];
+    for (int t = 0; t < threadCount; ++t) {
+      readers[t] =
+          new Thread(
+              () -> {
+                try {
+                  start.await();
+                  ThreadLocalRandom random = ThreadLocalRandom.current();
+                  for (int i = 0; i < iterationsPerThread && failure.get() == null; ++i) {
+                    String value = values[random.nextInt(values.length)];
+                    byte[] result = cache.getUtf8(value);
+                    if (!Arrays.equals(value.getBytes(StandardCharsets.UTF_8), result)) {
+                      failure.compareAndSet(
+                          null,
+                          new AssertionError(
+                              "getUtf8(\""
+                                  + value
+                                  + "\") returned bytes for \""
+                                  + new String(result, StandardCharsets.UTF_8)
+                                  + "\""));
+                      return;
+                    }
+                  }
+                } catch (Throwable e) {
+                  failure.compareAndSet(null, e);
+                } finally {
+                  readersRunning.decrementAndGet();
+                }
+              });
+    }
+
+    // Recalibrate in a tight loop for the duration, nulling decayed slots concurrently with reads.
+    Thread recalibrator =
+        new Thread(
+            () -> {
+              try {
+                start.await();
+                while (readersRunning.get() > 0 && failure.get() == null) {
+                  cache.recalibrate();
+                }
+              } catch (Throwable e) {
+                failure.compareAndSet(null, e);
+              }
+            });
+
+    for (Thread reader : readers) {
+      reader.start();
+    }
+    recalibrator.start();
+    start.countDown();
+
+    for (Thread reader : readers) {
+      reader.join();
+    }
+    recalibrator.join();
+
+    if (failure.get() != null) {
+      throw new AssertionError("concurrent getUtf8() failed", failure.get());
+    }
+  }
+
   @Test
   public void bigString_dont_cache() {
     String lorem = "Lorem ipsum dolor sit amet";

dd-trace-core/src/jmh/java/datadog/trace/common/writer/ddagent/Utf8Benchmark.java

@@ -1,75 +1,31 @@
 package datadog.trace.common.writer.ddagent;
 
+import static datadog.trace.common.writer.ddagent.Utf8Workload.NUM_LOOKUPS;
+import static datadog.trace.common.writer.ddagent.Utf8Workload.nextTag;
+import static datadog.trace.common.writer.ddagent.Utf8Workload.nextValue;
+
 import datadog.communication.serialization.GenerationalUtf8Cache;
 import datadog.communication.serialization.SimpleUtf8Cache;
 import java.nio.charset.StandardCharsets;
-import java.util.concurrent.ThreadLocalRandom;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Mode;
 import org.openjdk.jmh.infra.Blackhole;
 
 /**
- * This benchmark isn't really intended to used to measure throughput, but rather to be used with
- * "-prof gc" to check bytes / op.
+ * Single-threaded UTF8 cache benchmark. This reflects how the caches are actually driven today:
+ * trace serialization runs on a single thread, so one thread performs the lookups and drives {@link
+ * GenerationalUtf8Cache#recalibrate()} inline at a natural transaction boundary (here, once per
+ * op). This is the representative allocation/throughput number. See {@link Utf8ConcurrentBenchmark}
+ * for the multi-threaded contract/guardrail variant.
  *
- * <p>Since {@link String#getBytes(java.nio.charset.Charset)} is intrinsified the caches typically
- * perform worse throughput wise, the benefit of the caches is to reduce allocation. Intention of
- * this benchmark is to create data that roughly resembles what might be seen in a trace payload.
- * Tag names are quite static, tag values are mostly low cardinality, but some tag values have
- * infinite cardinality.
+ * <p>This benchmark isn't really intended to measure throughput, but rather to be used with "-prof
+ * gc" to check bytes / op. Since {@link String#getBytes(java.nio.charset.Charset)} is intrinsified
+ * the caches typically perform worse throughput wise; the benefit of the caches is to reduce
+ * allocation.
  */
 @BenchmarkMode(Mode.Throughput)
 public class Utf8Benchmark {
-  static final int NUM_LOOKUPS = 10_000;
-
-  static final String[] TAGS = {
-    "_dd.asm.keep",
-    "ci.provider",
-    "language",
-    "db.statement",
-    "ci.job.url",
-    "ci.pipeline.url",
-    "db.pool",
-    "http.forwarder",
-    "db.warehouse",
-    "custom"
-  };
-
-  static int pos = 0;
-  static int standardVal = 0;
-
-  static final String nextTag() {
-    if (pos == TAGS.length - 1) {
-      pos = 0;
-    } else {
-      pos += 1;
-    }
-    return TAGS[pos];
-  }
-
-  static final String nextValue(String tag) {
-    if (tag.equals("custom")) {
-      return nextCustomValue(tag);
-    } else {
-      return nextStandardValue(tag);
-    }
-  }
-
-  /*
-   * Produces a high cardinality value - > thousands of distinct values per tag - many 1-time values
-   */
-  static final String nextCustomValue(String tag) {
-    return tag + ThreadLocalRandom.current().nextInt();
-  }
-
-  /*
-   * Produces a moderate cardinality value - tens of distinct values per tag
-   */
-  static final String nextStandardValue(String tag) {
-    return tag + ThreadLocalRandom.current().nextInt(20);
-  }
-
   @Benchmark
   public static final String tagUtf8_baseline() {
     return nextTag();
@@ -109,7 +65,7 @@ public static final void valueUtf8_baseline(Blackhole bh) {
   @Benchmark
   public static final void valueUtf8_cache_generational(Blackhole bh) {
     GenerationalUtf8Cache valueCache = VALUE_CACHE;
-    valueCache.recalibrate();
+    valueCache.recalibrate(); // single thread drives recalibrate inline, at a transaction boundary
 
     for (int i = 0; i < NUM_LOOKUPS; ++i) {
       String tag = nextTag();
@@ -125,7 +81,7 @@ public static final void valueUtf8_cache_generational(Blackhole bh) {
   @Benchmark
   public static final void valueUtf8_cache_simple(Blackhole bh) {
     SimpleUtf8Cache valueCache = SIMPLE_VALUE_CACHE;
-    valueCache.recalibrate();
+    valueCache.recalibrate(); // single thread drives recalibrate inline, at a transaction boundary
 
     for (int i = 0; i < NUM_LOOKUPS; ++i) {
       String tag = nextTag();

dd-trace-core/src/jmh/java/datadog/trace/common/writer/ddagent/Utf8ConcurrentBenchmark.java

@@ -0,0 +1,72 @@
+package datadog.trace.common.writer.ddagent;
+
+import static datadog.trace.common.writer.ddagent.Utf8Workload.NUM_LOOKUPS;
+import static datadog.trace.common.writer.ddagent.Utf8Workload.nextTag;
+import static datadog.trace.common.writer.ddagent.Utf8Workload.nextValue;
+
+import datadog.communication.serialization.GenerationalUtf8Cache;
+import datadog.communication.serialization.SimpleUtf8Cache;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Group;
+import org.openjdk.jmh.annotations.GroupThreads;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.infra.Blackhole;
+
+/**
+ * Multi-threaded companion to {@link Utf8Benchmark}, exercising the caches' intended thread-safety
+ * contract against the shared caches.
+ *
+ * <p>The point of this variant is to illustrate how {@code recalibrate()} is driven under
+ * concurrency. Unlike the single-threaded case — where the lone serializer thread recalibrates
+ * inline — you would <em>not</em> have every worker recalibrate (that needlessly churns the shared
+ * cache and widens the lookup-then-read race window). Instead a single dedicated thread drives
+ * {@code recalibrate()} while the remaining threads perform lookups. JMH's {@code @Group} /
+ * {@code @GroupThreads} expresses exactly that: 7 lookup threads + 1 recalibrate thread on the same
+ * cache.
+ *
+ * <p>The recalibrate thread runs continuously, which is deliberately more aggressive than a real
+ * periodic cadence — it maximizes contention so this doubles as a concurrency guardrail.
+ */
+@BenchmarkMode(Mode.Throughput)
+@State(Scope.Group)
+public class Utf8ConcurrentBenchmark {
+  static final GenerationalUtf8Cache VALUE_CACHE = new GenerationalUtf8Cache(64, 128);
+  static final SimpleUtf8Cache SIMPLE_VALUE_CACHE = new SimpleUtf8Cache(128);
+
+  @Benchmark
+  @Group("generational")
+  @GroupThreads(7)
+  public void generational_lookup(Blackhole bh) {
+    for (int i = 0; i < NUM_LOOKUPS; ++i) {
+      String tag = nextTag();
+      bh.consume(VALUE_CACHE.getUtf8(nextValue(tag)));
+    }
+  }
+
+  @Benchmark
+  @Group("generational")
+  @GroupThreads(1)
+  public void generational_recalibrate() {
+    VALUE_CACHE.recalibrate();
+  }
+
+  @Benchmark
+  @Group("simple")
+  @GroupThreads(7)
+  public void simple_lookup(Blackhole bh) {
+    for (int i = 0; i < NUM_LOOKUPS; ++i) {
+      String tag = nextTag();
+      bh.consume(SIMPLE_VALUE_CACHE.getUtf8(nextValue(tag)));
+    }
+  }
+
+  @Benchmark
+  @Group("simple")
+  @GroupThreads(1)
+  public void simple_recalibrate() {
+    SIMPLE_VALUE_CACHE.recalibrate();
+  }
+}