diff --git a/CHANGES.txt b/CHANGES.txt index 4dfb429438..3fa97b0d1c 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -53,6 +53,7 @@ Merged from 6.0: * Introduce minimum_threshold for data resurrection startup check (CASSANDRA-21293) * Synchronously publish changes to local gossip state following metadata updates (CASSANDRA-21239) Merged from 5.0: + * Make synchronization on VectorMemoryIndex inserts more granular (CASSANDRA-21160) * Fix RequestFailureReason serializer and nits in a few others (CASSANDRA-21437) * Remove golang dependency in gen-doc and replace with python implementation (CASSANDRA-21432) * Use estimated compressed size for tables to check if there is enough free space for a compaction (CASSANDRA-21245) diff --git a/src/java/org/apache/cassandra/index/sai/disk/v1/vector/PrimaryKeyWithScore.java b/src/java/org/apache/cassandra/index/sai/disk/v1/vector/PrimaryKeyWithScore.java index c81c4ebd20..b582f3451d 100644 --- a/src/java/org/apache/cassandra/index/sai/disk/v1/vector/PrimaryKeyWithScore.java +++ b/src/java/org/apache/cassandra/index/sai/disk/v1/vector/PrimaryKeyWithScore.java @@ -18,6 +18,8 @@ package org.apache.cassandra.index.sai.disk.v1.vector; +import com.google.common.annotations.VisibleForTesting; + import org.apache.cassandra.db.CellSourceIdentifier; import org.apache.cassandra.db.rows.Cell; import org.apache.cassandra.db.rows.Row; @@ -52,6 +54,12 @@ public class PrimaryKeyWithScore implements Comparable return primaryKey; } + @VisibleForTesting + public float score() + { + return indexScore; + } + public boolean isIndexDataValid(Row row, long nowInSecs) { // If the indexed column is part of the primary key, we don't need this type of validation because we would have diff --git a/src/java/org/apache/cassandra/index/sai/memory/VectorMemoryIndex.java b/src/java/org/apache/cassandra/index/sai/memory/VectorMemoryIndex.java index b976b2038a..bf64396486 100644 --- a/src/java/org/apache/cassandra/index/sai/memory/VectorMemoryIndex.java +++ b/src/java/org/apache/cassandra/index/sai/memory/VectorMemoryIndex.java @@ -75,8 +75,7 @@ public class VectorMemoryIndex extends MemoryIndex private final Memtable memtable; private final LongAdder writeCount = new LongAdder(); - private PrimaryKey minimumKey; - private PrimaryKey maximumKey; + private volatile KeyBounds keyBounds; private final NavigableSet primaryKeys = new ConcurrentSkipListSet<>(); @@ -88,7 +87,7 @@ public class VectorMemoryIndex extends MemoryIndex } @Override - public synchronized long add(DecoratedKey key, Clustering clustering, ByteBuffer value) + public long add(DecoratedKey key, Clustering clustering, ByteBuffer value) { if (value == null || value.remaining() == 0 || !index.validateTermSize(key, value, false, null)) return 0; @@ -100,11 +99,11 @@ public class VectorMemoryIndex extends MemoryIndex private long index(PrimaryKey primaryKey, ByteBuffer value) { - updateKeyBounds(primaryKey); - + long bytesUsed = graph.add(value, primaryKey, OnHeapGraph.InvalidVectorBehavior.FAIL); writeCount.increment(); primaryKeys.add(primaryKey); - return graph.add(value, primaryKey, OnHeapGraph.InvalidVectorBehavior.FAIL); + updateKeyBounds(primaryKey); + return bytesUsed; } @Override @@ -131,9 +130,6 @@ public class VectorMemoryIndex extends MemoryIndex { PrimaryKey primaryKey = index.hasClustering() ? index.keyFactory().create(key, clustering) : index.keyFactory().create(key); - // update bounds because only rows with vectors are included in the key bounds, - // so if the vector was null before, we won't have included it - updateKeyBounds(primaryKey); // make the changes in this order, so we don't have a window where the row is not in the index at all if (newRemaining > 0) @@ -144,20 +140,18 @@ public class VectorMemoryIndex extends MemoryIndex // remove primary key if it's no longer indexed if (newRemaining <= 0 && oldRemaining > 0) primaryKeys.remove(primaryKey); + + // update bounds because only rows with vectors are included in the key bounds, + // so if the vector was null before, we won't have included it + updateKeyBounds(primaryKey); } return bytesUsed; } - private void updateKeyBounds(PrimaryKey primaryKey) + private synchronized void updateKeyBounds(PrimaryKey key) { - if (minimumKey == null) - minimumKey = primaryKey; - else if (primaryKey.compareTo(minimumKey) < 0) - minimumKey = primaryKey; - if (maximumKey == null) - maximumKey = primaryKey; - else if (primaryKey.compareTo(maximumKey) > 0) - maximumKey = primaryKey; + KeyBounds current = keyBounds; + keyBounds = current == null ? new KeyBounds(key, key) : current.withUpdated(key); } @Override @@ -213,15 +207,15 @@ public class VectorMemoryIndex extends MemoryIndex @Override public CloseableIterator orderResultsBy(QueryContext queryContext, List results, Expression orderer) { - if (minimumKey == null) - // This case implies maximumKey is empty too. + KeyBounds bounds = keyBounds; + if (bounds == null) return CloseableIterator.empty(); int limit = queryContext.limit(); List resultsInRange = results.stream() - .dropWhile(k -> k.compareTo(minimumKey) < 0) - .takeWhile(k -> k.compareTo(maximumKey) <= 0) + .dropWhile(k -> k.compareTo(bounds.minimum) < 0) + .takeWhile(k -> k.compareTo(bounds.maximum) <= 0) .collect(Collectors.toList()); int maxBruteForceRows = maxBruteForceRows(limit, resultsInRange.size(), graph.size()); @@ -420,4 +414,25 @@ public class VectorMemoryIndex extends MemoryIndex FileUtils.closeQuietly(nodeScores); } } + + private static final class KeyBounds + { + final PrimaryKey minimum; + final PrimaryKey maximum; + + KeyBounds(PrimaryKey minimum, PrimaryKey maximum) + { + this.minimum = minimum; + this.maximum = maximum; + } + + KeyBounds withUpdated(PrimaryKey key) + { + PrimaryKey newMin = minimum.compareTo(key) > 0 ? key : minimum; + PrimaryKey newMax = maximum.compareTo(key) < 0 ? key : maximum; + + // Avoid allocation if nothing changed + return newMin == minimum && newMax == maximum ? this : new KeyBounds(newMin, newMax); + } + } } diff --git a/test/simulator/main/org/apache/cassandra/simulator/asm/NemesisFieldSelectors.java b/test/simulator/main/org/apache/cassandra/simulator/asm/NemesisFieldSelectors.java index 0e1afa622f..d26b87f666 100644 --- a/test/simulator/main/org/apache/cassandra/simulator/asm/NemesisFieldSelectors.java +++ b/test/simulator/main/org/apache/cassandra/simulator/asm/NemesisFieldSelectors.java @@ -18,9 +18,9 @@ package org.apache.cassandra.simulator.asm; -import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import java.util.concurrent.atomic.AtomicLongFieldUpdater; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; @@ -39,10 +39,13 @@ import static org.apache.cassandra.simulator.asm.NemesisFieldKind.SIMPLE; /** * Define classes that receive special handling. * At present all instance methods invoked on such classes have nemesis points inserted either side of them. + * + * Tests that need nemesis behavior on fields without annotating the source class can use + * {@link #register(String, String, NemesisFieldKind)} to dynamically add entries. */ public class NemesisFieldSelectors { - public static final Map> classToFieldToNemesis; + public static final ConcurrentHashMap> classToFieldToNemesis; static { @@ -53,12 +56,49 @@ public class NemesisFieldSelectors Stream.of(AtomicIntegerFieldUpdater.class, AtomicLongFieldUpdater.class, AtomicReferenceFieldUpdater.class) .forEach(c -> byClass.put(c, NemesisFieldKind.ATOMICUPDATERX)); - Map> byField = new HashMap<>(); + ConcurrentHashMap> byField = new ConcurrentHashMap<>(); new Reflections(ConfigurationBuilder.build("org.apache.cassandra").addScanners(new FieldAnnotationsScanner())) .getFieldsAnnotatedWith(Nemesis.class) - .forEach(field -> byField.computeIfAbsent(dotsToSlashes(field.getDeclaringClass()), ignore -> new HashMap<>()) + .forEach(field -> byField.computeIfAbsent(dotsToSlashes(field.getDeclaringClass()), ignore -> new ConcurrentHashMap<>()) .put(field.getName(), byClass.getOrDefault(field.getType(), SIMPLE))); - classToFieldToNemesis = Collections.unmodifiableMap(byField); + classToFieldToNemesis = byField; + } + + /** + * Register a field for nemesis handling without requiring a {@link Nemesis} annotation on the source class. + * This allows tests to opt-in fields from classes they do not own. + * + * @param className the internal class name (slashes, e.g. "org/apache/cassandra/index/sai/disk/v1/vector/OnHeapGraph") + * @param fieldName the field name as declared in the class + * @param kind the nemesis field kind (typically {@link NemesisFieldKind#SIMPLE} for plain volatile fields, + * {@link NemesisFieldKind#ATOMICX} for AtomicInteger/AtomicLong/AtomicReference/AtomicBoolean fields) + */ + public static void register(String className, String fieldName, NemesisFieldKind kind) + { + classToFieldToNemesis.computeIfAbsent(className, ignore -> new ConcurrentHashMap<>()) + .put(fieldName, kind); + } + + /** + * Register a field for nemesis handling using the class object directly. + * + * @param clazz the class owning the field + * @param fieldName the field name as declared in the class + * @param kind the nemesis field kind + */ + public static void register(Class clazz, String fieldName, NemesisFieldKind kind) + { + register(dotsToSlashes(clazz), fieldName, kind); + } + + /** + * Remove a previously registered nemesis field. Useful for test cleanup. + */ + public static void unregister(Class clazz, String fieldName) + { + Map fields = classToFieldToNemesis.get(dotsToSlashes(clazz)); + if (fields != null) + fields.remove(fieldName); } public static NemesisFieldKind.Selector get() diff --git a/test/simulator/test/org/apache/cassandra/simulator/test/OnHeapGraphSimulationTest.java b/test/simulator/test/org/apache/cassandra/simulator/test/OnHeapGraphSimulationTest.java new file mode 100644 index 0000000000..04d9120ed3 --- /dev/null +++ b/test/simulator/test/org/apache/cassandra/simulator/test/OnHeapGraphSimulationTest.java @@ -0,0 +1,283 @@ +/* + * 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.simulator.test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.atomic.AtomicInteger; + +import org.junit.BeforeClass; +import org.junit.Test; + +import org.apache.cassandra.concurrent.ExecutorFactory; +import org.apache.cassandra.concurrent.ExecutorPlus; +import org.apache.cassandra.db.marshal.FloatType; +import org.apache.cassandra.db.marshal.VectorType; +import org.apache.cassandra.db.memtable.Memtable; +import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig; +import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph; +import org.apache.cassandra.index.sai.disk.v1.vector.VectorPostings; +import org.apache.cassandra.simulator.asm.NemesisFieldKind; +import org.apache.cassandra.simulator.asm.NemesisFieldSelectors; +import org.apache.cassandra.utils.CloseableIterator; + +import io.github.jbellis.jvector.graph.SearchResult; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; + +/** + * Simulation test for {@link OnHeapGraph} that uses the simulator's nemesis framework to inject + * adversarial scheduling around mutable field accesses, without modifying the source class. + *

+ * This test registers fields from {@link OnHeapGraph} and {@link VectorPostings} as nemesis targets + * via {@link NemesisFieldSelectors#register(Class, String, NemesisFieldKind)} so the bytecode + * transformer inserts scheduling perturbation points around those field accesses. + *

+ * The test exercises concurrent add + search workloads under the simulator, which exposes + * ordering-dependent bugs that are difficult to trigger with plain threads. + */ +public class OnHeapGraphSimulationTest extends SimulationTestBase +{ + private static final int DIMENSIONS = 8; + private static final int VECTORS_PER_THREAD = 200; + private static final int NUM_THREADS = 4; + + @BeforeClass + public static void registerNemesisFields() + { + // OnHeapGraph mutable fields + NemesisFieldSelectors.register(OnHeapGraph.class, "hasDeletions", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "postingsMap", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "postingsByOrdinal", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "vectorsByKey", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "vectorValues", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "builder", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "nextOrdinal", NemesisFieldKind.ATOMICX); + + // VectorPostings mutable fields + NemesisFieldSelectors.register(VectorPostings.class, "ordinal", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(VectorPostings.class, "postings", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(VectorPostings.class, "rowIds", NemesisFieldKind.SIMPLE); + } + + /** + * Concurrent adds under nemesis scheduling: multiple threads insert vectors into the same + * OnHeapGraph and the simulator adversarially reorders field accesses. After all inserts + * complete, a search must return the vast majority of inserted vectors (ANN recall tolerance). + */ + @Test + public void testConcurrentAddsUnderNemesis() + { + int totalInserted = NUM_THREADS * VECTORS_PER_THREAD; + + simulate(() -> { + OnHeapGraph graph = createGraph(); + AtomicInteger keyCounter = new AtomicInteger(0); + + ExecutorPlus executor = ExecutorFactory.Global.executorFactory().pooled("writers", NUM_THREADS); + + for (int t = 0; t < NUM_THREADS; t++) + { + executor.submit(() -> { + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int key = keyCounter.getAndIncrement(); + ByteBuffer vector = randomVector(DIMENSIONS); + graph.add(vector, key, OnHeapGraph.InvalidVectorBehavior.FAIL); + } + }); + } + + SharedGraphHolder.graph = graph; + SharedGraphHolder.totalInserted = totalInserted; + }, () -> { + @SuppressWarnings("unchecked") + OnHeapGraph graph = (OnHeapGraph) SharedGraphHolder.graph; + if (graph == null) + throw new AssertionError("Graph was not created"); + + int total = SharedGraphHolder.totalInserted; + int graphSize = graph.size(); + + // The graph should contain ALL inserted vectors - no data loss from concurrent adds. + if (graphSize != total) + throw new AssertionError(String.format( + "Graph size %d != total inserted %d — ordinal collision detected (lost %d vectors)", + graphSize, total, total - graphSize)); + }, DEFAULT_ITERATIONS); + } + + /** + * Concurrent adds and searches under nemesis: writers insert while readers search the graph. + * The simulator will adversarially schedule field accesses to expose races between + * hasDeletions reads in search() and writes in add()/remove(). + */ + @Test + public void testConcurrentAddsAndSearchesUnderNemesis() + { + simulate(() -> { + OnHeapGraph graph = createGraph(); + + // Pre-seed some vectors so search always has something to traverse + for (int i = 0; i < 50; i++) + { + ByteBuffer vector = randomVector(DIMENSIONS); + graph.add(vector, -(i + 1), OnHeapGraph.InvalidVectorBehavior.FAIL); + } + + ExecutorPlus executor = ExecutorFactory.Global.executorFactory().pooled("mixed", NUM_THREADS * 2); + + // Writers + AtomicInteger keyCounter = new AtomicInteger(0); + for (int t = 0; t < NUM_THREADS; t++) + { + executor.submit(() -> { + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int key = keyCounter.getAndIncrement(); + ByteBuffer vector = randomVector(DIMENSIONS); + graph.add(vector, key, OnHeapGraph.InvalidVectorBehavior.FAIL); + } + }); + } + + // Readers: search while writers are in progress + for (int t = 0; t < NUM_THREADS; t++) + { + executor.submit(() -> { + for (int i = 0; i < 10; i++) + { + float[] queryVector = new float[DIMENSIONS]; + for (int d = 0; d < DIMENSIONS; d++) + queryVector[d] = ThreadLocalRandom.current().nextFloat(); + + // Must not throw — safety invariant + try (CloseableIterator results = graph.search(queryVector, 50, new Bits.MatchAllBits(graph.size()))) + { + while (results.hasNext()) + { + SearchResult.NodeScore ns = results.next(); + if (!Float.isFinite(ns.score)) + throw new AssertionError("Non-finite score during concurrent search: " + ns.score); + } + } + } + }); + } + }, () -> {}, DEFAULT_ITERATIONS); + } + + /** + * Concurrent adds and removes under nemesis: exercises the hasDeletions volatile flag + * and the interaction between add() creating postings and remove() marking them deleted. + */ + @Test + public void testConcurrentAddsAndRemovesUnderNemesis() + { + simulate(() -> { + OnHeapGraph graph = createGraph(); + int insertCount = NUM_THREADS * VECTORS_PER_THREAD; + + // First, insert all vectors + List vectors = new ArrayList<>(insertCount); + for (int i = 0; i < insertCount; i++) + { + ByteBuffer vector = randomVector(DIMENSIONS); + vectors.add(vector); + graph.add(vector, i, OnHeapGraph.InvalidVectorBehavior.FAIL); + } + + ExecutorPlus executor = ExecutorFactory.Global.executorFactory().pooled("removers", NUM_THREADS); + + // Remove half the vectors concurrently + AtomicInteger removeCounter = new AtomicInteger(0); + for (int t = 0; t < NUM_THREADS; t++) + { + executor.submit(() -> { + int idx; + while ((idx = removeCounter.getAndIncrement()) < insertCount) + { + if (idx % 2 == 0) + graph.remove(vectors.get(idx), idx); + } + }); + } + + // After removes, search should still work without exceptions + float[] queryVector = new float[DIMENSIONS]; + for (int d = 0; d < DIMENSIONS; d++) + queryVector[d] = ThreadLocalRandom.current().nextFloat(); + + try (CloseableIterator results = graph.search(queryVector, 100, new Bits.MatchAllBits(graph.size()))) + { + while (results.hasNext()) + { + SearchResult.NodeScore ns = results.next(); + if (!Float.isFinite(ns.score)) + throw new AssertionError("Non-finite score with deletions: " + ns.score); + } + } + }, () -> {}, DEFAULT_ITERATIONS); + } + + @SuppressWarnings("unchecked") + private static OnHeapGraph createGraph() + { + VectorType vectorType = VectorType.getInstance(FloatType.instance, DIMENSIONS); + IndexWriterConfig config = new IndexWriterConfig( + IndexWriterConfig.DEFAULT_MAXIMUM_NODE_CONNECTIONS, + IndexWriterConfig.DEFAULT_CONSTRUCTION_BEAM_WIDTH, + VectorSimilarityFunction.DOT_PRODUCT, + null + ); + // Use a JDK proxy — Mockito cannot operate inside the InstanceClassLoader. + // OnHeapGraph only checks memtable != null and calls getClass().getSimpleName() + hashCode(). + Memtable memtable = (Memtable) java.lang.reflect.Proxy.newProxyInstance( + Memtable.class.getClassLoader(), + new Class[]{ Memtable.class }, + (proxy, method, args) -> { + 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]; + return null; + } + ); + return new OnHeapGraph<>(vectorType, config, memtable); + } + + private static ByteBuffer randomVector(int dimensions) + { + List rawVector = new ArrayList<>(dimensions); + for (int i = 0; i < dimensions; i++) + rawVector.add(ThreadLocalRandom.current().nextFloat()); + return VectorType.getInstance(FloatType.instance, dimensions).getSerializer().serialize(rawVector); + } + + /** + * Static holder so the graph created inside the simulated classloader can be shared + * between the action runnables and the check runnable. + */ + public static class SharedGraphHolder + { + public static volatile OnHeapGraph graph; + public static volatile int totalInserted; + } +} diff --git a/test/simulator/test/org/apache/cassandra/simulator/test/VectorMemoryIndexSimulationTest.java b/test/simulator/test/org/apache/cassandra/simulator/test/VectorMemoryIndexSimulationTest.java new file mode 100644 index 0000000000..b3690dc508 --- /dev/null +++ b/test/simulator/test/org/apache/cassandra/simulator/test/VectorMemoryIndexSimulationTest.java @@ -0,0 +1,750 @@ +/* + * 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.simulator.test; + +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import org.junit.BeforeClass; +import org.junit.Test; + +import org.apache.cassandra.cql3.Operator; +import org.apache.cassandra.cql3.statements.schema.IndexTarget; +import org.apache.cassandra.db.Clustering; +import org.apache.cassandra.db.ClusteringComparator; +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.DataRange; +import org.apache.cassandra.db.DecoratedKey; +import org.apache.cassandra.db.PartitionPosition; +import org.apache.cassandra.db.PartitionRangeReadCommand; +import org.apache.cassandra.db.filter.ColumnFilter; +import org.apache.cassandra.db.filter.DataLimits; +import org.apache.cassandra.db.filter.RowFilter; +import org.apache.cassandra.db.marshal.FloatType; +import org.apache.cassandra.db.marshal.Int32Type; +import org.apache.cassandra.db.marshal.VectorType; +import org.apache.cassandra.db.memtable.Memtable; +import org.apache.cassandra.dht.AbstractBounds; +import org.apache.cassandra.dht.Murmur3Partitioner; +import org.apache.cassandra.dht.Range; +import org.apache.cassandra.index.sai.QueryContext; +import org.apache.cassandra.index.sai.StorageAttachedIndex; +import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig; +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.disk.v1.vector.VectorPostings; +import org.apache.cassandra.index.sai.memory.VectorMemoryIndex; +import org.apache.cassandra.index.sai.plan.Expression; +import org.apache.cassandra.index.sai.utils.IndexTermType; +import org.apache.cassandra.index.sai.utils.PrimaryKey; +import org.apache.cassandra.schema.CachingParams; +import org.apache.cassandra.schema.ColumnMetadata; +import org.apache.cassandra.schema.TableMetadata; +import org.apache.cassandra.simulator.asm.NemesisFieldKind; +import org.apache.cassandra.simulator.asm.NemesisFieldSelectors; +import org.apache.cassandra.utils.CloseableIterator; + +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import sun.misc.Unsafe; + +/** + * Simulation test for {@link VectorMemoryIndex} that exercises the real {@code add()} and + * {@code orderBy()} / {@code orderResultsBy()} code paths under adversarial nemesis scheduling. + *

+ * The nemesis framework injects scheduling perturbation around mutable field accesses in + * {@link OnHeapGraph}, {@link VectorPostings}, and {@link VectorMemoryIndex}. + *

+ * Because the simulator's {@code InstanceClassLoader} cannot bootstrap full Cassandra infrastructure + * (DatabaseDescriptor, Keyspace, ColumnFamilyStore), this test constructs a minimal + * {@link StorageAttachedIndex} via {@code Unsafe.allocateInstance()} with only the fields + * that {@link VectorMemoryIndex} actually accesses set reflectively. + *

+ * Tests ported from {@link org.apache.cassandra.index.sai.memory.VectorMemoryIndexTest} + *

    + *
  • {@link #testConcurrentAddsWithRandomVectors()} — N writers with random vectors, verify no data loss
    • + *
  • {@link #testConcurrentAddsWithSharedVectors()} — N writers with shared vectors, verify no data loss
    • + *
  • {@link #testConcurrentAddsAndOrderByRandomVectors()} — writers + readers via orderBy() with random vectors
    • + *
  • {@link #testConcurrentAddsAndOrderBySharedVectors()} — writers + readers via orderBy() with shared vectors
    • + *
  • {@link #testConcurrentAddsAndOrderResultsByRandomVectors()} — writers + readers via orderResultsBy() with random vectors
    • + *
  • {@link #testConcurrentAddsAndOrderResultsBySharedVectors()} — writers + readers via orderResultsBy() with shared vectors
    • + *
+ */ +public class VectorMemoryIndexSimulationTest extends SimulationTestBase +{ + private static final int DIMENSIONS = 8; + private static final int VECTORS_PER_THREAD = 200; + private static final int NUM_WRITER_THREADS = 4; + private static final int NUM_READER_THREADS = 4; + private static final int PRE_SEED_COUNT = 50; + // Lower than VectorMemoryIndexTest (0.9) because simulation uses fewer vectors per thread + // and the adversarial scheduling can affect graph construction quality. + private static final double RECALL_THRESHOLD = 0.8; + + @BeforeClass + public static void registerNemesisFields() + { + // OnHeapGraph mutable fields + NemesisFieldSelectors.register(OnHeapGraph.class, "hasDeletions", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "postingsMap", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "postingsByOrdinal", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "vectorsByKey", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "vectorValues", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "builder", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(OnHeapGraph.class, "nextOrdinal", NemesisFieldKind.ATOMICX); + + // VectorPostings mutable fields + NemesisFieldSelectors.register(VectorPostings.class, "ordinal", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(VectorPostings.class, "postings", NemesisFieldKind.SIMPLE); + NemesisFieldSelectors.register(VectorPostings.class, "rowIds", NemesisFieldKind.SIMPLE); + + // NOTE: VectorMemoryIndex.keyBounds is NOT registered because it is written inside a + // synchronized(this) block (updateKeyBounds). Nemesis pausing inside a held monitor + // causes unresolvable deadlocks with the simulator's cooperative scheduling. + } + + /** + * Verifies that concurrent calls to {@link VectorMemoryIndex#add} with random vectors + * do not corrupt the graph or lose data. + */ + @Test + public void testConcurrentAddsWithRandomVectors() + { + SharedState.useSharedVectors = false; + testConcurrentAddsAreEventuallyConsistent(); + } + + /** + * Verifies that concurrent calls to {@link VectorMemoryIndex#add} with shared (near-duplicate) + * vectors do not corrupt the graph or lose data. + */ + @Test + public void testConcurrentAddsWithSharedVectors() + { + SharedState.useSharedVectors = true; + testConcurrentAddsAreEventuallyConsistent(); + } + + /** + * Verifies that concurrent calls to {@link VectorMemoryIndex#add} do not corrupt the graph + * or lose data. Each thread owns a disjoint range of partition key integers. + * After all writers complete, the graph must contain all inserted vectors. + *

+ * GraphIndexBuilder.addGraphNode() is designed for concurrent use: insertionsInProgress + * is a ConcurrentSkipListSet, and PoolingSupport gives each thread its own GraphSearcher + * and scratch arrays. This test validates the full stack from VectorMemoryIndex.index() + * through OnHeapGraph.add() through GraphIndexBuilder.addGraphNode(). + *

+ * After all writes complete, a full-ring search with limit == totalInserted + * must return the vast majority of distinct results. Every returned key must + * have been inserted by a worker thread, and every score must be a valid + * positive float (a zero or NaN score would indicate graph corruption). + */ + private void testConcurrentAddsAreEventuallyConsistent() + { + int totalInserted = NUM_WRITER_THREADS * VECTORS_PER_THREAD; + + simulate(() -> { + boolean useShared = SharedState.useSharedVectors; + VectorMemoryIndex memtableIndex = createVectorMemoryIndex(); + ConcurrentMap keyMap = new ConcurrentHashMap<>(); + + org.apache.cassandra.concurrent.ExecutorPlus executor = + org.apache.cassandra.concurrent.ExecutorFactory.Global.executorFactory().pooled("writers", NUM_WRITER_THREADS); + + for (int t = 0; t < NUM_WRITER_THREADS; t++) + { + final int threadId = t; + executor.submit(() -> { + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int pk = threadId * VECTORS_PER_THREAD + i; + DecoratedKey key = makeKey(pk); + ByteBuffer vector = useShared ? makeSharedVector(i) : randomVector(); + memtableIndex.add(key, Clustering.EMPTY, vector); + keyMap.put(pk, key); + } + }); + } + + SharedState.memtableIndex = memtableIndex; + SharedState.keyMap = keyMap; + SharedState.totalInserted = totalInserted; + }, () -> { + VectorMemoryIndex memtableIndex = (VectorMemoryIndex) SharedState.memtableIndex; + @SuppressWarnings("unchecked") + ConcurrentMap keyMap = (ConcurrentMap) SharedState.keyMap; + int total = SharedState.totalInserted; + + if (memtableIndex == null) + throw new AssertionError("VectorMemoryIndex was not created"); + if (memtableIndex.isEmpty()) + throw new AssertionError("VectorMemoryIndex is empty after concurrent adds"); + + StorageAttachedIndex index = getIndex(memtableIndex); + AbstractBounds fullRing = + new Range<>(Murmur3Partitioner.instance.getMinimumToken().minKeyBound(), + Murmur3Partitioner.instance.getMinimumToken().minKeyBound()); + + Expression expression = Expression.create(index); + expression.add(Operator.ANN, randomVector()); + + QueryContext queryContext = createQueryContext(total); + Set foundKeys = new HashSet<>(); + try (CloseableIterator results = memtableIndex.orderBy(queryContext, expression, fullRing)) + { + while (results.hasNext()) + { + PrimaryKeyWithScore result = results.next(); + if (result.primaryKey() == null) + throw new AssertionError("Null PrimaryKey in search results after concurrent adds"); + float score = result.score(); + if (!Float.isFinite(score)) + throw new AssertionError("Non-finite score after concurrent adds: " + score); + + // All vector components are drawn from [0, 1) via ThreadLocalRandom.nextFloat(), + // so every term in the dot product is non-negative and the sum is strictly positive. + // A score of 0f or below would indicate graph corruption, not a valid similarity result. + if (score <= 0f) + throw new AssertionError("Non-positive score after concurrent adds: " + score); + + int pk = Int32Type.instance.compose(result.primaryKey().partitionKey().getKey()); + if (foundKeys.contains(pk)) + throw new AssertionError("Duplicate key returned after concurrent adds: " + pk); + if (!keyMap.containsKey(pk)) + throw new AssertionError("Returned key " + pk + " was not inserted by any worker thread"); + foundKeys.add(pk); + } + } + + int expectedMinimum = (int) (total * RECALL_THRESHOLD); + if (foundKeys.size() < expectedMinimum) + throw new AssertionError(String.format( + "Search returned %d of %d results after concurrent adds (expected at least %d)", + foundKeys.size(), total, expectedMinimum)); + }, DEFAULT_ITERATIONS); + } + + /** + * Verifies that {@link VectorMemoryIndex#orderBy} never throws while concurrent add() calls + * with random vectors are in progress. + */ + @Test + public void testConcurrentAddsAndOrderByRandomVectors() + { + SharedState.useSharedVectors = false; + testConcurrentAddsAndOrderByNeverThrow(); + } + + /** + * Verifies that {@link VectorMemoryIndex#orderBy} never throws while concurrent add() calls + * with shared (near-duplicate) vectors are in progress. + */ + @Test + public void testConcurrentAddsAndOrderBySharedVectors() + { + SharedState.useSharedVectors = true; + testConcurrentAddsAndOrderByNeverThrow(); + } + + /** + * Verifies that {@link VectorMemoryIndex#orderBy} never throws while concurrent add() calls + * are in progress, and that the index reaches a consistent state once writes settle. + *

+ * Missing results during concurrent writes are expected and correct — a read that + * races with a write is allowed to miss that write (valid linearization). The only + * invariant asserted during the write window is safety: no exceptions, no null PKs, + * no non-finite scores from results that *are* returned. + *

+ * After all writers complete, a final search verifies full consistency at rest. + */ + private void testConcurrentAddsAndOrderByNeverThrow() + { + int totalInserted = NUM_WRITER_THREADS * VECTORS_PER_THREAD; + + simulate(() -> { + boolean useShared = SharedState.useSharedVectors; + VectorMemoryIndex memtableIndex = createVectorMemoryIndex(); + ConcurrentMap keyMap = new ConcurrentHashMap<>(); + + // Pre-seed enough rows that orderBy() always has a non-empty graph to search, + // avoiding the early-return in OnHeapGraph.search() when vectorValues.size() == 0 + // which would prevent readers from exercising any real code paths. + for (int i = 1; i <= PRE_SEED_COUNT; i++) + { + DecoratedKey dk = makeKey(-i); + memtableIndex.add(dk, Clustering.EMPTY, randomVector()); + keyMap.put(-i, dk); + } + + org.apache.cassandra.concurrent.ExecutorPlus executor = + org.apache.cassandra.concurrent.ExecutorFactory.Global.executorFactory().pooled("mixed", NUM_WRITER_THREADS + NUM_READER_THREADS); + + // Writers: each inserts into a disjoint PK range + for (int t = 0; t < NUM_WRITER_THREADS; t++) + { + final int threadId = t; + executor.submit(() -> { + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int pk = threadId * VECTORS_PER_THREAD + i; + DecoratedKey dk = makeKey(pk); + ByteBuffer vector = useShared ? makeSharedVector(i) : randomVector(); + memtableIndex.add(dk, Clustering.EMPTY, vector); + keyMap.put(pk, dk); + } + }); + } + + // Readers: call real orderBy() while writers are in progress + // Safety assertions only. Missing results are a valid linearization + // of concurrent read/write and are not asserted against here. + StorageAttachedIndex index = getIndex(memtableIndex); + for (int t = 0; t < NUM_READER_THREADS; t++) + { + executor.submit(() -> { + AbstractBounds fullRing = + new Range<>(Murmur3Partitioner.instance.getMinimumToken().minKeyBound(), + Murmur3Partitioner.instance.getMinimumToken().minKeyBound()); + + for (int i = 0; i < 10; i++) + { + Expression expression = Expression.create(index); + expression.add(Operator.ANN, randomVector()); + QueryContext ctx = createQueryContext(totalInserted + PRE_SEED_COUNT); + + // Must not throw — safety invariant + try (CloseableIterator results = memtableIndex.orderBy(ctx, expression, fullRing)) + { + while (results.hasNext()) + { + PrimaryKeyWithScore result = results.next(); + if (result.primaryKey() == null) + throw new AssertionError("Null PrimaryKey during concurrent add() + orderBy()"); + if (!Float.isFinite(result.score())) + throw new AssertionError("Non-finite score during concurrent add() + orderBy(): " + result.score()); + } + } + } + }); + } + + SharedState.memtableIndex = memtableIndex; + SharedState.keyMap = keyMap; + SharedState.totalInserted = totalInserted; + }, () -> { + VectorMemoryIndex memtableIndex = (VectorMemoryIndex) SharedState.memtableIndex; + int total = SharedState.totalInserted; + StorageAttachedIndex index = getIndex(memtableIndex); + + AbstractBounds fullRing = + new Range<>(Murmur3Partitioner.instance.getMinimumToken().minKeyBound(), + Murmur3Partitioner.instance.getMinimumToken().minKeyBound()); + + Expression expression = Expression.create(index); + expression.add(Operator.ANN, randomVector()); + QueryContext ctx = createQueryContext(total + PRE_SEED_COUNT); + + Set foundKeys = new HashSet<>(); + try (CloseableIterator results = memtableIndex.orderBy(ctx, expression, fullRing)) + { + while (results.hasNext()) + { + PrimaryKeyWithScore result = results.next(); + if (result.primaryKey() == null) + throw new AssertionError("Null PrimaryKey after writes settled"); + if (!Float.isFinite(result.score())) + throw new AssertionError("Non-finite score after writes settled: " + result.score()); + + // All vector components are drawn from [0, 1) via ThreadLocalRandom.nextFloat(), + // so every term in the dot product is non-negative and the sum is strictly positive. + // A score of 0f or below would indicate graph corruption, not a valid similarity result. + if (result.score() <= 0f) + throw new AssertionError("Non-positive score after writes settled: " + result.score()); + + int pk = Int32Type.instance.compose(result.primaryKey().partitionKey().getKey()); + foundKeys.add(pk); + } + } + + // ANN recall is approximate, so we allow a small miss rate rather + // than asserting exact equality. Pre-seeded keys (negative PKs) are + // included in the limit so they do not crowd out writer-inserted keys. + long writerKeysFound = foundKeys.stream().filter(pk -> pk >= 0).count(); + int expectedMinimum = (int) (total * RECALL_THRESHOLD); + if (writerKeysFound < expectedMinimum) + throw new AssertionError(String.format( + "Only %d of %d writer-inserted keys found after writes settled (expected at least %d)", + writerKeysFound, total, expectedMinimum)); + }, DEFAULT_ITERATIONS); + } + + /** + * Verifies that {@link VectorMemoryIndex#orderResultsBy} never throws while concurrent + * add() calls with random vectors are in progress. + */ + @Test + public void testConcurrentAddsAndOrderResultsByRandomVectors() + { + SharedState.useSharedVectors = false; + testConcurrentAddsAndOrderResultsByNeverThrow(); + } + + /** + * Verifies that {@link VectorMemoryIndex#orderResultsBy} never throws while concurrent + * add() calls with shared (near-duplicate) vectors are in progress. + */ + @Test + public void testConcurrentAddsAndOrderResultsBySharedVectors() + { + SharedState.useSharedVectors = true; + testConcurrentAddsAndOrderResultsByNeverThrow(); + } + + /** + * Verifies that orderResultsBy() never throws while concurrent add() calls are in + * progress, and that the index reaches a consistent state once writes settle. + *

+ * The materialized key list passed to orderResultsBy() is built from keyMap, which is a + * ConcurrentHashMap updated by every add() call. A snapshot taken mid-write may be + * incomplete — this is intentional and mirrors the production path where the source + * KeyRangeIterator only sees keys committed before the non-ANN index scan ran. + */ + private void testConcurrentAddsAndOrderResultsByNeverThrow() + { + int totalInserted = NUM_WRITER_THREADS * VECTORS_PER_THREAD; + + simulate(() -> { + boolean useShared = SharedState.useSharedVectors; + VectorMemoryIndex memtableIndex = createVectorMemoryIndex(); + StorageAttachedIndex index = getIndex(memtableIndex); + ConcurrentMap keyMap = new ConcurrentHashMap<>(); + + // Pre-seed rows so orderResultsBy() always has a non-empty [minimumKey, maximumKey] + // window and a non-trivial resultsInRange list on the first reader pass. + for (int i = 1; i <= PRE_SEED_COUNT; i++) + { + DecoratedKey dk = makeKey(-i); + memtableIndex.add(dk, Clustering.EMPTY, randomVector()); + keyMap.put(-i, dk); + } + + org.apache.cassandra.concurrent.ExecutorPlus executor = + org.apache.cassandra.concurrent.ExecutorFactory.Global.executorFactory().pooled("mixed", NUM_WRITER_THREADS + NUM_READER_THREADS); + + // Writers: each inserts into a disjoint PK range [threadId*VECTORS_PER_THREAD, (threadId+1)*VECTORS_PER_THREAD) + for (int t = 0; t < NUM_WRITER_THREADS; t++) + { + final int threadId = t; + executor.submit(() -> { + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int pk = threadId * VECTORS_PER_THREAD + i; + DecoratedKey dk = makeKey(pk); + ByteBuffer vector = useShared ? makeSharedVector(i) : randomVector(); + memtableIndex.add(dk, Clustering.EMPTY, vector); + keyMap.put(pk, dk); + } + }); + } + + // Readers: call real orderResultsBy() with a snapshot of current keys + // Safety assertions only during concurrent writes. Missing results are a valid + // linearization and are not asserted against here. + for (int t = 0; t < NUM_READER_THREADS; t++) + { + executor.submit(() -> { + for (int i = 0; i < 10; i++) + { + // Snapshot current keys and build sorted PrimaryKey list + List snapshotKeys = keyMap.values() + .stream() + .map(dk -> index.keyFactory().create(dk)) + .sorted() + .collect(Collectors.toList()); + if (snapshotKeys.isEmpty()) + continue; + + Expression expression = Expression.create(index); + expression.add(Operator.ANN, randomVector()); + QueryContext ctx = createQueryContext(snapshotKeys.size()); + + // Must not throw + try (CloseableIterator results = memtableIndex.orderResultsBy(ctx, snapshotKeys, expression)) + { + while (results.hasNext()) + { + PrimaryKeyWithScore result = results.next(); + if (result.primaryKey() == null) + throw new AssertionError("Null PrimaryKey during concurrent add() + orderResultsBy()"); + if (!Float.isFinite(result.score())) + throw new AssertionError("Non-finite score during concurrent add() + orderResultsBy(): " + result.score()); + } + } + } + }); + } + + SharedState.memtableIndex = memtableIndex; + SharedState.keyMap = keyMap; + SharedState.totalInserted = totalInserted; + }, () -> { + VectorMemoryIndex memtableIndex = (VectorMemoryIndex) SharedState.memtableIndex; + StorageAttachedIndex index = getIndex(memtableIndex); + @SuppressWarnings("unchecked") + ConcurrentMap keyMap = (ConcurrentMap) SharedState.keyMap; + int total = SharedState.totalInserted; + + List allKeys = keyMap.values() + .stream() + .map(dk -> index.keyFactory().create(dk)) + .sorted() + .collect(Collectors.toList()); + + Expression expression = Expression.create(index); + expression.add(Operator.ANN, randomVector()); + QueryContext ctx = createQueryContext(total + PRE_SEED_COUNT); + + Set foundKeys = new HashSet<>(); + try (CloseableIterator results = memtableIndex.orderResultsBy(ctx, allKeys, expression)) + { + while (results.hasNext()) + { + PrimaryKeyWithScore result = results.next(); + if (result.primaryKey() == null) + throw new AssertionError("Null PrimaryKey after writes settled in orderResultsBy()"); + if (!Float.isFinite(result.score())) + throw new AssertionError("Non-finite score after writes settled in orderResultsBy(): " + result.score()); + + // All vector components are drawn from [0, 1) via ThreadLocalRandom.nextFloat(), + // so every term in the dot product is non-negative and the sum is strictly positive. + // A score of 0f or below would indicate graph corruption, not a valid similarity result. + if (result.score() <= 0f) + throw new AssertionError("Non-positive score after writes settled in orderResultsBy(): " + result.score()); + + int pk = Int32Type.instance.compose(result.primaryKey().partitionKey().getKey()); + foundKeys.add(pk); + } + } + + long writerKeysFound = foundKeys.stream().filter(pk -> pk >= 0).count(); + int expectedMinimum = (int) (total * RECALL_THRESHOLD); + if (writerKeysFound < expectedMinimum) + throw new AssertionError(String.format( + "orderResultsBy() returned %d of %d writer-inserted keys after writes settled (expected at least %d)", + writerKeysFound, total, expectedMinimum)); + }, DEFAULT_ITERATIONS); + } + + // ---- Infrastructure: create VectorMemoryIndex without DatabaseDescriptor ---- + + /** + * Creates a {@link VectorMemoryIndex} by constructing a minimal {@link StorageAttachedIndex} + * via {@code Unsafe.allocateInstance()} (bypassing the constructor which requires ColumnFamilyStore + * and triggers DatabaseDescriptor initialization). + *

+ * Only the fields actually accessed by VectorMemoryIndex methods are set: + *

    + *
  • {@code indexTermType} — for decomposeVector(), indexType(), columnMetadata()
    • + *
  • {@code indexWriterConfig} — for graph construction and getSimilarityFunction()
    • + *
  • {@code primaryKeyFactory} — for creating PrimaryKey instances
    • + *
+ */ + private static VectorMemoryIndex createVectorMemoryIndex() + { + try + { + // RangeUtil has a static initializer that calls DatabaseDescriptor.getPartitioner(). + // Set the partitioner field directly to avoid full DD initialization. + setField(org.apache.cassandra.config.DatabaseDescriptor.class, "partitioner", null, Murmur3Partitioner.instance); + + VectorType vectorType = VectorType.getInstance(FloatType.instance, DIMENSIONS); + ColumnMetadata column = ColumnMetadata.regularColumn("ks_sim", "tbl_sim", "vec", vectorType, 0); + IndexTermType termType = IndexTermType.create(column, Collections.singletonList( + ColumnMetadata.regularColumn("ks_sim", "tbl_sim", "pk", Int32Type.instance, 1)), + IndexTarget.Type.SIMPLE); + + IndexWriterConfig writerConfig = new IndexWriterConfig( + IndexWriterConfig.DEFAULT_MAXIMUM_NODE_CONNECTIONS, + IndexWriterConfig.DEFAULT_CONSTRUCTION_BEAM_WIDTH, + VectorSimilarityFunction.DOT_PRODUCT, + null); + + ClusteringComparator emptyComparator = new ClusteringComparator(); + PrimaryKey.Factory keyFactory = new PrimaryKey.Factory(Murmur3Partitioner.instance, emptyComparator); + + // Allocate StorageAttachedIndex without calling its constructor + Unsafe unsafe = getUnsafe(); + StorageAttachedIndex index = (StorageAttachedIndex) unsafe.allocateInstance(StorageAttachedIndex.class); + + // Set the fields VectorMemoryIndex accesses + setField(StorageAttachedIndex.class, "indexTermType", index, termType); + setField(StorageAttachedIndex.class, "indexWriterConfig", index, writerConfig); + setField(StorageAttachedIndex.class, "primaryKeyFactory", index, keyFactory); + + // hasClustering() calls baseCfs.getComparator().size() — needs a minimal CFS with metadata + TableMetadata tableMetadata = TableMetadata.builder("ks_sim", "tbl_sim") + .addPartitionKeyColumn("pk", Int32Type.instance) + .addRegularColumn("val", vectorType) + .partitioner(Murmur3Partitioner.instance) + .caching(CachingParams.CACHE_NOTHING) + .build(); + // Create a minimal CFS via Unsafe with just the metadata ref set + ColumnFamilyStore fakeCfs = (ColumnFamilyStore) unsafe.allocateInstance(ColumnFamilyStore.class); + org.apache.cassandra.schema.TableMetadataRef metadataRef = + org.apache.cassandra.schema.TableMetadataRef.forOfflineTools(tableMetadata); + setField(ColumnFamilyStore.class, "metadata", fakeCfs, metadataRef); + setField(StorageAttachedIndex.class, "baseCfs", index, fakeCfs); + + // validateTermSize needs maxTermSizeGuardrail (non-null). + // Create a disabled guardrail via Unsafe rather than loading Guardrails (which may trigger DD init) + org.apache.cassandra.db.guardrails.MaxThreshold fakeGuardrail = + (org.apache.cassandra.db.guardrails.MaxThreshold) unsafe.allocateInstance(org.apache.cassandra.db.guardrails.MaxThreshold.class); + setField(StorageAttachedIndex.class, "maxTermSizeGuardrail", index, fakeGuardrail); + + // Proxy Memtable — Mockito cannot be used inside InstanceClassLoader + Memtable memtable = (Memtable) java.lang.reflect.Proxy.newProxyInstance( + Memtable.class.getClassLoader(), + new Class[]{ Memtable.class }, + (proxy, method, args) -> { + 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]; + return null; + } + ); + + return new VectorMemoryIndex(index, memtable); + } + catch (Exception e) + { + throw new RuntimeException("Failed to create VectorMemoryIndex for simulation", e); + } + } + + private static DecoratedKey makeKey(int pk) + { + ByteBuffer key = Int32Type.instance.decompose(pk); + return Murmur3Partitioner.instance.decorateKey(key); + } + + private static ByteBuffer randomVector() + { + List rawVector = new ArrayList<>(DIMENSIONS); + for (int i = 0; i < DIMENSIONS; i++) + rawVector.add(ThreadLocalRandom.current().nextFloat()); + return VectorType.getInstance(FloatType.instance, DIMENSIONS).getSerializer().serialize(rawVector); + } + + /** + * Creates a shared vector where most dimensions are 0.5f and the last dimension + * varies by index. This tests that the graph handles duplicate/near-duplicate vectors correctly. + */ + private static ByteBuffer makeSharedVector(int i) + { + List raw = new ArrayList<>(Collections.nCopies(DIMENSIONS - 1, 0.5f)); + raw.add(i / (float) VECTORS_PER_THREAD); + return VectorType.getInstance(FloatType.instance, DIMENSIONS).getSerializer().serialize(raw); + } + + private static QueryContext createQueryContext(int limit) + { + // Create a minimal ReadCommand for QueryContext's limit() method. + // TableMetadata and PartitionRangeReadCommand do not require DatabaseDescriptor. + VectorType vectorType = VectorType.getInstance(FloatType.instance, DIMENSIONS); + TableMetadata metadata = TableMetadata.builder("ks_sim", "tbl_sim") + .addPartitionKeyColumn("pk", Int32Type.instance) + .addRegularColumn("val", vectorType) + .partitioner(Murmur3Partitioner.instance) + .caching(CachingParams.CACHE_NOTHING) + .build(); + return new QueryContext( + PartitionRangeReadCommand.create(metadata, + (int) (System.currentTimeMillis() / 1000), + ColumnFilter.all(metadata), + RowFilter.none(), + DataLimits.cqlLimits(limit), + DataRange.allData(metadata.partitioner)), + TimeUnit.SECONDS.toMillis(60)); + } + + private static StorageAttachedIndex getIndex(VectorMemoryIndex memtableIndex) + { + try + { + Field f = org.apache.cassandra.index.sai.memory.MemoryIndex.class.getDeclaredField("index"); + f.setAccessible(true); + return (StorageAttachedIndex) f.get(memtableIndex); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + } + + private static Unsafe getUnsafe() + { + try + { + Field f = Unsafe.class.getDeclaredField("theUnsafe"); + f.setAccessible(true); + return (Unsafe) f.get(null); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + } + + private static void setField(Class clazz, String fieldName, Object target, Object value) + { + try + { + Field f = clazz.getDeclaredField(fieldName); + f.setAccessible(true); + f.set(target, value); + } + catch (Exception e) + { + throw new RuntimeException("Failed to set field " + fieldName + " on " + clazz.getSimpleName(), e); + } + } + + /** + * Static holder for sharing state between the action and check runnables. + * Fields are volatile because they are written by the test method (outer classloader) + * and read inside the InstanceClassLoader. + */ + public static class SharedState + { + public static volatile Object memtableIndex; + public static volatile Object keyMap; + public static volatile int totalInserted; + public static volatile boolean useSharedVectors; + } +} diff --git a/test/unit/org/apache/cassandra/index/sai/memory/VectorMemoryIndexTest.java b/test/unit/org/apache/cassandra/index/sai/memory/VectorMemoryIndexTest.java index 1c87978a63..7f89e46d1b 100644 --- a/test/unit/org/apache/cassandra/index/sai/memory/VectorMemoryIndexTest.java +++ b/test/unit/org/apache/cassandra/index/sai/memory/VectorMemoryIndexTest.java @@ -20,13 +20,27 @@ package org.apache.cassandra.index.sai.memory; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.TreeMap; +import java.util.concurrent.BrokenBarrierException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiFunction; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.junit.Before; @@ -64,6 +78,7 @@ 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.PrimaryKeyWithScore; import org.apache.cassandra.index.sai.plan.Expression; +import org.apache.cassandra.index.sai.utils.PrimaryKey; import org.apache.cassandra.inject.Injections; import org.apache.cassandra.inject.InvokePointBuilder; import org.apache.cassandra.schema.TableMetadata; @@ -74,8 +89,17 @@ import org.apache.cassandra.utils.FBUtilities; import static org.apache.cassandra.config.CassandraRelevantProperties.MEMTABLE_SHARD_COUNT; import static org.apache.cassandra.config.CassandraRelevantProperties.ORG_APACHE_CASSANDRA_DISABLE_MBEAN_REGISTRATION; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +/** + * Randomized tests around the functionality of {@link VectorMemoryIndex}. + *

+ * Note that the multithreaded tests here have been ported to use the simulator in + * {@link org.apache.cassandra.simulator.test.VectorMemoryIndexSimulationTest}. Although they should catch many of + * the same kinds of bugs, this test persists partially as an example of how this porting can be done. + */ public class VectorMemoryIndexTest extends SAITester { private static final Injections.Counter indexSearchCounter = Injections.newCounter("IndexSearchCounter") @@ -84,11 +108,14 @@ public class VectorMemoryIndexTest extends SAITester .onMethod("search")) .build(); + private static final double RECALL_THRESHOLD = 0.9; + private static final int VECTORS_PER_THREAD = 2000; + private ColumnFamilyStore cfs; private StorageAttachedIndex index; private VectorMemoryIndex memtableIndex; private IPartitioner partitioner; - private Map keyMap; + private ConcurrentMap keyMap; private Map rowMap; private int dimensionCount; @@ -119,16 +146,12 @@ public class VectorMemoryIndexTest extends SAITester cfs = index.baseCfs(); partitioner = cfs.getPartitioner(); indexSearchCounter.reset(); - keyMap = new TreeMap<>(); + keyMap = new ConcurrentHashMap<>(); rowMap = new HashMap<>(); Injections.inject(indexSearchCounter); } - public static void reassignLocalTokens() - { - } - @Test public void randomQueryTest() throws Exception { @@ -149,7 +172,6 @@ public class VectorMemoryIndexTest extends SAITester List keys = new ArrayList<>(keyMap.keySet()); long actualVectorsReturned = 0; long expectedVectorsReturned = 0; - double expectedRecall = 0.9; for (int executionCount = 0; executionCount < 1000; executionCount++) { @@ -161,14 +183,7 @@ public class VectorMemoryIndexTest extends SAITester Set foundKeys = new HashSet<>(); int limit = getRandom().nextIntBetween(1, 100); - - ReadCommand command = PartitionRangeReadCommand.create(cfs.metadata(), - FBUtilities.nowInSeconds(), - ColumnFilter.all(cfs.metadata()), - RowFilter.none(), - DataLimits.cqlLimits(limit), - DataRange.allData(cfs.metadata().partitioner)); - + ReadCommand command = createRangeRead(limit); long expectedResults = Math.min(limit, keysInRange.size()); try (CloseableIterator iterator = memtableIndex.orderBy(new QueryContext(command, @@ -198,18 +213,504 @@ public class VectorMemoryIndexTest extends SAITester expectedVectorsReturned += expectedResults; if (foundKeys.size() < expectedResults) assertTrue("Expected at least " + expectedResults + " results but got " + foundKeys.size(), - foundKeys.size() >= expectedResults * expectedRecall); + foundKeys.size() >= expectedResults * RECALL_THRESHOLD); } } assertTrue("Expected at least " + expectedVectorsReturned + " results but got " + actualVectorsReturned, - actualVectorsReturned >= expectedVectorsReturned * expectedRecall); + actualVectorsReturned >= expectedVectorsReturned * RECALL_THRESHOLD); + } + + /** + * Verifies that expectedNodesVisited() always returns a value within its documented + * bounds: at least min(limit, graphSize) and at most graphSize. + *

+ * This is a pure arithmetic test with no index infrastructure required. It exercises + * the boundary conditions that matter for the brute-force/ANN threshold decision in + * maxBruteForceRows(): if the formula underflows its lower bound, small queries will + * incorrectly use ANN; if it overflows its upper bound, the result is nonsensical. + */ + @Test + public void testExpectedNodesVisitedRespectsBounds() + { + int[] graphSizes = { 1, 2, 10, 100, 1000, 10000 }; + int[] limits = { 1, 2, 5, 10, 50, 100 }; + double[] permittedFractions = { 0.01, 0.1, 0.5, 1.0, 2.0 }; + + for (int graphSize : graphSizes) + { + for (int limit : limits) + { + for (double fraction : permittedFractions) + { + int permitted = Math.max(1, (int) (graphSize * fraction)); + int result = VectorMemoryIndex.expectedNodesVisited(limit, permitted, graphSize); + int lowerBound = Math.min(limit, graphSize); + + assertTrue(String.format("expectedNodesVisited(%d, %d, %d) = %d is below lower bound %d", limit, permitted, graphSize, result, lowerBound), + result >= lowerBound); + + assertTrue(String.format("expectedNodesVisited(%d, %d, %d) = %d exceeds graphSize %d", limit, permitted, graphSize, result, graphSize), + result <= graphSize); + } + } + } } @Test - public void indexIteratorTest() + public void testConcurrentAddsWithRandomVectors() throws Exception { - // VSTODO + testConcurrentAddsAreEventuallyConsistent((threadId, i) -> randomVectorFromThreadLocal()); + } + + @Test + public void testConcurrentAddsWithSharedVectors() throws Exception + { + testConcurrentAddsAreEventuallyConsistent((threadId, i) -> makeSharedVector(i)); + } + + /** + * Verifies that concurrent calls to add() do not corrupt the graph or lose data. + *

+ * GraphIndexBuilder.addGraphNode() is designed for concurrent use: insertionsInProgress + * is a ConcurrentSkipListSet, and PoolingSupport gives each thread its own GraphSearcher + * and scratch arrays. This test validates the full stack from VectorMemoryIndex.index() + * through OnHeapGraph.add() through GraphIndexBuilder.addGraphNode(). + *

+ * 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 vectorFactory) throws Exception + { + Memtable memtable = Mockito.mock(Memtable.class); + memtableIndex = new VectorMemoryIndex(index, memtable); + + int numThreads = Runtime.getRuntime().availableProcessors(); + int totalInserted = numThreads * VECTORS_PER_THREAD; + + ExecutorService executor = Executors.newFixedThreadPool(numThreads); + + // CyclicBarrier ensures all threads begin inserting simultaneously, + // maximizing contention on GraphIndexBuilder and ConcurrentVectorValues. + CyclicBarrier barrier = new CyclicBarrier(numThreads); + List> futures = new ArrayList<>(); + + for (int t = 0; t < numThreads; t++) + { + final int threadId = t; + futures.add(executor.submit(() -> { + try + { + barrier.await(); + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int pk = threadId * VECTORS_PER_THREAD + i; + addRow(pk, vectorFactory.apply(threadId, i)); + } + } + catch (BrokenBarrierException | InterruptedException e) + { + throw new RuntimeException(e); + } + })); + } + + executor.shutdown(); + assertTrue("Timed out waiting for concurrent adds", executor.awaitTermination(60, TimeUnit.SECONDS)); + + // Rethrow any exception from worker threads — assertion failures inside a + // Runnable are otherwise silently swallowed by the ExecutorService. + for (Future f : futures) + { + try + { + f.get(); + } + catch (ExecutionException e) + { + fail("Worker thread threw during concurrent add(): " + e.getCause()); + } + } + + // After all writes complete, a full-ring search with limit == totalInserted + // must return the vast majority of distinct results. Every returned key must + // have been inserted by a worker thread, and every score must be a valid + // positive float (a zero or NaN score would indicate graph corruption). + AbstractBounds fullRing = new Range<>(partitioner.getMinimumToken().minKeyBound(), partitioner.getMinimumToken().minKeyBound()); + Expression expression = generateRandomExpression(); + ReadCommand command = createRangeRead(totalInserted); + + QueryContext queryContext = new QueryContext(command, DatabaseDescriptor.getRangeRpcTimeout(TimeUnit.MILLISECONDS)); + Set foundKeys = new HashSet<>(); + + try (CloseableIterator iterator = memtableIndex.orderBy(queryContext, expression, fullRing)) + { + while (iterator.hasNext()) + { + PrimaryKeyWithScore result = iterator.next(); + assertNotNull("Null PrimaryKey in search results after concurrent adds", result.primaryKey()); + float score = result.score(); + assertTrue("Non-finite score after concurrent adds: " + score, Float.isFinite(score)); + + // All vector components are drawn from [0, 1) via ThreadLocalRandom.nextFloat(), + // so every term in the dot product is non-negative and the sum is strictly positive. + // A score of 0f or below would indicate graph corruption, not a valid similarity result. + assertTrue("Non-positive score after concurrent adds: " + score, score > 0f); + + int pk = Int32Type.instance.compose(result.primaryKey().partitionKey().getKey()); + assertFalse("Duplicate key returned after concurrent adds: " + pk, foundKeys.contains(pk)); + assertTrue("Returned key " + pk + " was not inserted by any worker thread", keyMap.containsKey(result.primaryKey().partitionKey())); + foundKeys.add(pk); + } + } + + assertTrue("Search returned " + foundKeys.size() + " of " + totalInserted + " results after concurrent adds (expected at least " + (int) (totalInserted * RECALL_THRESHOLD) + ')', + foundKeys.size() >= totalInserted * RECALL_THRESHOLD); + } + + @Test + public void testConcurrentAddsAndOrderByRandomVectors() throws Exception + { + testConcurrentAddsAndOrderByNeverThrow((threadId, i) -> randomVectorFromThreadLocal()); + } + + @Test + public void testConcurrentAddsAndOrderBySharedVectors() throws Exception + { + testConcurrentAddsAndOrderByNeverThrow((threadId, i) -> makeSharedVector(i)); + } + + /** + * Verifies that orderBy() never throws an exception while concurrent add() calls + * are in progress, and that the index reaches a consistent state once writes settle. + *

+ * Missing results during concurrent writes are expected and correct — a read that + * races with a write is allowed to miss that write (valid linearization). The only + * invariant asserted during the write window is safety: no exceptions, no null PKs, + * no non-finite scores from results that *are* returned. + *

+ * Readers block on writersFinished after the barrier release and perform one final + * search pass after all writers have joined, confirming full consistency at rest. + */ + public void testConcurrentAddsAndOrderByNeverThrow(BiFunction vectorFactory) throws Exception + { + Memtable memtable = Mockito.mock(Memtable.class); + memtableIndex = new VectorMemoryIndex(index, memtable); + + int numWriterThreads = Runtime.getRuntime().availableProcessors(); + int numReaderThreads = Runtime.getRuntime().availableProcessors(); + int totalInserted = numWriterThreads * VECTORS_PER_THREAD; + + // Pre-seed enough rows that orderBy() always has a non-empty graph to search, + // avoiding the early-return in OnHeapGraph.search() when vectorValues.size() == 0 + // which would prevent readers from exercising any real code paths. + int preSeedCount = 50; + for (int i = 1; i <= preSeedCount; i++) + addRow(-i, randomVector()); // negative PKs, disjoint from writer range [0, totalInserted) + + // each reader blocks here until every writer has completed, + // then performs one final search to verify post-settlement consistency. + CountDownLatch writersFinished = new CountDownLatch(numWriterThreads); + + // phase1Executed: confirms at least one reader searched during the concurrent + // write window. If this latch is never counted down, the write window was too + // short and the concurrent safety assertions in Phase 1 were never exercised. + AtomicBoolean phase1Executed = new AtomicBoolean(false); + + CopyOnWriteArrayList errors = new CopyOnWriteArrayList<>(); + ExecutorService executor = Executors.newFixedThreadPool(numWriterThreads + numReaderThreads); + CyclicBarrier barrier = new CyclicBarrier(numWriterThreads + numReaderThreads); + + // Writers: each inserts into a disjoint PK range + for (int t = 0; t < numWriterThreads; t++) + { + final int threadId = t; + executor.submit(() -> { + try + { + barrier.await(); + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int pk = threadId * VECTORS_PER_THREAD + i; + addRow(pk, vectorFactory.apply(threadId, i)); + } + } + catch (Throwable e) + { + errors.add(e); + } + finally + { + writersFinished.countDown(); + } + }); + } + + // Readers: issue one search while writers are running (safety only), then + // block on writersFinished and issue one final search for correctness. + for (int t = 0; t < numReaderThreads; t++) + { + executor.submit(() -> { + try + { + barrier.await(); + + AbstractBounds fullRing = new Range<>(partitioner.getMinimumToken().minKeyBound(), partitioner.getMinimumToken().minKeyBound()); + ReadCommand command = createRangeRead(totalInserted + preSeedCount); + QueryContext queryContext = new QueryContext(command, DatabaseDescriptor.getRangeRpcTimeout(TimeUnit.MILLISECONDS)); + + // Build query vectors inline — getRandom() is not thread-safe from + // worker threads, so we use ThreadLocalRandom directly. + ByteBuffer queryBuf = randomVectorFromThreadLocal(); + Expression concurrentExpression = Expression.create(index); + concurrentExpression.add(Operator.ANN, queryBuf); + + // --- Phase 1: one search while writers are still running --- + // Safety assertions only. Missing results are a valid linearization + // of concurrent read/write and are not asserted against here. + try (CloseableIterator it = memtableIndex.orderBy(queryContext, concurrentExpression, fullRing)) + { + while (it.hasNext()) + { + PrimaryKeyWithScore result = it.next(); + assertNotNull("Null PrimaryKey during concurrent add() + orderBy()", result.primaryKey()); + assertTrue("Non-finite score during concurrent add() + orderBy(): " + result.score(), Float.isFinite(result.score())); + } + } + phase1Executed.set(true); + + // --- Phase 2: block until all writers finish, then verify consistency --- + writersFinished.await(); + + ByteBuffer settledQueryBuf = randomVectorFromThreadLocal(); + Expression settledExpression = Expression.create(index); + settledExpression.add(Operator.ANN, settledQueryBuf); + + Set foundAfterSettle = new HashSet<>(); + try (CloseableIterator it = memtableIndex.orderBy(queryContext, settledExpression, fullRing)) + { + while (it.hasNext()) + { + PrimaryKeyWithScore result = it.next(); + assertNotNull("Null PrimaryKey after writes settled", result.primaryKey()); + assertTrue("Non-finite score after writes settled: " + result.score(), Float.isFinite(result.score())); + + // All vector components are drawn from [0, 1) via ThreadLocalRandom.nextFloat(), + // so every term in the dot product is non-negative and the sum is strictly positive. + // A score of 0f or below would indicate graph corruption, not a valid similarity result. + assertTrue("Non-positive score after writes settled: " + result.score(), result.score() > 0f); + + int pk = Int32Type.instance.compose(result.primaryKey().partitionKey().getKey()); + foundAfterSettle.add(pk); + } + } + + // ANN recall is approximate, so we allow a small miss rate rather + // than asserting exact equality. Pre-seeded keys (negative PKs) are + // included in the limit so they do not crowd out writer-inserted keys. + int expectedMinimum = (int) (totalInserted * RECALL_THRESHOLD); + long writerKeysFound = foundAfterSettle.stream().filter(pk -> pk >= 0).count(); + assertTrue("Only " + writerKeysFound + " of " + totalInserted + " writer-inserted keys found after writes settled" + " (expected at least " + expectedMinimum + ')', + writerKeysFound >= expectedMinimum); + } + catch (Throwable e) + { + errors.add(e); + } + }); + } + + executor.shutdown(); + assertTrue("Timed out waiting for concurrent add() + orderBy()", executor.awaitTermination(60, TimeUnit.SECONDS)); + + // Verify Phase 1 actually executed — if this fails, increase vectorsPerWriter + // so the write window is wide enough for readers to search concurrently. + assertTrue("No reader executed a search during the concurrent write window; increase vectorsPerWriter to widen the write window", phase1Executed.get()); + + if (!errors.isEmpty()) + { + AssertionError failure = new AssertionError("Concurrent add() + orderBy() produced " + errors.size() + " error(s); first: " + errors.get(0)); + errors.forEach(failure::addSuppressed); + throw failure; + } + } + + @Test + public void testConcurrentAddsAndOrderResultsByRandomVectors() throws Exception + { + testConcurrentAddsAndOrderResultsByNeverThrow((threadId, i) -> randomVectorFromThreadLocal()); + } + + @Test + public void testConcurrentAddsAndOrderResultsBySharedVectors() throws Exception + { + testConcurrentAddsAndOrderResultsByNeverThrow((threadId, i) -> makeSharedVector(i)); + } + + /** + * Verifies that orderResultsBy() never throws while concurrent add() calls are in + * progress, and that the index reaches a consistent state once writes settle. + *

+ * The materialized key list passed to orderResultsBy() is built from keyMap, which is a + * ConcurrentHashMap updated by every addRow() call. A snapshot taken mid-write may be + * incomplete — this is intentional and mirrors the production path where the source + * KeyRangeIterator only sees keys committed before the non-ANN index scan ran. + */ + private void testConcurrentAddsAndOrderResultsByNeverThrow(BiFunction vectorFactory) throws Exception + { + Memtable memtable = Mockito.mock(Memtable.class); + memtableIndex = new VectorMemoryIndex(index, memtable); + + int numWriterThreads = Runtime.getRuntime().availableProcessors(); + int numReaderThreads = Runtime.getRuntime().availableProcessors(); + int totalInserted = numWriterThreads * VECTORS_PER_THREAD; + + // Pre-seed rows so orderResultsBy() always has a non-empty [minimumKey, maximumKey] + // window and a non-trivial resultsInRange list on the first reader pass. + int preSeedCount = 50; + for (int i = 1; i <= preSeedCount; i++) + addRow(-i, randomVector()); // negative PKs, disjoint from writer range [0, totalInserted) + + CountDownLatch writersFinished = new CountDownLatch(numWriterThreads); + AtomicBoolean phase1Executed = new AtomicBoolean(false); + CopyOnWriteArrayList errors = new CopyOnWriteArrayList<>(); + + ExecutorService executor = Executors.newFixedThreadPool(numWriterThreads + numReaderThreads); + CyclicBarrier barrier = new CyclicBarrier(numWriterThreads + numReaderThreads); + + // Writers: each inserts into a disjoint PK range [threadId*vectorsPerWriter, (threadId+1)*vectorsPerWriter) + for (int t = 0; t < numWriterThreads; t++) + { + final int threadId = t; + executor.submit(() -> { + try + { + barrier.await(); + for (int i = 0; i < VECTORS_PER_THREAD; i++) + { + int pk = threadId * VECTORS_PER_THREAD + i; + addRow(pk, vectorFactory.apply(threadId, i)); + } + } + catch (Throwable e) + { + errors.add(e); + } + finally + { + writersFinished.countDown(); + } + }); + } + + for (int t = 0; t < numReaderThreads; t++) + { + executor.submit(() -> { + try + { + barrier.await(); + + ReadCommand command = createRangeRead(totalInserted + preSeedCount); + QueryContext queryContext = new QueryContext(command, DatabaseDescriptor.getRangeRpcTimeout(TimeUnit.MILLISECONDS)); + + // --- Phase 1: search during concurrent writes --- + // Snapshot the keys visible so far; the list may be incomplete, which is + // a valid linearization. We only assert safety here, not completeness. + List snapshotKeys = buildSortedPrimaryKeySnapshot(); + + ByteBuffer queryBuf = randomVectorFromThreadLocal(); + Expression concurrentExpression = Expression.create(index); + concurrentExpression.add(Operator.ANN, queryBuf); + + if (!snapshotKeys.isEmpty()) + { + try (CloseableIterator it = memtableIndex.orderResultsBy(queryContext, snapshotKeys, concurrentExpression)) + { + while (it.hasNext()) + { + PrimaryKeyWithScore result = it.next(); + assertNotNull("Null PrimaryKey during concurrent add() + orderResultsBy()", result.primaryKey()); + assertTrue("Non-finite score during concurrent add() + orderResultsBy(): " + result.score(), Float.isFinite(result.score())); + } + } + phase1Executed.set(true); + } + + // --- Phase 2: wait for all writers, then verify correctness --- + writersFinished.await(); + + List allKeys = buildSortedPrimaryKeySnapshot(); + ByteBuffer settledQueryBuf = randomVectorFromThreadLocal(); + Expression settledExpression = Expression.create(index); + settledExpression.add(Operator.ANN, settledQueryBuf); + + Set foundAfterSettle = new HashSet<>(); + try (CloseableIterator it = memtableIndex.orderResultsBy(queryContext, allKeys, settledExpression)) + { + while (it.hasNext()) + { + PrimaryKeyWithScore result = it.next(); + assertNotNull("Null PrimaryKey after writes settled in orderResultsBy()", result.primaryKey()); + assertTrue("Non-finite score after writes settled in orderResultsBy(): " + result.score(), Float.isFinite(result.score())); + + // All vector components are drawn from [0, 1) via ThreadLocalRandom.nextFloat(), + // so every term in the dot product is non-negative and the sum is strictly positive. + // A score of 0f or below would indicate graph corruption, not a valid similarity result. + assertTrue("Non-positive score after writes settled in orderResultsBy(): " + result.score(), result.score() > 0f); + + int pk = Int32Type.instance.compose(result.primaryKey().partitionKey().getKey()); + foundAfterSettle.add(pk); + } + } + + long writerKeysFound = foundAfterSettle.stream().filter(pk -> pk >= 0).count(); + int expectedMinimum = (int) (totalInserted * RECALL_THRESHOLD); + assertTrue("orderResultsBy() returned " + writerKeysFound + " of " + totalInserted + " writer-inserted keys after writes settled (expected at least " + expectedMinimum + ')', + writerKeysFound >= expectedMinimum); + } + catch (Throwable e) + { + errors.add(e); + } + }); + } + + executor.shutdown(); + assertTrue("Timed out waiting for concurrent add() + orderResultsBy()", executor.awaitTermination(60, TimeUnit.SECONDS)); + assertTrue("No reader executed a Phase 1 search during the concurrent write window; increase vectorsPerWriter to widen the write window", phase1Executed.get()); + + if (!errors.isEmpty()) + { + AssertionError failure = new AssertionError("Concurrent add() + orderResultsBy() produced " + errors.size() + " error(s); first: " + errors.get(0)); + errors.forEach(failure::addSuppressed); + throw failure; + } + } + + private PartitionRangeReadCommand createRangeRead(int limit) + { + return PartitionRangeReadCommand.create(cfs.metadata(), + FBUtilities.nowInSeconds(), + ColumnFilter.all(cfs.metadata()), + RowFilter.none(), + DataLimits.cqlLimits(limit), + DataRange.allData(cfs.metadata().partitioner)); + } + + private List buildSortedPrimaryKeySnapshot() + { + return keyMap.keySet() + .stream() + .map(dk -> index.hasClustering() ? index.keyFactory().create(dk, Clustering.EMPTY) : index.keyFactory().create(dk)) + .sorted() + .collect(Collectors.toList()); + } + + private ByteBuffer makeSharedVector(int i) + { + List raw = new ArrayList<>(Collections.nCopies(dimensionCount - 1, 0.5f)); + raw.add(i / (float) VECTORS_PER_THREAD); + return VectorType.getInstance(FloatType.instance, dimensionCount).getSerializer().serialize(raw); } private Expression generateRandomExpression() @@ -219,11 +720,21 @@ public class VectorMemoryIndexTest extends SAITester return expression; } - private ByteBuffer randomVector() { + private ByteBuffer randomVector() + { + return randomVector(() -> getRandom().nextFloat()); + } + + private ByteBuffer randomVectorFromThreadLocal() + { + return randomVector(() -> ThreadLocalRandom.current().nextFloat()); + } + + private ByteBuffer randomVector(Supplier supplier) + { List rawVector = new ArrayList<>(dimensionCount); - for (int i = 0; i < dimensionCount; i++) { - rawVector.add(getRandom().nextFloat()); - } + for (int i = 0; i < dimensionCount; i++) + rawVector.add(supplier.get()); return VectorType.getInstance(FloatType.instance, dimensionCount).getSerializer().serialize(rawVector); }