Merge branch 'cassandra-5.0' into cassandra-6.0

* cassandra-5.0:
  Make synchronization on VectorMemoryIndex inserts more granular

Author: maedhroz

CHANGES.txt

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

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 PrimaryKey primaryKey()
         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

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<PrimaryKey> primaryKeys = new ConcurrentSkipListSet<>();
 
@@ -88,7 +87,7 @@ public VectorMemoryIndex(StorageAttachedIndex index, Memtable memtable)
     }
 
     @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 synchronized long add(DecoratedKey key, Clustering<?> clustering, ByteBuf
 
     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 long update(DecoratedKey key, Clustering<?> clustering, ByteBuffer oldVal
         {
             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 long update(DecoratedKey key, Clustering<?> clustering, ByteBuffer oldVal
             // 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 CloseableIterator<PrimaryKeyWithScore> orderBy(QueryContext queryContext,
     @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 void close()
             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);
+        }
+    }
 }

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 @@
 /**
  * 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()