CASSANDRA-21644 Fix concurrent SAI vector inserts failing once jvector's per-graph pool limit is exceeded

CHANGES.txt

@@ -22,6 +22,7 @@
  * Make cqlsh prompt to reset to no keyspace set by USE after dropping that keyspace (CASSANDRA-21548)
  * Implement CMS rediscovery and recovery protocol (CASSANDRA-20476)
 Merged from 5.0:
+ * Fix concurrent SAI vector inserts failing once jvector's per-graph pool limit is exceeded (CASSANDRA-21644)
  * Unwrap LongType properly when calculating min/max terms in V1SSTableIndex (CASSANDRA-21635)
  * Force repair should ignore min_repair_interval (CASSANDRA-21552)
  * Render SubnetGroups as JSON in system_views.settings (CASSANDRA-21579)

src/java/org/apache/cassandra/db/memtable/Memtable.java

@@ -204,6 +204,17 @@ default long put(PartitionUpdate update, UpdateTransaction indexer, OpOrder.Grou
      */
     long put(PartitionUpdate update, UpdateTransaction indexer, OpOrder.Group opGroup, boolean assumeMissing);
 
+    /**
+     * Whether this memtable guarantees that no more than {@code maxWriters} threads are inside {@link #put} at once,
+     * and therefore that no more than that many threads concurrently apply the update's effect to the {@code indexer}.
+     * Memtables that serialize writes per shard can return true when they have at most {@code maxWriters} shards.
+     * Indexes whose implementation has a limit on concurrent writers can skip their own bounding when this is true.
+     */
+    default boolean limitsConcurrentWritesTo(int maxWriters)
+    {
+        return false;
+    }
+
     // Read operations are provided by the UnfilteredSource interface.
 
     // Statistics

src/java/org/apache/cassandra/db/memtable/Memtable_API.md

@@ -135,6 +135,10 @@ implementations) are provided as the `AbstractMemtable` (statistics tracking), `
 commit log span tracking) and `AbstractAllocatorMemtable` (adds memory management via the `Allocator` class, together
 with flush triggering on memory use and time interval expiration).
 
+A memtable that serializes writes, for example per shard, can say so through `limitsConcurrentWritesTo()`. Secondary
+indexes whose implementation limits the number of concurrent writers use this to skip their own bounding when the
+memtable already provides it. The default answer is false.
+
 The memtable API also gives the memtable some control over flushing and the functioning of the commit log. The former
 is there to permit memtables that operate long-term and/or can handle some events internally, without a need to flush.
 The latter enables memtables that have an internal durability mechanism, such as ones using persistent memory or a

src/java/org/apache/cassandra/db/memtable/ShardedSkipListMemtable.java

@@ -524,6 +524,11 @@ public long put(PartitionUpdate update, UpdateTransaction indexer, OpOrder.Group
             }
         }
 
+        @Override
+        public boolean limitsConcurrentWritesTo(int maxWriters)
+        {
+            return boundaries.shardCount() <= maxWriters;
+        }
     }
 
     public static Factory factory(Map<String, String> optionsCopy)

src/java/org/apache/cassandra/db/memtable/TrieMemtable.java

@@ -176,6 +176,13 @@ public void discard()
         }
     }
 
+    @Override
+    public boolean limitsConcurrentWritesTo(int maxWriters)
+    {
+        // each shard applies one update at a time under its write lock
+        return boundaries.shardCount() <= maxWriters;
+    }
+
     /**
      * Should only be called by ColumnFamilyStore.apply via Keyspace.apply, which supplies the appropriate
      * OpOrdering.

src/java/org/apache/cassandra/index/sai/disk/v1/vector/OnHeapGraph.java

@@ -31,6 +31,10 @@
 import java.util.function.Function;
 import java.util.stream.IntStream;
 
+import javax.annotation.Nullable;
+
+import com.google.common.annotations.VisibleForTesting;
+
 import org.apache.lucene.util.StringHelper;
 import org.cliffc.high_scale_lib.NonBlockingHashMap;
 import org.cliffc.high_scale_lib.NonBlockingHashMapLong;
@@ -51,6 +55,7 @@
 import org.apache.cassandra.io.util.SequentialWriter;
 import org.apache.cassandra.utils.ByteBufferUtil;
 import org.apache.cassandra.utils.CloseableIterator;
+import org.apache.cassandra.utils.concurrent.Semaphore;
 
 import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.graph.GraphIndex;
@@ -72,6 +77,13 @@ public class OnHeapGraph<T>
 
     public static final int MIN_PQ_ROWS = 1024;
 
+    /**
+     * jvector's {@link GraphIndexBuilder#addGraphNode} throws if more than this many threads are inside it at once
+     * (see {@code PoolingSupport}). Uses {@link Runtime} rather than {@code DatabaseDescriptor} because that is what jvector reads.
+     */
+    @VisibleForTesting
+    public static final int MAX_CONCURRENT_GRAPH_INSERTS = Runtime.getRuntime().availableProcessors() + 1;
+
     private final RamAwareVectorValues vectorValues;
     private final GraphIndexBuilder<float[]> builder;
     private final VectorType<?> vectorType;
@@ -80,6 +92,10 @@ public class OnHeapGraph<T>
     private final NonBlockingHashMapLong<VectorPostings<T>> postingsByOrdinal;
     private final NonBlockingHashMap<T, float[]> vectorsByKey;
     private final AtomicInteger nextOrdinal = new AtomicInteger();
+    // memtable inserts are concurrent (CASSANDRA-21160); excess writers wait here rather than fail in jvector.
+    // Null when the memtable itself already limits writers to that many, e.g. a TrieMemtable with few enough shards.
+    @Nullable
+    private final Semaphore graphInsertPermits;
     private volatile boolean hasDeletions;
     private String source;
 
@@ -108,6 +124,9 @@ public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriter
         postingsMap = new ConcurrentSkipListMap<>(Arrays::compare);
         postingsByOrdinal = new NonBlockingHashMapLong<>();
         vectorsByKey = memtable != null ? new NonBlockingHashMap<>() : null;
+        graphInsertPermits = memtable != null && memtable.limitsConcurrentWritesTo(MAX_CONCURRENT_GRAPH_INSERTS)
+                             ? null
+                             : Semaphore.newSemaphore(MAX_CONCURRENT_GRAPH_INSERTS);
 
         builder = new GraphIndexBuilder<>(vectorValues,
                                           VectorEncoding.FLOAT32,
@@ -188,7 +207,17 @@ public long add(ByteBuffer term, T key, InvalidVectorBehavior behavior)
                              : ((CompactionVectorValues) vectorValues).add(ordinal, term);
                 bytesUsed += VectorPostings.emptyBytesUsed() + VectorPostings.bytesPerPosting();
                 postingsByOrdinal.put(ordinal, postings);
-                bytesUsed += builder.addGraphNode(ordinal, vectorValues);
+                if (graphInsertPermits != null)
+                    graphInsertPermits.acquireThrowUncheckedOnInterrupt(1);
+                try
+                {
+                    bytesUsed += builder.addGraphNode(ordinal, vectorValues);
+                }
+                finally
+                {
+                    if (graphInsertPermits != null)
+                        graphInsertPermits.release(1);
+                }
                 return bytesUsed;
             }
             else

test/simulator/test/org/apache/cassandra/simulator/test/VectorMemoryIndexSimulationTest.java

@@ -635,6 +635,7 @@ private static VectorMemoryIndex createVectorMemoryIndex()
                     if ("hashCode".equals(method.getName())) return System.identityHashCode(proxy);
                     if ("toString".equals(method.getName())) return "SimulatedMemtable";
                     if ("equals".equals(method.getName())) return proxy == args[0];
+                    if (method.getReturnType() == boolean.class) return false;
                     return null;
                 }
             );

test/unit/org/apache/cassandra/db/memtable/MemtableConcurrentWriteLimitTest.java

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.db.memtable;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import org.apache.cassandra.config.Config;
+import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.config.InheritingClass;
+import org.apache.cassandra.cql3.CQLTester;
+import org.apache.cassandra.db.ColumnFamilyStore;
+
+import static org.apache.cassandra.db.memtable.AbstractShardedMemtable.SHARDS_OPTION;
+import static org.apache.cassandra.db.memtable.ShardedSkipListMemtable.LOCKING_OPTION;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * {@link Memtable#limitsConcurrentWritesTo} must be true only for memtables that serialize writes per shard, and
+ * only when the shard count is within the limit asked about.
+ */
+public class MemtableConcurrentWriteLimitTest extends CQLTester
+{
+    private static final int SHARDS = 4;
+
+    // Overrides CQLTester.setUpClass so the memtable configurations are registered before the server is prepared
+    @BeforeClass
+    public static void setUpClass()
+    {
+        prePrepareServer();
+
+        LinkedHashMap<String, InheritingClass> memtableConfig = new LinkedHashMap<>();
+        memtableConfig.put("skiplist", new InheritingClass(null, SkipListMemtable.class.getName(), Map.of()));
+        memtableConfig.put("trie", new InheritingClass(null, TrieMemtable.class.getName(), Map.of(SHARDS_OPTION, String.valueOf(SHARDS))));
+        memtableConfig.put("sharded", new InheritingClass(null, ShardedSkipListMemtable.class.getName(), Map.of(SHARDS_OPTION, String.valueOf(SHARDS))));
+        memtableConfig.put("sharded_locking", new InheritingClass(null, ShardedSkipListMemtable.class.getName(), Map.of(SHARDS_OPTION, String.valueOf(SHARDS), LOCKING_OPTION, "true")));
+        DatabaseDescriptor.getRawConfig().memtable = new Config.MemtableOptions();
+        DatabaseDescriptor.getRawConfig().memtable.configurations = memtableConfig;
+
+        prepareServer();
+    }
+
+    @Test
+    public void unshardedMemtableDoesNotLimitWriters()
+    {
+        assertFalse(memtableFor("skiplist").limitsConcurrentWritesTo(Integer.MAX_VALUE));
+    }
+
+    @Test
+    public void shardedSkipListWithoutLockingDoesNotLimitWriters()
+    {
+        assertFalse(memtableFor("sharded").limitsConcurrentWritesTo(Integer.MAX_VALUE));
+    }
+
+    @Test
+    public void trieMemtableLimitsWritersToShardCount()
+    {
+        assertLimitsToShardCount(memtableFor("trie"));
+    }
+
+    @Test
+    public void lockingShardedSkipListLimitsWritersToShardCount()
+    {
+        assertLimitsToShardCount(memtableFor("sharded_locking"));
+    }
+
+    private void assertLimitsToShardCount(Memtable memtable)
+    {
+        // the memtable may get fewer shards than requested, e.g. if local ranges cannot be split that finely
+        int shards = getCurrentColumnFamilyStore().localRangeSplits(SHARDS).shardCount();
+        assertTrue(memtable.limitsConcurrentWritesTo(shards));
+        assertTrue(memtable.limitsConcurrentWritesTo(shards + 1));
+        assertFalse(memtable.limitsConcurrentWritesTo(shards - 1));
+    }
+
+    private Memtable memtableFor(String memtableConfig)
+    {
+        createTable("CREATE TABLE %s (pk int PRIMARY KEY, v int) WITH memtable = '" + memtableConfig + '\'');
+        ColumnFamilyStore cfs = getCurrentColumnFamilyStore();
+        return cfs.getTracker().getView().getCurrentMemtable();
+    }
+}

test/unit/org/apache/cassandra/index/sai/memory/VectorMemoryIndexTest.java

@@ -76,6 +76,7 @@
 import org.apache.cassandra.index.sai.SAITester;
 import org.apache.cassandra.index.sai.StorageAttachedIndex;
 import org.apache.cassandra.index.sai.disk.format.Version;
+import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph;
 import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
 import org.apache.cassandra.index.sai.plan.Expression;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
@@ -110,6 +111,7 @@ public class VectorMemoryIndexTest extends SAITester
 
     private static final double RECALL_THRESHOLD = 0.9;
     private static final int VECTORS_PER_THREAD = 2000;
+    private static final int VECTORS_PER_THREAD_BEYOND_POOL_CAP = 500; // more writers, so fewer vectors each
 
     private ColumnFamilyStore cfs;
     private StorageAttachedIndex index;
@@ -156,8 +158,7 @@ public void setup() throws Throwable
     public void randomQueryTest() throws Exception
     {
         // A non-null memtable tells it to track the mapping from primary key to vector, needed for brute force search
-        Memtable memtable = Mockito.mock(Memtable.class);
-        memtableIndex = new VectorMemoryIndex(index, memtable);
+        memtableIndex = new VectorMemoryIndex(index, mockMemtable(1));
 
         for (int row = 0; row < getRandom().nextIntBetween(1000, 5000); row++)
         {
@@ -260,13 +261,26 @@ public void testExpectedNodesVisitedRespectsBounds()
     @Test
     public void testConcurrentAddsWithRandomVectors() throws Exception
     {
-        testConcurrentAddsAreEventuallyConsistent((threadId, i) -> randomVectorFromThreadLocal());
+        testConcurrentAddsAreEventuallyConsistent(Runtime.getRuntime().availableProcessors(), VECTORS_PER_THREAD, (threadId, i) -> randomVectorFromThreadLocal());
     }
 
     @Test
     public void testConcurrentAddsWithSharedVectors() throws Exception
     {
-        testConcurrentAddsAreEventuallyConsistent((threadId, i) -> makeSharedVector(i));
+        testConcurrentAddsAreEventuallyConsistent(Runtime.getRuntime().availableProcessors(), VECTORS_PER_THREAD, (threadId, i) -> makeSharedVector(i));
+    }
+
+    /**
+     * More writers than jvector's GraphIndexBuilder can serve at once must wait, not fail the insert. The other
+     * concurrent tests use exactly availableProcessors writers, within jvector's limit, so their memtable reports that
+     * it bounds writers and OnHeapGraph skips its semaphore; this test's memtable cannot make that promise, so
+     * OnHeapGraph bounds the writers itself.
+     */
+    @Test
+    public void testConcurrentAddsExceedingJVectorPoolCap() throws Exception
+    {
+        int numThreads = 2 * OnHeapGraph.MAX_CONCURRENT_GRAPH_INSERTS;
+        testConcurrentAddsAreEventuallyConsistent(numThreads, VECTORS_PER_THREAD_BEYOND_POOL_CAP, (threadId, i) -> randomVectorFromThreadLocal());
     }
 
     /**
@@ -280,13 +294,11 @@ public void testConcurrentAddsWithSharedVectors() throws Exception
      * After all writers complete, a full-ring search must return the vast majority of
      * inserted keys with valid scores, confirming no data was lost or corrupted.
      */
-    private void testConcurrentAddsAreEventuallyConsistent(BiFunction<Integer, Integer, ByteBuffer> vectorFactory) throws Exception
+    private void testConcurrentAddsAreEventuallyConsistent(int numThreads, int vectorsPerThread, BiFunction<Integer, Integer, ByteBuffer> vectorFactory) throws Exception
     {
-        Memtable memtable = Mockito.mock(Memtable.class);
-        memtableIndex = new VectorMemoryIndex(index, memtable);
+        memtableIndex = new VectorMemoryIndex(index, mockMemtable(numThreads));
 
-        int numThreads = Runtime.getRuntime().availableProcessors();
-        int totalInserted = numThreads * VECTORS_PER_THREAD;
+        int totalInserted = numThreads * vectorsPerThread;
 
         ExecutorService executor = Executors.newFixedThreadPool(numThreads);
 
@@ -302,9 +314,9 @@ private void testConcurrentAddsAreEventuallyConsistent(BiFunction<Integer, Integ
                 try
                 {
                     barrier.await();
-                    for (int i = 0; i < VECTORS_PER_THREAD; i++)
+                    for (int i = 0; i < vectorsPerThread; i++)
                     {
-                        int pk = threadId * VECTORS_PER_THREAD + i;
+                        int pk = threadId * vectorsPerThread + i;
                         addRow(pk, vectorFactory.apply(threadId, i));
                     }
                 }
@@ -394,11 +406,9 @@ public void testConcurrentAddsAndOrderBySharedVectors() throws Exception
      */
     public void testConcurrentAddsAndOrderByNeverThrow(BiFunction<Integer, Integer, ByteBuffer> vectorFactory) throws Exception
     {
-        Memtable memtable = Mockito.mock(Memtable.class);
-        memtableIndex = new VectorMemoryIndex(index, memtable);
-
         int numWriterThreads = Runtime.getRuntime().availableProcessors();
         int numReaderThreads = Runtime.getRuntime().availableProcessors();
+        memtableIndex = new VectorMemoryIndex(index, mockMemtable(numWriterThreads));
         int totalInserted = numWriterThreads * VECTORS_PER_THREAD;
 
         // Pre-seed enough rows that orderBy() always has a non-empty graph to search,
@@ -558,11 +568,9 @@ public void testConcurrentAddsAndOrderResultsBySharedVectors() throws Exception
      */
     private void testConcurrentAddsAndOrderResultsByNeverThrow(BiFunction<Integer, Integer, ByteBuffer> vectorFactory) throws Exception
     {
-        Memtable memtable = Mockito.mock(Memtable.class);
-        memtableIndex = new VectorMemoryIndex(index, memtable);
-
         int numWriterThreads = Runtime.getRuntime().availableProcessors();
         int numReaderThreads = Runtime.getRuntime().availableProcessors();
+        memtableIndex = new VectorMemoryIndex(index, mockMemtable(numWriterThreads));
         int totalInserted = numWriterThreads * VECTORS_PER_THREAD;
 
         // Pre-seed rows so orderResultsBy() always has a non-empty [minimumKey, maximumKey]
@@ -775,6 +783,18 @@ private void addRow(int pk, ByteBuffer value)
         keyMap.put(key, pk);
     }
 
+    /**
+     * A memtable that answers {@link Memtable#limitsConcurrentWritesTo} truthfully for the number of writer threads the
+     * test will use. With that many writers at or under {@link OnHeapGraph#MAX_CONCURRENT_GRAPH_INSERTS}, OnHeapGraph
+     * relies on the memtable and creates no semaphore; with more, it bounds the writers itself.
+     */
+    private static Memtable mockMemtable(int writers)
+    {
+        Memtable memtable = Mockito.mock(Memtable.class);
+        Mockito.when(memtable.limitsConcurrentWritesTo(Mockito.anyInt())).thenAnswer(invocation -> writers <= (int) invocation.getArgument(0));
+        return memtable;
+    }
+
     private DecoratedKey makeKey(TableMetadata table, Integer partitionKey)
     {
         ByteBuffer key = table.partitionKeyType.fromString(partitionKey.toString());