Fix lookup-then-read race in GenerationalUtf8Cache.getUtf8

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

@@ -219,32 +219,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/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";