Cover object-store contention and computation failures

field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreDispatchFunctionalTest.java

@@ -2,8 +2,10 @@
 
 import static datadog.instrument.fieldinject.ObjectStoreIds.objectStoreId;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 import java.util.ArrayList;
@@ -22,6 +24,8 @@
 import java.util.function.Function;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /**
  * Loads and actually invokes the generated {@code ObjectStoreDispatch} bytecode, to check its
@@ -213,14 +217,16 @@ void injectedAccessorFallsBackViaWeakGetPutForStoresItDoesNotHave() throws Excep
     assertThat(GlobalObjectStore.get(key, injectedId)).isNull();
   }
 
-  @Test
-  void concurrentGetOrPutOnSameKeyConvergesToASingleWinner() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {"getOrPut", "getOrCompute"})
+  void concurrentInsertionOnSameKeyConvergesToASingleWinner(String operation) throws Exception {
     int storeId = objectStoreId("test.Dispatch.Concurrent.Key", "test.Dispatch.Concurrent.Value");
     Map<Integer, Object> backing = Collections.synchronizedMap(new HashMap<>());
     int threadCount = 8;
     CountDownLatch initialReads = new CountDownLatch(threadCount);
     AtomicInteger getterCalls = new AtomicInteger();
     AtomicInteger setterCalls = new AtomicInteger();
+    AtomicInteger computeCalls = new AtomicInteger();
     InvocationHandler handler =
         (proxy, method, args) -> {
           if (args.length == 1) {
@@ -229,6 +235,8 @@ void concurrentGetOrPutOnSameKeyConvergesToASingleWinner() throws Exception {
               if (!initialReads.await(10, TimeUnit.SECONDS)) {
                 throw new AssertionError("not every contender reached the initial accessor read");
               }
+              // Every caller must take the initial miss, even if another has since stored a value.
+              return null;
             }
             return backing.get((Integer) args[0]);
           }
@@ -240,7 +248,6 @@ void concurrentGetOrPutOnSameKeyConvergesToASingleWinner() throws Exception {
         Proxy.newProxyInstance(
             keyWithValueClass.getClassLoader(), new Class<?>[] {keyWithValueClass}, handler);
 
-    CountDownLatch start = new CountDownLatch(1);
     ExecutorService executor = Executors.newFixedThreadPool(threadCount);
     List<Object> results = new ArrayList<>();
     try {
@@ -250,18 +257,29 @@ void concurrentGetOrPutOnSameKeyConvergesToASingleWinner() throws Exception {
         futures.add(
             executor.submit(
                 () -> {
-                  start.await();
-                  return getOrPutMethod.invoke(null, key, storeId, candidate);
+                  if (operation.equals("getOrPut")) {
+                    return getOrPutMethod.invoke(null, key, storeId, candidate);
+                  }
+                  return getOrComputeMethod.invoke(
+                      null,
+                      key,
+                      storeId,
+                      (Function<Object, Object>)
+                          k -> {
+                            assertThat(k).isSameAs(key);
+                            computeCalls.incrementAndGet();
+                            return candidate;
+                          });
                 }));
       }
-      start.countDown();
       for (java.util.concurrent.Future<Object> future : futures) {
         results.add(future.get(10, TimeUnit.SECONDS));
       }
     } finally {
       executor.shutdownNow();
     }
 
+    assertThat(computeCalls).hasValue(operation.equals("getOrCompute") ? 1 : 0);
     Object winner = backing.get(storeId);
     assertThat(setterCalls)
         .as("only the first contender inside the lock should set the value")
@@ -271,6 +289,81 @@ void concurrentGetOrPutOnSameKeyConvergesToASingleWinner() throws Exception {
         .allMatch(winner::equals);
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = {"compute", "get", "put"})
+  void failureInsideMonitorReleasesLockAndAllowsRetry(String failurePoint) throws Exception {
+    int storeId = objectStoreId("test.Dispatch.Failure.Key", "test.Dispatch.Failure.Value");
+    Map<Integer, Object> backing = Collections.synchronizedMap(new HashMap<>());
+    RuntimeException failure = new IllegalStateException("expected failure");
+    AtomicBoolean failOnce = new AtomicBoolean(true);
+    Object key =
+        Proxy.newProxyInstance(
+            keyWithValueClass.getClassLoader(),
+            new Class<?>[] {keyWithValueClass},
+            (proxy, method, args) -> {
+              boolean getter = args.length == 1;
+              if (Thread.holdsLock(proxy)
+                  && failurePoint.equals(getter ? "get" : "put")
+                  && failOnce.getAndSet(false)) {
+                throw failure;
+              }
+              if (getter) {
+                return backing.get((Integer) args[0]);
+              }
+              backing.put((Integer) args[0], args[1]);
+              return null;
+            });
+    Function<Object, Object> compute =
+        k -> {
+          if (failurePoint.equals("compute") && failOnce.getAndSet(false)) {
+            throw failure;
+          }
+          return "computed";
+        };
+
+    assertThatThrownBy(() -> getOrComputeMethod.invoke(null, key, storeId, compute))
+        .isInstanceOf(InvocationTargetException.class)
+        .hasCause(failure);
+    assertThat(backing).isEmpty();
+
+    // Retry on a different thread so a reentrant acquisition cannot hide a leaked monitor.
+    ExecutorService executor =
+        Executors.newSingleThreadExecutor(
+            r -> {
+              Thread thread = new Thread(r, "dispatch-retry");
+              // A broken monitor exit must not prevent the test JVM from stopping.
+              thread.setDaemon(true);
+              return thread;
+            });
+    try {
+      assertThat(
+              executor
+                  .submit(() -> getOrComputeMethod.invoke(null, key, storeId, compute))
+                  .get(10, TimeUnit.SECONDS))
+          .isEqualTo("computed");
+    } finally {
+      executor.shutdownNow();
+    }
+    assertThat(backing.get(storeId)).isEqualTo("computed");
+  }
+
+  @Test
+  void nullComputedValueIsNotCachedAndCanBeRetried() throws Exception {
+    int storeId = objectStoreId("test.Dispatch.NullCompute.Key", "test.Dispatch.NullCompute.Value");
+    Map<Integer, Object> backing = Collections.synchronizedMap(new HashMap<>());
+    Object key = fastPathKey(backing);
+    AtomicInteger calls = new AtomicInteger();
+    Function<Object, Object> compute = k -> calls.incrementAndGet() <= 2 ? null : "computed";
+
+    assertThat(getOrComputeMethod.invoke(null, key, storeId, compute)).isNull();
+    assertThat(getMethod.invoke(null, key, storeId)).isNull();
+    assertThat(getOrComputeMethod.invoke(null, key, storeId, compute)).isNull();
+    assertThat(getOrComputeMethod.invoke(null, key, storeId, compute)).isEqualTo("computed");
+    assertThat(getOrComputeMethod.invoke(null, key, storeId, compute)).isEqualTo("computed");
+    assertThat(calls).hasValue(3);
+    assertThat(backing.get(storeId)).isEqualTo("computed");
+  }
+
   /**
    * Creates a {@code KeyWithValue} proxy backed by a plain map, simulating a field-injected key.
    */

field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java

@@ -6,6 +6,7 @@
 import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_SOFT_LIMIT;
 import static datadog.instrument.fieldinject.ObjectStoreIds.objectStoreId;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
@@ -18,11 +19,15 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BooleanSupplier;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 class ObjectStoreTest {
 
@@ -100,6 +105,63 @@ void getOrPutInsertsWhenAbsent() {
     assertEquals("new", store.get(key));
   }
 
+  @Test
+  void getOrPutNullDoesNotInsertOrReplaceValue() {
+    Object key = new Object();
+    assertNull(store.getOrPut(key, null));
+    assertNull(store.get(key));
+    assertEquals("first", store.getOrPut(key, "first"));
+    assertEquals("first", store.getOrPut(key, null));
+    assertEquals("first", store.get(key));
+  }
+
+  @ParameterizedTest
+  @ValueSource(strings = {"getOrPut", "getOrCompute"})
+  void concurrentInsertionOnSameKeyReturnsSingleWinner(String operation) throws Exception {
+    Object key = new Object();
+    int threadCount = 8;
+    CountDownLatch ready = new CountDownLatch(threadCount);
+    CountDownLatch start = new CountDownLatch(1);
+    AtomicInteger computeCalls = new AtomicInteger();
+    ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+    List<Future<String>> results = new ArrayList<>();
+    try {
+      for (int i = 0; i < threadCount; i++) {
+        String candidate = "candidate-" + i;
+        results.add(
+            executor.submit(
+                () -> {
+                  ready.countDown();
+                  if (!start.await(10, TimeUnit.SECONDS)) {
+                    throw new AssertionError("contenders did not start");
+                  }
+                  if (operation.equals("getOrPut")) {
+                    return store.getOrPut(key, candidate);
+                  }
+                  return store.getOrCompute(
+                      key,
+                      k -> {
+                        assertSame(key, k);
+                        computeCalls.incrementAndGet();
+                        return candidate;
+                      });
+                }));
+      }
+      assertTrue(ready.await(10, TimeUnit.SECONDS));
+      start.countDown();
+      String winner = results.get(0).get(10, TimeUnit.SECONDS);
+      assertThat(winner).startsWith("candidate-");
+      for (Future<String> result : results) {
+        assertEquals(winner, result.get(10, TimeUnit.SECONDS));
+      }
+      assertEquals(winner, store.get(key));
+      assertEquals(operation.equals("getOrCompute") ? 1 : 0, computeCalls.get());
+    } finally {
+      start.countDown();
+      executor.shutdownNow();
+    }
+  }
+
   // --- getOrCompute ---
 
   @Test
@@ -130,6 +192,49 @@ void getOrComputePassesCorrectKeyToFunction() {
     assertSame(key, received.get());
   }
 
+  @Test
+  void failedComputationLeavesKeyAbsentAndAllowsRetry() {
+    Object key = new Object();
+    RuntimeException failure = new IllegalStateException("expected failure");
+    assertThatThrownBy(
+            () ->
+                store.getOrCompute(
+                    key,
+                    k -> {
+                      throw failure;
+                    }))
+        .isSameAs(failure);
+    assertNull(store.get(key));
+    assertEquals("retry", store.getOrCompute(key, k -> "retry"));
+    assertEquals("retry", store.get(key));
+    assertEquals(
+        "retry",
+        store.getOrCompute(
+            key,
+            k -> {
+              throw new AssertionError("already cached");
+            }));
+  }
+
+  @Test
+  void nullComputedValueIsNotCachedAndCanBeRetried() {
+    Object key = new Object();
+    AtomicInteger calls = new AtomicInteger();
+    for (int i = 0; i < 2; i++) {
+      assertNull(
+          store.getOrCompute(
+              key,
+              k -> {
+                calls.incrementAndGet();
+                return null;
+              }));
+      assertNull(store.get(key));
+    }
+    assertEquals(2, calls.get());
+    assertEquals("retry", store.getOrCompute(key, k -> "retry"));
+    assertEquals("retry", store.get(key));
+  }
+
   // --- isolation ---
 
   @Test
@@ -325,6 +430,7 @@ void singleThreadCyclingAcrossAllShardsStillSamplesEachShard() {
       }
     }
 
+    @SuppressWarnings("unchecked")
     ObjectStore<Object, Integer>[] cycleStores = new ObjectStore[SHARD_COUNT];
     for (int i = 0; i < SHARD_COUNT; i++) {
       cycleStores[i] = ObjectStore.of(chosenKeyTypes[i], "test.ShardCycle.Value");