NoSQL: fix concurrency issue in IndexImpl + add tests + adopt JMH

Author: snazy

components/persistence/index/src/jmh/java/org/apache/polaris/persistence/indexes/RandomUuidKeyIndexImplBench.java

@@ -44,7 +44,7 @@
 @Warmup(iterations = 3, time = 2000, timeUnit = MILLISECONDS)
 @Measurement(iterations = 5, time = 1000, timeUnit = MILLISECONDS)
 @Fork(1)
-@Threads(1) // Do NOT use multiple threads StoreIndex is NOT thread safe!
+@Threads(4)
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(MICROSECONDS)
 public class RandomUuidKeyIndexImplBench {

components/persistence/index/src/jmh/java/org/apache/polaris/persistence/indexes/RealisticKeyIndexImplBench.java

@@ -48,7 +48,7 @@
 @Warmup(iterations = 3, time = 2000, timeUnit = MILLISECONDS)
 @Measurement(iterations = 5, time = 1000, timeUnit = MILLISECONDS)
 @Fork(1)
-@Threads(1) // Do NOT use multiple threads StoreIndex is NOT thread safe!
+@Threads(4)
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(MICROSECONDS)
 public class RealisticKeyIndexImplBench {

components/persistence/index/src/main/java/org/apache/polaris/persistence/indexes/IndexImpl.java

@@ -516,7 +516,7 @@ public int estimatedSerializedSize() {
 
       target = target.flip();
     } else {
-      target = serialized.duplicate().position(0).limit(originalSerializedSize);
+      target = serializedThreadSafe().position(0).limit(originalSerializedSize);
     }
 
     return target;
@@ -603,8 +603,8 @@ private IndexImpl(ByteBuffer serialized, IndexValueSerializer<V> ser) {
 
     this.elements = elements;
     this.serializer = ser;
-    this.serialized = serialized.duplicate().clear();
     this.originalSerializedSize = serialized.position();
+    this.serialized = serialized.duplicate().clear();
   }
 
   /**
@@ -708,14 +708,13 @@ ByteBuffer serializeKey(ByteBuffer keySerBuffer, ByteBuffer previousKey) {
         previousKey.limit(limitSave).position(0);
       }
 
-      IndexImpl<V> index = IndexImpl.this;
-      var serialized = requireNonNull(index.serialized);
-      return keySerBuffer.put(serialized.limit(valueOffset).position(keyOffset)).flip();
+      return keySerBuffer
+          .put(serializedNotThreadSafe().limit(valueOffset).position(keyOffset))
+          .flip();
     }
 
     private IndexKey materializeKey() {
-      IndexImpl<V> index = IndexImpl.this;
-      var serialized = requireNonNull(index.serialized);
+      var serialized = serializedThreadSafe();
 
       var suffix = serialized.limit(valueOffset).position(keyOffset);
 
@@ -761,19 +760,9 @@ private ByteBuffer prefixKey(ByteBuffer serialized, LazyIndexElement me, int rem
       return keyBuffer;
     }
 
-    private ByteBuffer serializedForContent() {
-      IndexImpl<V> index = IndexImpl.this;
-      var serialized = requireNonNull(index.serialized);
-      return serialized.limit(endOffset).position(valueOffset);
-    }
-
-    private V materializeContent() {
-      return serializer.deserialize(serializedForContent());
-    }
-
     @Override
     public void serializeContent(IndexValueSerializer<V> ser, ByteBuffer target) {
-      target.put(serializedForContent());
+      target.put(serializedNotThreadSafe().limit(endOffset).position(valueOffset));
     }
 
     @Override
@@ -800,7 +789,10 @@ public V getValue() {
       var c = content;
       if (c == null) {
         if (!hasContent) {
-          c = content = materializeContent();
+          c =
+              content =
+                  serializer.deserialize(
+                      serializedThreadSafe().limit(endOffset).position(valueOffset));
           hasContent = true;
         }
       }
@@ -842,6 +834,32 @@ static ByteBuffer newKeyBuffer() {
     return ByteBuffer.allocate(MAX_KEY_BYTES);
   }
 
+  /**
+   * Non-thread-safe variant to retrieve {@link #serialized}. Used in these scenarios, which are not
+   * thread safe by contract:
+   *
+   * <ul>
+   *   <li>serializing a <em>modified</em>
+   * </ul>
+   */
+  private ByteBuffer serializedNotThreadSafe() {
+    return requireNonNull(serialized);
+  }
+
+  /**
+   * Thread-safe variant to retrieve {@link #serialized}. Used in these scenarios, which are not
+   * thread safe by contract:
+   *
+   * <ul>
+   *   <li>serializing a non-modified index
+   *   <li>lazy materialization of a key (deserialization)
+   *   <li>lazy materialization of a value (deserialization)
+   * </ul>
+   */
+  private ByteBuffer serializedThreadSafe() {
+    return requireNonNull(serialized).duplicate();
+  }
+
   private static <V> int search(List<IndexElement<V>> e, @Nonnull IndexKey key) {
     // Need a StoreIndexElement for the sake of 'binarySearch()' (the content value isn't used)
     return search(e, indexElement(key, ""));

components/persistence/index/src/test/java/org/apache/polaris/persistence/indexes/TestIndexImpl.java

@@ -43,6 +43,13 @@
 import java.util.HexFormat;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -785,4 +792,88 @@ public void stateRelated() {
     soft.assertThat(index.isMutable()).isTrue();
     soft.assertThatCode(() -> index.divide(3)).doesNotThrowAnyException();
   }
+
+  // The following multithreaded "tests" are only there to verify that no ByteBuffer related
+  // exceptions are thrown.
+
+  @Test
+  public void multithreadedGetKey() throws Exception {
+    multithreaded(KeyIndexTestSet::randomGetKey, true);
+  }
+
+  @Test
+  public void multithreadedSerialize() throws Exception {
+    multithreaded(KeyIndexTestSet::serialize, false);
+  }
+
+  @Test
+  public void multithreadedFirst() throws Exception {
+    multithreaded(ts -> ts.keyIndex().first(), false);
+  }
+
+  @Test
+  public void multithreadedLast() throws Exception {
+    multithreaded(ts -> ts.keyIndex().last(), false);
+  }
+
+  @Test
+  public void multithreadedKeys() throws Exception {
+    multithreaded(ts -> ts.keyIndex().asKeyList(), false);
+  }
+
+  @Test
+  public void multithreadedElementIterator() throws Exception {
+    multithreaded(ts -> ts.keyIndex().elementIterator().forEachRemaining(el -> {}), false);
+  }
+
+  @Test
+  public void multithreadedIterator() throws Exception {
+    multithreaded(ts -> ts.keyIndex().iterator().forEachRemaining(el -> {}), false);
+  }
+
+  void multithreaded(Consumer<KeyIndexTestSet<ObjRef>> worker, boolean longTest) throws Exception {
+    var indexTestSet =
+        KeyIndexTestSet.<ObjRef>newGenerator()
+            .keySet(ImmutableRandomUuidKeySet.builder().numKeys(100_000).build())
+            .elementSupplier(key -> indexElement(key, Util.randomObjId()))
+            .elementSerializer(OBJ_REF_SERIALIZER)
+            .build()
+            .generateIndexTestSet();
+
+    var threads = Runtime.getRuntime().availableProcessors();
+
+    try (var executor = Executors.newFixedThreadPool(threads)) {
+      var latch = new CountDownLatch(threads);
+      var start = new Semaphore(0);
+      var stop = new AtomicBoolean();
+
+      var futures =
+          IntStream.range(0, threads)
+              .mapToObj(
+                  i ->
+                      CompletableFuture.runAsync(
+                          () -> {
+                            latch.countDown();
+                            try {
+                              start.acquire();
+                            } catch (InterruptedException e) {
+                              throw new RuntimeException(e);
+                            }
+                            while (!stop.get()) {
+                              worker.accept(indexTestSet);
+                            }
+                          },
+                          executor))
+              .toArray(CompletableFuture[]::new);
+
+      latch.await();
+      start.release(threads);
+
+      Thread.sleep(longTest ? TimeUnit.SECONDS.toMillis(3) : 500L);
+
+      stop.set(true);
+
+      CompletableFuture.allOf(futures).get();
+    }
+  }
 }