natas_hw
/
cassandra
mirror of https://github.com/apache/cassandra
0
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'cassandra-5.0' into cassandra-6.0 

* cassandra-5.0:
  Make synchronization on VectorMemoryIndex inserts more granular
This commit is contained in:
Caleb Rackliffe 2026-07-02 15:40:27 -05:00
parent 5282f81d0a 448d98ce31
commit bf4437296b
7 changed files with 1659 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGES.txt
View File

@ -59,6 +59,7 @@
* Fix a removed TTLed row re-appearance in a materialized view after a cursor compaction (CASSANDRA-21152)
* Rework ZSTD dictionary compression logic to create a trainer per training (CASSANDRA-21209)
Merged from 5.0:
* Make synchronization on VectorMemoryIndex inserts more granular (CASSANDRA-21160)
* putShortVolatile is not volatile in InMemoryTrie (CASSANDRA-21353)
* Fix RequestFailureReason serializer and nits in a few others (CASSANDRA-21437)
* Remove golang dependency in gen-doc and replace with python implementation (CASSANDRA-21432)

8
src/java/org/apache/cassandra/index/sai/disk/v1/vector/PrimaryKeyWithScore.java
View File

@ -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<PrimaryKeyWithScore>
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

59
src/java/org/apache/cassandra/index/sai/memory/VectorMemoryIndex.java
View File

@ -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<PrimaryKey> 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<PrimaryKeyWithScore> orderResultsBy(QueryContext queryContext, List<PrimaryKey> 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<PrimaryKey> 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);
}
}
}

50
test/simulator/main/org/apache/cassandra/simulator/asm/NemesisFieldSelectors.java
View File

@ -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<String, Map<String, NemesisFieldKind>> classToFieldToNemesis;
public static final ConcurrentHashMap<String, Map<String, NemesisFieldKind>> 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<String, Map<String, NemesisFieldKind>> byField = new HashMap<>();
ConcurrentHashMap<String, Map<String, NemesisFieldKind>> 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<String, NemesisFieldKind> fields = classToFieldToNemesis.get(dotsToSlashes(clazz));
if (fields != null)
fields.remove(fieldName);
}
public static NemesisFieldKind.Selector get()

283
test/simulator/test/org/apache/cassandra/simulator/test/OnHeapGraphSimulationTest.java Normal file
View File

@ -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.
* <p>
* 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.
* <p>
* 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<Integer> 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<Integer> graph = (OnHeapGraph<Integer>) 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<Integer> 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<SearchResult.NodeScore> 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<Integer> graph = createGraph();
int insertCount = NUM_THREADS * VECTORS_PER_THREAD;
// First, insert all vectors
List<ByteBuffer> 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<SearchResult.NodeScore> 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<Integer> 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<Float> 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;
}
}

750
test/simulator/test/org/apache/cassandra/simulator/test/VectorMemoryIndexSimulationTest.java Normal file
View File

@ -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.
* <p>
* The nemesis framework injects scheduling perturbation around mutable field accesses in
* {@link OnHeapGraph}, {@link VectorPostings}, and {@link VectorMemoryIndex}.
* <p>
* 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.
* <p>
* Tests ported from {@link org.apache.cassandra.index.sai.memory.VectorMemoryIndexTest}
* <ul>
* <li>{@link #testConcurrentAddsWithRandomVectors()} — N writers with random vectors, verify no data loss</li>
* <li>{@link #testConcurrentAddsWithSharedVectors()} — N writers with shared vectors, verify no data loss</li>
* <li>{@link #testConcurrentAddsAndOrderByRandomVectors()} — writers + readers via orderBy() with random vectors</li>
* <li>{@link #testConcurrentAddsAndOrderBySharedVectors()} — writers + readers via orderBy() with shared vectors</li>
* <li>{@link #testConcurrentAddsAndOrderResultsByRandomVectors()} — writers + readers via orderResultsBy() with random vectors</li>
* <li>{@link #testConcurrentAddsAndOrderResultsBySharedVectors()} — writers + readers via orderResultsBy() with shared vectors</li>
* </ul>
*/
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.
* <p>
* 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().
* <p>
* 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<Integer, DecoratedKey> 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<Integer, DecoratedKey> keyMap = (ConcurrentMap<Integer, DecoratedKey>) 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<PartitionPosition> 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<Integer> foundKeys = new HashSet<>();
try (CloseableIterator<PrimaryKeyWithScore> 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.
* <p>
* 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.
* <p>
* 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<Integer, DecoratedKey> 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<PartitionPosition> 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<PrimaryKeyWithScore> 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<PartitionPosition> 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<Integer> foundKeys = new HashSet<>();
try (CloseableIterator<PrimaryKeyWithScore> 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.
* <p>
* 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<Integer, DecoratedKey> 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<PrimaryKey> 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<PrimaryKeyWithScore> 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<Integer, DecoratedKey> keyMap = (ConcurrentMap<Integer, DecoratedKey>) SharedState.keyMap;
int total = SharedState.totalInserted;
List<PrimaryKey> 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<Integer> foundKeys = new HashSet<>();
try (CloseableIterator<PrimaryKeyWithScore> 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).
* <p>
* Only the fields actually accessed by VectorMemoryIndex methods are set:
* <ul>
* <li>{@code indexTermType} — for decomposeVector(), indexType(), columnMetadata()</li>
* <li>{@code indexWriterConfig} — for graph construction and getSimilarityFunction()</li>
* <li>{@code primaryKeyFactory} — for creating PrimaryKey instances</li>
* </ul>
*/
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<Float> 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<Float> 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;
}
}

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

@ -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}.
* <p>
* 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<DecoratedKey, Integer> keyMap;
private ConcurrentMap<DecoratedKey, Integer> keyMap;
private Map<Integer, ByteBuffer> 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<DecoratedKey> 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<Integer> 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<PrimaryKeyWithScore> 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.
* <p>
* 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.
* <p>
* 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().
* <p>
* 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
{
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<Future<?>> 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<PartitionPosition> 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<Integer> foundKeys = new HashSet<>();
try (CloseableIterator<PrimaryKeyWithScore> 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.
* <p>
* 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.
* <p>
* 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<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();
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<Throwable> 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<PartitionPosition> 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<PrimaryKeyWithScore> 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<Integer> foundAfterSettle = new HashSet<>();
try (CloseableIterator<PrimaryKeyWithScore> 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.
* <p>
* 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<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();
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<Throwable> 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<PrimaryKey> snapshotKeys = buildSortedPrimaryKeySnapshot();
ByteBuffer queryBuf = randomVectorFromThreadLocal();
Expression concurrentExpression = Expression.create(index);
concurrentExpression.add(Operator.ANN, queryBuf);
if (!snapshotKeys.isEmpty())
{
try (CloseableIterator<PrimaryKeyWithScore> 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<PrimaryKey> allKeys = buildSortedPrimaryKeySnapshot();
ByteBuffer settledQueryBuf = randomVectorFromThreadLocal();
Expression settledExpression = Expression.create(index);
settledExpression.add(Operator.ANN, settledQueryBuf);
Set<Integer> foundAfterSettle = new HashSet<>();
try (CloseableIterator<PrimaryKeyWithScore> 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<PrimaryKey> 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<Float> 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<Float> supplier)
{
List<Float> 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);
}