Fix race condition

Author: henningandersen

x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/KeyMapping.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.blobcache.shared;
 
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
 
@@ -37,12 +38,13 @@ public Value get(Key1 key1, Key2 key2) {
      * @return the resulting value.
      */
     public Value computeIfAbsent(Key1 key1, Key2 key2, Function<Key2, Value> function) {
-        var inner = mapping.compute(key1, (k, current) -> {
+        AtomicReference<Value> result = new AtomicReference<>();
+        mapping.compute(key1, (k, current) -> {
             ConcurrentHashMap<Key2, Value> map = current == null ? new ConcurrentHashMap<>() : current;
-            map.computeIfAbsent(key2, function);
+            result.setPlain(map.computeIfAbsent(key2, function));
             return map;
         });
-        return inner.get(key2);
+        return result.getPlain();
     }
 
     public boolean remove(Key1 key1, Key2 key2, Value value) {

x-pack/plugin/blob-cache/src/test/java/org/elasticsearch/blobcache/shared/KeyMappingTests.java

@@ -62,7 +62,7 @@ public void testMultiThreaded() {
             for (int j = 0; j < 1000; ++j) {
                 Integer finalJ = j;
                 assertNull(mapping.get(k1, k2));
-                mapping.computeIfAbsent(k1, k2, (kx) -> finalJ);
+                assertSame(finalJ, mapping.computeIfAbsent(k1, k2, (kx) -> finalJ));
                 assertEquals(finalJ, mapping.get(k1, k2));
                 assertTrue(mapping.remove(k1, k2, finalJ));
                 if ((j & 1) == 0) {
@@ -84,4 +84,30 @@ public void testMultiThreaded() {
 
         assertEquals(Set.of(), mapping.key1s());
     }
+
+    public void testMultiThreadedSameKey() {
+        final String k1 = randomAlphanumericOfLength(10);
+        KeyMapping<String, String, Integer> mapping = new KeyMapping<>();
+
+        List<Thread> threads = IntStream.range(0, 10).mapToObj(i -> new Thread(() -> {
+            for (int j = 0; j < 1000; ++j) {
+                Integer computeValue = i * 1000 + j;
+                Integer value = mapping.computeIfAbsent(k1, k1, (kx) -> computeValue);
+                assertNotNull(value);
+                // either our value or another threads value.
+                assertTrue(value == computeValue || value / 1000 != i);
+                mapping.remove(k1, k1, value);
+            }
+        })).toList();
+        threads.forEach(Thread::start);
+        threads.forEach(t -> {
+            try {
+                t.join(10000);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        });
+
+        assertEquals(Set.of(), mapping.key1s());
+    }
 }