CEP-45: Add mutation tracking support for secondary index reads

This patch enables mutation tracking for queries using secondary indexes (both
legacy 2i and SAI). Key changes include the addition of the MultiStepSearcher
interface, which allows existing index implementations to iterate over matches
and filter augmented data, proper Memtable snapshotting, and the integration of
these in PartialTrackedIndexRead.

patch by Blake Eggleston; reviewed by Aleksey Yeschenko and Caleb Rackliffe for CASSANDRA-20374

Co-authored-by: Blake Eggleston <blake@ultrablake.com>
Co-authored-by: Caleb Rackliffe <calebrackliffe@gmail.com>
Co-authored-by: Aleksey Yeschenko <aleksey@apache.org>

src/java/org/apache/cassandra/cql3/statements/SelectStatement.java

@@ -116,6 +116,7 @@ import org.apache.cassandra.index.Index;
 import org.apache.cassandra.index.IndexRegistry;
 import org.apache.cassandra.metrics.ClientRequestSizeMetrics;
 import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.ReplicationType;
 import org.apache.cassandra.schema.Schema;
 import org.apache.cassandra.schema.SchemaConstants;
 import org.apache.cassandra.schema.TableMetadata;
@@ -483,7 +484,7 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement,
             return getRangeCommand(options, state, columnFilter, rowFilter, limit, nowInSec, potentialTxnConflicts);
         }
 
-        if (restrictions.usesSecondaryIndexing() && !rowFilter.isStrict())
+        if (restrictions.usesSecondaryIndexing() && !rowFilter.isStrict() && table.keyspaceReplicationType != ReplicationType.tracked)
             return getRangeCommand(options, state, columnFilter, rowFilter, limit, nowInSec, potentialTxnConflicts);
 
         return getSliceCommands(options, state, columnFilter, rowFilter, limit, nowInSec, potentialTxnConflicts);

src/java/org/apache/cassandra/db/ColumnFamilyStore.java

@@ -101,6 +101,7 @@ import org.apache.cassandra.db.partitions.CachedPartition;
 import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.repair.CassandraTableRepairManager;
 import org.apache.cassandra.db.rows.CellPath;
+import org.apache.cassandra.db.rows.UnfilteredSource;
 import org.apache.cassandra.db.streaming.CassandraStreamManager;
 import org.apache.cassandra.db.view.TableViews;
 import org.apache.cassandra.dht.AbstractBounds;
@@ -3098,7 +3099,7 @@ public <T> T withAllSSTables(final OperationType operationType, Function<Lifecyc
         return compactionStrategyManager.getLevelFanoutSize();
     }
 
-    public static class ViewFragment
+    public static class ViewFragment implements ReadableView
     {
         public final List<SSTableReader> sstables;
         public final Iterable<Memtable> memtables;
@@ -3108,6 +3109,18 @@ public <T> T withAllSSTables(final OperationType operationType, Function<Lifecyc
             this.sstables = sstables;
             this.memtables = memtables;
         }
+
+        @Override
+        public Iterable<? extends UnfilteredSource> memtables()
+        {
+            return memtables;
+        }
+
+        @Override
+        public List<SSTableReader> sstables()
+        {
+            return sstables;
+        }
     }
 
     public static class RefViewFragment extends ViewFragment implements AutoCloseable

src/java/org/apache/cassandra/db/PartitionRangeReadCommand.java

@@ -449,11 +449,10 @@ public class PartitionRangeReadCommand extends ReadCommand implements PartitionR
     @Override
     protected PartialTrackedRead createInProgressRead(UnfilteredPartitionIterator iterator,
                                                       ReadExecutionController executionController,
-                                                      Index.Searcher searcher,
                                                       ColumnFamilyStore cfs,
                                                       long startTimeNanos)
     {
-        return PartialTrackedRangeRead.create(executionController, searcher, cfs, startTimeNanos, this, iterator);
+        return PartialTrackedRangeRead.create(executionController, null, cfs, startTimeNanos, this, iterator);
     }
 
     @Override

src/java/org/apache/cassandra/db/ReadCommand.java

@@ -130,7 +130,42 @@ public abstract class ReadCommand extends AbstractReadQuery
 {
     private interface ReadCompleter<T>
     {
-        T complete(UnfilteredPartitionIterator iterator, ReadExecutionController executionController, Index.Searcher searcher, ColumnFamilyStore cfs, long startTimeNanos);
+        T complete(ReadCommand command, UnfilteredPartitionIterator iterator, ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos);
+        T complete(ReadCommand command, Index.Searcher searcher, ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos);
+
+        ReadCompleter<UnfilteredPartitionIterator> IMMEDIATE = new ReadCompleter<>()
+        {
+            @Override
+            public UnfilteredPartitionIterator complete(ReadCommand command, UnfilteredPartitionIterator iterator, ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
+            {
+                return command.completeRead(iterator, executionController, null, cfs, startTimeNanos);
+            }
+
+            @Override
+            public UnfilteredPartitionIterator complete(ReadCommand command, Index.Searcher searcher, ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
+            {
+                UnfilteredPartitionIterator iterator = searcher.search(executionController);
+                return command.completeRead(iterator, executionController, searcher, cfs, startTimeNanos);
+            }
+        };
+
+        ReadCompleter<PartialTrackedRead> TRACKED = new ReadCompleter<>()
+        {
+            @Override
+            public PartialTrackedRead complete(ReadCommand command, UnfilteredPartitionIterator iterator, ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
+            {
+                return command.createInProgressRead(iterator, executionController, cfs, startTimeNanos);
+            }
+
+            @Override
+            public PartialTrackedRead complete(ReadCommand command, Index.Searcher searcher, ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
+            {
+                if (!searcher.isMultiStep())
+                    throw new IllegalStateException("Cannot use " + searcher.getClass().getName() + " with tracked reads");
+
+                return searcher.asMultiStep().beginRead(executionController, cfs, startTimeNanos);
+            }
+        };
     }
 
     private static final int TEST_ITERATION_DELAY_MILLIS = CassandraRelevantProperties.TEST_READ_ITERATION_DELAY_MS.getInt();
@@ -483,28 +518,26 @@ public abstract class ReadCommand extends AbstractReadQuery
                 ConsensusRequestRouter.validateSafeToReadNonTransactionally(this, cm);
             Index.QueryPlan indexQueryPlan = indexQueryPlan();
 
-            Index.Searcher searcher = null;
             if (indexQueryPlan != null)
             {
                 cfs.indexManager.checkQueryability(indexQueryPlan);
 
-                searcher = indexQueryPlan.searcherFor(this);
-                Tracing.trace("Executing read on {}.{} using index{} {}",
-                              cfs.metadata.keyspace,
-                              cfs.metadata.name,
-                              indexQueryPlan.getIndexes().size() == 1 ? "" : "es",
-                              indexQueryPlan.getIndexes()
-                                            .stream()
-                                            .map(i -> i.getIndexMetadata().name)
-                                            .collect(Collectors.joining(",")));
+                if (Tracing.isTracing())
+                {
+                    Tracing.trace("Executing read on {}.{} using index{} {}",
+                                  cfs.metadata.keyspace,
+                                  cfs.metadata.name,
+                                  indexQueryPlan.getIndexes().size() == 1 ? "" : "es",
+                                  indexQueryPlan.getIndexes()
+                                                .stream()
+                                                .map(i -> i.getIndexMetadata().name)
+                                                .collect(Collectors.joining(",")));
+                }
+                Index.Searcher searcher = indexQueryPlan.searcherFor(this);
+                return completer.complete(this, searcher, executionController, cfs, startTimeNanos);
             }
-
-            if (searcher != null && metadata().replicationType().isTracked())
-                throw new UnsupportedOperationException("TODO: support tracked index reads");
-
-            UnfilteredPartitionIterator iterator = (null == searcher) ? queryStorage(cfs, executionController) : searcher.search(executionController);
-
-            return completer.complete(iterator, executionController, searcher, cfs, startTimeNanos);
+            UnfilteredPartitionIterator iterator = queryStorage(cfs, executionController);
+            return completer.complete(this, iterator, executionController, cfs, startTimeNanos);
         }
         finally
         {
@@ -584,13 +617,12 @@ public abstract class ReadCommand extends AbstractReadQuery
      */
     protected abstract PartialTrackedRead createInProgressRead(UnfilteredPartitionIterator iterator,
                                                                ReadExecutionController executionController,
-                                                               Index.Searcher searcher,
                                                                ColumnFamilyStore cfs,
                                                                long startTimeNanos);
 
     public PartialTrackedRead beginTrackedRead(ReadExecutionController executionController)
     {
-        return beginRead(executionController, null, this::createInProgressRead);
+        return beginRead(executionController, null, ReadCompleter.TRACKED);
     }
 
     public UnfilteredPartitionIterator completeTrackedRead(UnfilteredPartitionIterator iterator, PartialTrackedRead read)
@@ -608,12 +640,12 @@ public abstract class ReadCommand extends AbstractReadQuery
     // iterators created inside the try as long as we do close the original resultIterator), or by closing the result.
     public UnfilteredPartitionIterator executeLocally(ReadExecutionController executionController)
     {
-        return beginRead(executionController, null, this::completeRead);
+        return beginRead(executionController, null, ReadCompleter.IMMEDIATE);
     }
 
     public UnfilteredPartitionIterator executeLocally(ReadExecutionController executionController, @Nullable ClusterMetadata cm)
     {
-        return beginRead(executionController, cm, this::completeRead);
+        return beginRead(executionController, cm, ReadCompleter.IMMEDIATE);
     }
 
     protected abstract void recordLatency(TableMetrics metric, long latencyNanos);
@@ -1148,7 +1180,7 @@ public abstract class ReadCommand extends AbstractReadQuery
         return toCQLString();
     }
 
-    InputCollector<UnfilteredRowIterator> iteratorsForPartition(ColumnFamilyStore.ViewFragment view, ReadExecutionController controller)
+    InputCollector<UnfilteredRowIterator> iteratorsForPartition(ReadableView view, ReadExecutionController controller)
     {
         final BiFunction<List<UnfilteredRowIterator>, RepairedDataInfo, UnfilteredRowIterator> merge =
             (unfilteredRowIterators, repairedDataInfo) -> {
@@ -1198,7 +1230,7 @@ public abstract class ReadCommand extends AbstractReadQuery
         List<T> repairedIters;
         List<T> unrepairedIters;
 
-        InputCollector(ColumnFamilyStore.ViewFragment view,
+        InputCollector(ReadableView view,
                        ReadExecutionController controller,
                        BiFunction<List<T>, RepairedDataInfo, T> repairedMerger,
                        Function<T, UnfilteredPartitionIterator> postLimitAdditionalPartitions)
@@ -1208,12 +1240,12 @@ public abstract class ReadCommand extends AbstractReadQuery
             
             if (isTrackingRepairedStatus)
             {
-                for (SSTableReader sstable : view.sstables)
+                for (SSTableReader sstable : view.sstables())
                 {
                     if (considerRepairedForTracking(sstable))
                     {
                         if (repairedSSTables == null)
-                            repairedSSTables = Sets.newHashSetWithExpectedSize(view.sstables.size());
+                            repairedSSTables = Sets.newHashSetWithExpectedSize(view.sstables().size());
                         repairedSSTables.add(sstable);
                     }
                 }
@@ -1221,14 +1253,14 @@ public abstract class ReadCommand extends AbstractReadQuery
             if (repairedSSTables == null)
             {
                 repairedIters = Collections.emptyList();
-                unrepairedIters = new ArrayList<>(view.sstables.size());
+                unrepairedIters = new ArrayList<>(view.sstables().size());
             }
             else
             {
                 repairedIters = new ArrayList<>(repairedSSTables.size());
                 // when we're done collating, we'll merge the repaired iters and add the
                 // result to the unrepaired list, so size that list accordingly
-                unrepairedIters = new ArrayList<>((view.sstables.size() - repairedSSTables.size()) + Iterables.size(view.memtables) + 1);
+                unrepairedIters = new ArrayList<>((view.sstables().size() - repairedSSTables.size()) + Iterables.size(view.memtables()) + 1);
             }
             this.repairedMerger = repairedMerger;
             this.postLimitAdditionalPartitions = postLimitAdditionalPartitions;

src/java/org/apache/cassandra/db/ReadableView.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.db;
+
+import java.util.List;
+
+import org.apache.cassandra.db.rows.UnfilteredSource;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+
+public interface ReadableView
+{
+    Iterable<? extends UnfilteredSource> memtables();
+    List<SSTableReader> sstables();
+}

src/java/org/apache/cassandra/db/SinglePartitionReadCommand.java

@@ -65,6 +65,7 @@ import org.apache.cassandra.db.rows.Unfiltered;
 import org.apache.cassandra.db.rows.UnfilteredRowIterator;
 import org.apache.cassandra.db.rows.UnfilteredRowIteratorWithLowerBound;
 import org.apache.cassandra.db.rows.UnfilteredRowIterators;
+import org.apache.cassandra.db.rows.UnfilteredSource;
 import org.apache.cassandra.db.rows.WrappingUnfilteredRowIterator;
 import org.apache.cassandra.db.transform.RTBoundValidator;
 import org.apache.cassandra.db.transform.Transformation;
@@ -558,11 +559,10 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
     @Override
     protected PartialTrackedRead createInProgressRead(UnfilteredPartitionIterator iterator,
                                                       ReadExecutionController executionController,
-                                                      Index.Searcher searcher,
                                                       ColumnFamilyStore cfs,
                                                       long startTimeNanos)
     {
-        return PartialTrackedSinglePartitionRead.create(executionController, searcher, cfs, startTimeNanos, this, iterator);
+        return PartialTrackedSinglePartitionRead.create(executionController, null, cfs, startTimeNanos, this, iterator);
     }
 
     /**
@@ -726,29 +726,32 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
      * Also note that one must have created a {@code ReadExecutionController} on the queried table and we require it as
      * a parameter to enforce that fact, even though it's not explicitlly used by the method.
      */
-    public UnfilteredRowIterator queryMemtableAndDisk(ColumnFamilyStore cfs, ReadExecutionController executionController)
+    public UnfilteredRowIterator queryMemtableAndDisk(ReadableView view, ColumnFamilyStore cfs, ReadExecutionController executionController)
     {
         assert executionController != null && executionController.validForReadOn(cfs);
         Tracing.trace("Executing single-partition query on {}", cfs.name);
 
-        Tracing.trace("Acquiring sstable references");
-        ColumnFamilyStore.ViewFragment view = cfs.select(View.select(SSTableSet.LIVE, partitionKey()));
-        return queryMemtableAndDiskInternal(cfs, view, null, executionController);
+        return queryMemtableAndDiskInternal(view, cfs, null, executionController);
     }
 
-    public UnfilteredRowIterator queryMemtableAndDisk(ColumnFamilyStore cfs,
-                                                      ColumnFamilyStore.ViewFragment view,
+    public UnfilteredRowIterator queryMemtableAndDisk(ColumnFamilyStore cfs, ReadExecutionController executionController)
+    {
+        return queryMemtableAndDisk(cfs.select(View.select(SSTableSet.LIVE, partitionKey())), cfs, executionController);
+    }
+
+    public UnfilteredRowIterator queryMemtableAndDisk(ReadableView view,
+                                                      ColumnFamilyStore cfs,
                                                       Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer,
                                                       ReadExecutionController executionController)
     {
         assert executionController != null && executionController.validForReadOn(cfs);
         Tracing.trace("Executing single-partition query on {}", cfs.name);
 
-        return queryMemtableAndDiskInternal(cfs, view, rowTransformer, executionController);
+        return queryMemtableAndDiskInternal(view, cfs, rowTransformer, executionController);
     }
 
-    private UnfilteredRowIterator queryMemtableAndDiskInternal(ColumnFamilyStore cfs,
-                                                               ColumnFamilyStore.ViewFragment view,
+    private UnfilteredRowIterator queryMemtableAndDiskInternal(ReadableView view,
+                                                               ColumnFamilyStore cfs,
                                                                Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer,
                                                                ReadExecutionController controller)
     {
@@ -774,10 +777,12 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
             && !queriesMulticellType()
             && !controller.isTrackingRepairedStatus())
         {
-            return queryMemtableAndSSTablesInTimestampOrder(cfs, view, rowTransformer, (ClusteringIndexNamesFilter)clusteringIndexFilter(), controller);
+            return queryMemtableAndSSTablesInTimestampOrder(view, cfs, rowTransformer, (ClusteringIndexNamesFilter)clusteringIndexFilter(), controller);
         }
 
-        view.sstables.sort(SSTableReader.maxTimestampDescending);
+        Tracing.trace("Acquiring sstable references");
+        List<SSTableReader> sstables = view.sstables();
+        sstables.sort(SSTableReader.maxTimestampDescending);
         ClusteringIndexFilter filter = clusteringIndexFilter();
         long minTimestamp = Long.MAX_VALUE;
         long mostRecentPartitionTombstone = Long.MIN_VALUE;
@@ -786,7 +791,7 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
         {
             SSTableReadMetricsCollector metricsCollector = new SSTableReadMetricsCollector();
 
-            for (Memtable memtable : view.memtables)
+            for (UnfilteredSource memtable : view.memtables())
             {
                 UnfilteredRowIterator iter = memtable.rowIterator(partitionKey(), filter.getSlices(metadata()), columnFilter(), filter.isReversed(), metricsCollector);
                 if (iter == null)
@@ -795,8 +800,8 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
                 if (memtable.getMinTimestamp() != Memtable.NO_MIN_TIMESTAMP)
                     minTimestamp = Math.min(minTimestamp, memtable.getMinTimestamp());
 
-                if (rowTransformer != null)
-                    iter = Transformation.apply(iter, rowTransformer.apply(memtable));
+                if (rowTransformer != null && memtable instanceof CellSourceIdentifier)
+                    iter = Transformation.apply(iter, rowTransformer.apply((CellSourceIdentifier) memtable));
 
                 // Memtable data is always considered unrepaired
                 controller.updateMinOldestUnrepairedTombstone(memtable.getMinLocalDeletionTime());
@@ -818,14 +823,14 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
              * In other words, iterating in descending maxTimestamp order allow to do our mostRecentPartitionTombstone
              * elimination in one pass, and minimize the number of sstables for which we read a partition tombstone.
             */
-            view.sstables.sort(SSTableReader.maxTimestampDescending);
+            sstables.sort(SSTableReader.maxTimestampDescending);
             int nonIntersectingSSTables = 0;
             int includedDueToTombstones = 0;
 
             if (controller.isTrackingRepairedStatus())
                 Tracing.trace("Collecting data from sstables and tracking repaired status");
 
-            for (SSTableReader sstable : view.sstables)
+            for (SSTableReader sstable : sstables)
             {
                 // if we've already seen a partition tombstone with a timestamp greater
                 // than the most recent update to this sstable, we can skip it
@@ -900,7 +905,7 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
 
             if (Tracing.isTracing())
                 Tracing.trace("Skipped {}/{} non-slice-intersecting sstables, included {} due to tombstones",
-                               nonIntersectingSSTables, view.sstables.size(), includedDueToTombstones);
+                               nonIntersectingSSTables, sstables.size(), includedDueToTombstones);
 
             if (inputCollector.isEmpty())
                 return EmptyIterators.unfilteredRow(cfs.metadata(), partitionKey(), filter.isReversed());
@@ -1030,21 +1035,23 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
      * no collection or counters are included).
      * This method assumes the filter is a {@code ClusteringIndexNamesFilter}.
      */
-    private UnfilteredRowIterator queryMemtableAndSSTablesInTimestampOrder(ColumnFamilyStore cfs, ColumnFamilyStore.ViewFragment view, Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer, ClusteringIndexNamesFilter filter, ReadExecutionController controller)
+    private UnfilteredRowIterator queryMemtableAndSSTablesInTimestampOrder(ReadableView view, ColumnFamilyStore cfs, Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer, ClusteringIndexNamesFilter filter, ReadExecutionController controller)
     {
+        Tracing.trace("Acquiring sstable references");
         ImmutableBTreePartition result = null;
         SSTableReadMetricsCollector metricsCollector = new SSTableReadMetricsCollector();
 
         Tracing.trace("Merging memtable contents");
-        for (Memtable memtable : view.memtables)
+        for (UnfilteredSource memtable : view.memtables())
         {
             try (UnfilteredRowIterator iter = memtable.rowIterator(partitionKey, filter.getSlices(metadata()), columnFilter(), isReversed(), metricsCollector))
             {
                 if (iter == null)
                     continue;
 
-                UnfilteredRowIterator wrapped = rowTransformer != null ? Transformation.apply(iter, rowTransformer.apply(memtable))
-                                                                       : iter;
+                UnfilteredRowIterator wrapped = rowTransformer != null && memtable instanceof CellSourceIdentifier
+                                                ? Transformation.apply(iter, rowTransformer.apply((CellSourceIdentifier) memtable))
+                                                : iter;
                 result = add(RTBoundValidator.validate(wrapped, RTBoundValidator.Stage.MEMTABLE, false),
                              result,
                              filter,
@@ -1054,9 +1061,10 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
         }
 
         /* add the SSTables on disk */
-        view.sstables.sort(SSTableReader.maxTimestampDescending);
+        List<SSTableReader> sstables = view.sstables();
+        sstables.sort(SSTableReader.maxTimestampDescending);
         // read sorted sstables
-        for (SSTableReader sstable : view.sstables)
+        for (SSTableReader sstable : sstables)
         {
             // if we've already seen a partition tombstone with a timestamp greater
             // than the most recent update to this sstable, we're done, since the rest of the sstables

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

@@ -26,6 +26,7 @@ import javax.annotation.concurrent.NotThreadSafe;
 
 import org.apache.cassandra.db.CellSourceIdentifier;
 import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.PartitionPosition;
 import org.apache.cassandra.db.RegularAndStaticColumns;
 import org.apache.cassandra.db.commitlog.CommitLogPosition;
@@ -206,7 +207,17 @@ public interface Memtable extends Comparable<Memtable>, UnfilteredSource, CellSo
      */
     long put(MutationId mutationId, PartitionUpdate update, UpdateTransaction indexer, OpOrder.Group opGroup, boolean assumeMissing);
 
-    // Read operations are provided by the UnfilteredSource interface.
+    /**
+     * Creates a point-in-time snapshot of a partition in this memtable.
+     * <p>
+     * This method returns an immutable view of the partition as it exists at the time of the call.
+     * The snapshot is isolated from subsequent writes to the memtable and can be safely read
+     * concurrently with ongoing mutations.
+     *
+     * @param key the partition key to snapshot
+     * @return an immutable snapshot of the partition, or {@code null} if the partition does not exist
+     */
+    Partition snapshotPartition(DecoratedKey key);
 
     // Statistics
 

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

@@ -262,7 +262,7 @@ public class ShardedSkipListMemtable extends AbstractShardedMemtable
         return iterator;
     }
 
-    private Partition getPartition(DecoratedKey key)
+    private AtomicBTreePartition getPartition(DecoratedKey key)
     {
         int shardIndex = boundaries.getShardForKey(key);
         return shards[shardIndex].partitions.get(key);
@@ -285,6 +285,13 @@ public class ShardedSkipListMemtable extends AbstractShardedMemtable
         return p != null ? p.unfilteredIterator() : null;
     }
 
+    @Override
+    public Partition snapshotPartition(DecoratedKey partitionKey)
+    {
+        AtomicBTreePartition partition = getPartition(partitionKey);
+        return partition == null ? null : partition.asImmutable();
+    }
+
     public FlushablePartitionSet<AtomicBTreePartition> getFlushSet(PartitionPosition from, PartitionPosition to)
     {
         long keySize = 0;

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

@@ -207,7 +207,7 @@ public class SkipListMemtable extends AbstractAllocatorMemtable
         }
     }
 
-    Partition getPartition(DecoratedKey key)
+    AtomicBTreePartition getPartition(DecoratedKey key)
     {
         return partitions.get(key);
     }
@@ -229,6 +229,13 @@ public class SkipListMemtable extends AbstractAllocatorMemtable
         return p != null ? p.unfilteredIterator() : null;
     }
 
+    @Override
+    public Partition snapshotPartition(DecoratedKey partitionKey)
+    {
+        AtomicBTreePartition partition = getPartition(partitionKey);
+        return partition == null ? null : partition.asImmutable();
+    }
+
     private static int estimateRowOverhead(final int count)
     {
         // calculate row overhead

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

@@ -319,15 +319,26 @@ public class TrieMemtable extends AbstractShardedMemtable
     }
 
     private Partition getPartition(DecoratedKey key)
+    {
+        return getPartition(key, false);
+    }
+
+    private Partition getPartition(DecoratedKey key, boolean canHaveShadowedData)
     {
         int shardIndex = boundaries.getShardForKey(key);
         BTreePartitionData data = shards[shardIndex].data.get(key);
         if (data != null)
-            return createPartition(metadata(), allocator.ensureOnHeap(), key, data);
+            return createPartition(metadata(), allocator.ensureOnHeap(), key, data, canHaveShadowedData);
         else
             return null;
     }
 
+    @Override
+    public Partition snapshotPartition(DecoratedKey partitionKey)
+    {
+        return getPartition(partitionKey, true);
+    }
+
     @Override
     public UnfilteredRowIterator rowIterator(DecoratedKey key, Slices slices, ColumnFilter selectedColumns, boolean reversed, SSTableReadsListener listener)
     {
@@ -345,9 +356,9 @@ public class TrieMemtable extends AbstractShardedMemtable
         return p != null ? p.unfilteredIterator() : null;
     }
 
-    private static MemtablePartition createPartition(TableMetadata metadata, EnsureOnHeap ensureOnHeap, DecoratedKey key, BTreePartitionData data)
+    private static MemtablePartition createPartition(TableMetadata metadata, EnsureOnHeap ensureOnHeap, DecoratedKey key, BTreePartitionData data, boolean canHaveShadowedData)
     {
-        return new MemtablePartition(metadata, ensureOnHeap, key, data);
+        return new MemtablePartition(metadata, ensureOnHeap, key, data, canHaveShadowedData);
     }
 
     private static MemtablePartition getPartitionFromTrieEntry(TableMetadata metadata, EnsureOnHeap ensureOnHeap, Map.Entry<ByteComparable, BTreePartitionData> en)
@@ -355,7 +366,7 @@ public class TrieMemtable extends AbstractShardedMemtable
         DecoratedKey key = BufferDecoratedKey.fromByteComparable(en.getKey(),
                                                                  BYTE_COMPARABLE_VERSION,
                                                                  metadata.partitioner);
-        return createPartition(metadata, ensureOnHeap, key, en.getValue());
+        return createPartition(metadata, ensureOnHeap, key, en.getValue(), false);
     }
 
 
@@ -696,9 +707,9 @@ public class TrieMemtable extends AbstractShardedMemtable
 
         private final EnsureOnHeap ensureOnHeap;
 
-        private MemtablePartition(TableMetadata table, EnsureOnHeap ensureOnHeap, DecoratedKey key, BTreePartitionData data)
+        private MemtablePartition(TableMetadata table, EnsureOnHeap ensureOnHeap, DecoratedKey key, BTreePartitionData data, boolean canHaveShadowedData)
         {
-            super(table, key, data);
+            super(table, key, data, canHaveShadowedData);
             this.ensureOnHeap = ensureOnHeap;
         }
 

src/java/org/apache/cassandra/db/partitions/AtomicBTreePartition.java

@@ -46,7 +46,7 @@ import static org.apache.cassandra.utils.Clock.Global.nanoTime;
 
 /**
  * A thread-safe and atomic Partition implementation.
- *
+ * <p>
  * Operations (in particular addAll) on this implementation are atomic and
  * isolated (in the sense of ACID). Typically a addAll is guaranteed that no
  * other thread can see the state where only parts but not all rows have
@@ -77,7 +77,7 @@ public final class AtomicBTreePartition extends AbstractBTreePartition
     /**
      * (clock + allocation) granularity are combined to give us an acceptable (waste) allocation rate that is defined by
      * the passage of real time of ALLOCATION_GRANULARITY_BYTES/CLOCK_GRANULARITY, or in this case 7.63KiB/ms, or 7.45Mb/s
-     *
+     * <p>
      * in wasteTracker we maintain within EXCESS_WASTE_OFFSET before the current time; whenever we waste bytes
      * we increment the current value if it is within this window, and set it to the min of the window plus our waste
      * otherwise.
@@ -101,6 +101,11 @@ public final class AtomicBTreePartition extends AbstractBTreePartition
         this.ref = BTreePartitionData.EMPTY;
     }
 
+    public ImmutableBTreePartition asImmutable()
+    {
+        return new ImmutableBTreePartition(metadata(), partitionKey, holder(), true);
+    }
+
     protected BTreePartitionData holder()
     {
         return ref;

src/java/org/apache/cassandra/db/partitions/ImmutableBTreePartition.java

@@ -32,6 +32,8 @@ public class ImmutableBTreePartition extends AbstractBTreePartition
     protected final BTreePartitionData holder;
     protected final TableMetadata metadata;
 
+    private final boolean canHaveShadowedData;
+
     public ImmutableBTreePartition(TableMetadata metadata,
                                    DecoratedKey partitionKey,
                                    RegularAndStaticColumns columns,
@@ -43,15 +45,25 @@ public class ImmutableBTreePartition extends AbstractBTreePartition
         super(partitionKey);
         this.metadata = metadata;
         this.holder = new BTreePartitionData(columns, tree, deletionInfo, staticRow, stats);
+        this.canHaveShadowedData = false;
     }
 
     protected ImmutableBTreePartition(TableMetadata metadata,
-                                      DecoratedKey partitionKey,
-                                      BTreePartitionData holder)
+                                   DecoratedKey partitionKey,
+                                   BTreePartitionData holder)
+    {
+        this(metadata, partitionKey, holder, false);
+    }
+
+    public ImmutableBTreePartition(TableMetadata metadata,
+                                   DecoratedKey partitionKey,
+                                   BTreePartitionData holder,
+                                   boolean canHaveShadowedData)
     {
         super(partitionKey);
         this.metadata = metadata;
         this.holder = holder;
+        this.canHaveShadowedData = canHaveShadowedData;
     }
 
     /**
@@ -128,6 +140,6 @@ public class ImmutableBTreePartition extends AbstractBTreePartition
 
     protected boolean canHaveShadowedData()
     {
-        return false;
+        return canHaveShadowedData;
     }
 }

src/java/org/apache/cassandra/db/partitions/UnfilteredPartitionIterators.java

@@ -427,4 +427,16 @@ public abstract class UnfilteredPartitionIterators
             };
         }
     }
+
+    public static void consume(UnfilteredPartitionIterator iterator)
+    {
+        while (iterator.hasNext())
+        {
+            try (UnfilteredRowIterator partition = iterator.next())
+            {
+                while (partition.hasNext())
+                    partition.next();
+            }
+        }
+    }
 }

src/java/org/apache/cassandra/index/Index.java

@@ -29,6 +29,7 @@ import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
@@ -36,6 +37,8 @@ import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
+import com.google.common.collect.PeekingIterator;
+
 import org.apache.cassandra.cql3.Operator;
 import org.apache.cassandra.cql3.QueryOptions;
 import org.apache.cassandra.cql3.restrictions.Restriction;
@@ -45,6 +48,7 @@ import org.apache.cassandra.db.DeletionTime;
 import org.apache.cassandra.db.RangeTombstone;
 import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
+import org.apache.cassandra.db.ReadableView;
 import org.apache.cassandra.db.RegularAndStaticColumns;
 import org.apache.cassandra.db.WriteContext;
 import org.apache.cassandra.db.filter.RowFilter;
@@ -55,6 +59,7 @@ import org.apache.cassandra.db.partitions.PartitionIterator;
 import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
 import org.apache.cassandra.db.rows.Row;
+import org.apache.cassandra.db.rows.UnfilteredRowIterator;
 import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.apache.cassandra.index.internal.CollatedViewIndexBuilder;
 import org.apache.cassandra.index.transactions.IndexTransaction;
@@ -68,6 +73,8 @@ import org.apache.cassandra.schema.ColumnMetadata;
 import org.apache.cassandra.schema.IndexMetadata;
 import org.apache.cassandra.schema.TableMetadata;
 import org.apache.cassandra.service.ClientState;
+import org.apache.cassandra.service.reads.tracked.PartialTrackedRead;
+import org.apache.cassandra.utils.CloseablePeekingIterator;
 
 /**
  * Consisting of a top level Index interface and two sub-interfaces which handle read and write operations,
@@ -750,6 +757,73 @@ public interface Index
         {
             return command().rowFilter().filter(fullResponse, command().metadata(), command().nowInSec());
         }
+
+        default boolean isMultiStep()
+        {
+            return false;
+        }
+
+        default MultiStepSearcher<?> asMultiStep()
+        {
+            throw new IllegalStateException(getClass().getSimpleName() + " is not a multi-step searcher");
+        }
+    }
+
+    interface IndexMatch
+    {
+        DecoratedKey key();
+    }
+
+    /**
+     * Extended searcher capable of participating in tracked reads
+     */
+    interface MultiStepSearcher<Match extends IndexMatch> extends Searcher
+    {
+
+        PartialTrackedRead beginRead(ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos);
+
+        CloseablePeekingIterator<Match> matchIterator(ReadExecutionController executionController);
+        MatchIndexer<Match> matchIndexer();
+
+        UnfilteredRowIterator queryNextMatches(ReadExecutionController executionController, DecoratedKey partitionKey, ReadableView view, PeekingIterator<Match> matches);
+
+        /**
+         * Since partition updates may not contain all the info the index query needs to know if it will create a hit
+         * it may return false positives. This filter is meant to catch and remove them from the augmented result
+         */
+        UnfilteredPartitionIterator filterCompletedRead(UnfilteredPartitionIterator iterator);
+
+        MatchComparator<Match> matchComparator();
+
+        @Override
+        default boolean isMultiStep()
+        {
+            return true;
+        }
+
+        @Override
+        default MultiStepSearcher<?> asMultiStep()
+        {
+            return this;
+        }
+
+        interface MatchIndexer<Match extends IndexMatch>
+        {
+            void index(PartitionUpdate update, Consumer<Match> indexTo);
+        }
+
+        interface MatchComparator<Match extends IndexMatch> extends Comparator<Match>
+        {
+            int compare(Match a, Match b, boolean strict);
+
+            @Override
+            default int compare(Match a, Match b)
+            {
+                return compare(a, b, true);
+            }
+
+            default void consumeDuplicates(Match original, PeekingIterator<Match> iterator) { }
+        }
     }
 
     /**

src/java/org/apache/cassandra/index/internal/CassandraIndex.java

@@ -356,6 +356,162 @@ public abstract class CassandraIndex implements Index
         }
     }
 
+    static abstract class AbstractIndexer implements Indexer
+    {
+        public void begin()
+        {
+        }
+
+        public void partitionDelete(DeletionTime deletionTime)
+        {
+        }
+
+        public void rangeTombstone(RangeTombstone tombstone)
+        {
+        }
+
+        abstract CassandraIndex index();
+
+        ColumnMetadata indexedColumn()
+        {
+            return index().indexedColumn;
+        }
+
+        boolean isPrimaryKeyIndex()
+        {
+            return index().isPrimaryKeyIndex();
+        }
+
+        abstract long nowInSec();
+        abstract DecoratedKey key();
+
+        abstract void insert(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info);
+        abstract void delete(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, long nowInSec);
+        abstract void delete(DecoratedKey rowKey, Clustering<?> clustering, DeletionTime deletion);
+
+        public void insertRow(Row row)
+        {
+            if (row.isStatic() && !indexedColumn().isStatic() && !indexedColumn().isPartitionKey())
+                return;
+
+            if (isPrimaryKeyIndex())
+            {
+                indexPrimaryKey(row.clustering(),
+                                getPrimaryKeyIndexLiveness(row),
+                                row.deletion());
+            }
+            else
+            {
+                if (indexedColumn().isComplex())
+                    indexCells(row.clustering(), row.getComplexColumnData(indexedColumn()));
+                else
+                    indexCell(row.clustering(), row.getCell(indexedColumn()));
+            }
+        }
+
+        public void removeRow(Row row)
+        {
+            if (isPrimaryKeyIndex())
+                return;
+
+            if (indexedColumn().isComplex())
+                removeCells(row.clustering(), row.getComplexColumnData(indexedColumn()));
+            else
+                removeCell(row.clustering(), row.getCell(indexedColumn()));
+        }
+
+        public void updateRow(Row oldRow, Row newRow)
+        {
+            assert oldRow.isStatic() == newRow.isStatic();
+            if (newRow.isStatic() != indexedColumn().isStatic())
+                return;
+
+            if (isPrimaryKeyIndex())
+                indexPrimaryKey(newRow.clustering(),
+                                getPrimaryKeyIndexLiveness(newRow),
+                                newRow.deletion());
+
+            if (indexedColumn().isComplex())
+            {
+                indexCells(newRow.clustering(), newRow.getComplexColumnData(indexedColumn()));
+                removeCells(oldRow.clustering(), oldRow.getComplexColumnData(indexedColumn()));
+            }
+            else
+            {
+                indexCell(newRow.clustering(), newRow.getCell(indexedColumn()));
+                removeCell(oldRow.clustering(), oldRow.getCell(indexedColumn()));
+            }
+        }
+
+        public void finish()
+        {
+        }
+
+        private void indexCells(Clustering<?> clustering, Iterable<Cell<?>> cells)
+        {
+            if (cells == null)
+                return;
+
+            for (Cell<?> cell : cells)
+                indexCell(clustering, cell);
+        }
+
+        private void indexCell(Clustering<?> clustering, Cell<?> cell)
+        {
+            if (cell == null || !cell.isLive(nowInSec()))
+                return;
+
+            insert(key(),
+                   clustering,
+                   cell,
+                   LivenessInfo.withExpirationTime(cell.timestamp(), cell.ttl(), cell.localDeletionTime()));
+        }
+
+        private void removeCells(Clustering<?> clustering, Iterable<Cell<?>> cells)
+        {
+            if (cells == null)
+                return;
+
+            for (Cell<?> cell : cells)
+                removeCell(clustering, cell);
+        }
+
+        private void removeCell(Clustering<?> clustering, Cell<?> cell)
+        {
+            if (cell == null || !cell.isLive(nowInSec()))
+                return;
+
+            delete(key(), clustering, cell, nowInSec());
+        }
+
+        private void indexPrimaryKey(final Clustering<?> clustering,
+                                     final LivenessInfo liveness,
+                                     final Row.Deletion deletion)
+        {
+            if (liveness.timestamp() != LivenessInfo.NO_TIMESTAMP)
+                insert(key(), clustering, null, liveness);
+
+            if (!deletion.isLive())
+                delete(key(), clustering, deletion.time());
+        }
+
+        private LivenessInfo getPrimaryKeyIndexLiveness(Row row)
+        {
+            long timestamp = row.primaryKeyLivenessInfo().timestamp();
+            int ttl = row.primaryKeyLivenessInfo().ttl();
+            for (Cell<?> cell : row.cells())
+            {
+                long cellTimestamp = cell.timestamp();
+                if (cell.isLive(nowInSec()))
+                {
+                    if (cellTimestamp > timestamp)
+                        timestamp = cellTimestamp;
+                }
+            }
+            return LivenessInfo.create(timestamp, ttl, nowInSec());
+        }
+    }
+
     public Indexer indexerFor(final DecoratedKey key,
                               final RegularAndStaticColumns columns,
                               final long nowInSec,
@@ -375,141 +531,42 @@ public abstract class CassandraIndex implements Index
         if (!isPrimaryKeyIndex() && !columns.contains(indexedColumn))
             return null;
 
-        return new Indexer()
+        return new AbstractIndexer()
         {
-            public void begin()
+            @Override
+            CassandraIndex index()
             {
+                return CassandraIndex.this;
             }
 
-            public void partitionDelete(DeletionTime deletionTime)
+            @Override
+            long nowInSec()
             {
+                return nowInSec;
             }
 
-            public void rangeTombstone(RangeTombstone tombstone)
+            @Override
+            DecoratedKey key()
             {
+                return key;
             }
 
-            public void insertRow(Row row)
+            @Override
+            void insert(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info)
             {
-                if (row.isStatic() && !indexedColumn.isStatic() && !indexedColumn.isPartitionKey())
-                    return;
-
-                if (isPrimaryKeyIndex())
-                {
-                    indexPrimaryKey(row.clustering(),
-                                    getPrimaryKeyIndexLiveness(row),
-                                    row.deletion());
-                }
-                else
-                {
-                    if (indexedColumn.isComplex())
-                        indexCells(row.clustering(), row.getComplexColumnData(indexedColumn));
-                    else
-                        indexCell(row.clustering(), row.getCell(indexedColumn));
-                }
+                CassandraIndex.this.insert(rowKey.getKey(), clustering, cell, info, ctx);
             }
 
-            public void removeRow(Row row)
+            @Override
+            void delete(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, long nowInSec)
             {
-                if (isPrimaryKeyIndex())
-                    return;
-
-                if (indexedColumn.isComplex())
-                    removeCells(row.clustering(), row.getComplexColumnData(indexedColumn));
-                else
-                    removeCell(row.clustering(), row.getCell(indexedColumn));
+                CassandraIndex.this.delete(rowKey.getKey(), clustering, cell, ctx, nowInSec);
             }
 
-            public void updateRow(Row oldRow, Row newRow)
+            @Override
+            void delete(DecoratedKey rowKey, Clustering<?> clustering, DeletionTime deletion)
             {
-                assert oldRow.isStatic() == newRow.isStatic();
-                if (newRow.isStatic() != indexedColumn.isStatic())
-                    return;
-
-                if (isPrimaryKeyIndex())
-                    indexPrimaryKey(newRow.clustering(),
-                                    getPrimaryKeyIndexLiveness(newRow),
-                                    newRow.deletion());
-
-                if (indexedColumn.isComplex())
-                {
-                    indexCells(newRow.clustering(), newRow.getComplexColumnData(indexedColumn));
-                    removeCells(oldRow.clustering(), oldRow.getComplexColumnData(indexedColumn));
-                }
-                else
-                {
-                    indexCell(newRow.clustering(), newRow.getCell(indexedColumn));
-                    removeCell(oldRow.clustering(), oldRow.getCell(indexedColumn));
-                }
-            }
-
-            public void finish()
-            {
-            }
-
-            private void indexCells(Clustering<?> clustering, Iterable<Cell<?>> cells)
-            {
-                if (cells == null)
-                    return;
-
-                for (Cell<?> cell : cells)
-                    indexCell(clustering, cell);
-            }
-
-            private void indexCell(Clustering<?> clustering, Cell<?> cell)
-            {
-                if (cell == null || !cell.isLive(nowInSec))
-                    return;
-
-                insert(key.getKey(),
-                       clustering,
-                       cell,
-                       LivenessInfo.withExpirationTime(cell.timestamp(), cell.ttl(), cell.localDeletionTime()),
-                       ctx);
-            }
-
-            private void removeCells(Clustering<?> clustering, Iterable<Cell<?>> cells)
-            {
-                if (cells == null)
-                    return;
-
-                for (Cell<?> cell : cells)
-                    removeCell(clustering, cell);
-            }
-
-            private void removeCell(Clustering<?> clustering, Cell<?> cell)
-            {
-                if (cell == null || !cell.isLive(nowInSec))
-                    return;
-
-                delete(key.getKey(), clustering, cell, ctx, nowInSec);
-            }
-
-            private void indexPrimaryKey(final Clustering<?> clustering,
-                                         final LivenessInfo liveness,
-                                         final Row.Deletion deletion)
-            {
-                if (liveness.timestamp() != LivenessInfo.NO_TIMESTAMP)
-                    insert(key.getKey(), clustering, null, liveness, ctx);
-
-                if (!deletion.isLive())
-                    delete(key.getKey(), clustering, deletion.time(), ctx);
-            }
-
-            private LivenessInfo getPrimaryKeyIndexLiveness(Row row)
-            {
-                long timestamp = row.primaryKeyLivenessInfo().timestamp();
-                int ttl = row.primaryKeyLivenessInfo().ttl();
-                for (Cell<?> cell : row.cells())
-                {
-                    long cellTimestamp = cell.timestamp();
-                    if (cell.isLive(nowInSec))
-                    {
-                        if (cellTimestamp > timestamp)
-                            timestamp = cellTimestamp;
-                    }
-                }
-                return LivenessInfo.create(timestamp, ttl, nowInSec);
+                CassandraIndex.this.delete(rowKey.getKey(), clustering, deletion, ctx);
             }
         };
     }
@@ -531,6 +588,15 @@ public abstract class CassandraIndex implements Index
         logger.trace("Removed index entry for stale value {}", indexKey);
     }
 
+    public IndexEntry createIndexEntry(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info)
+    {
+        DecoratedKey indexKey = getIndexKeyFor(getIndexedValue(rowKey.getKey(),
+                                                               clustering,
+                                                               cell));
+        Clustering<?> indexClustering = buildIndexClustering(rowKey.getKey(), clustering, cell);
+        return new IndexEntry(indexKey, indexClustering, info.timestamp(), rowKey, clustering);
+    }
+
     /**
      * Called when adding a new entry to the index
      */

src/java/org/apache/cassandra/index/internal/CassandraIndexSearcher.java

@@ -21,8 +21,11 @@
 package org.apache.cassandra.index.internal;
 
 import java.nio.ByteBuffer;
+import java.util.Optional;
 import java.util.SortedSet;
+import java.util.function.Consumer;
 
+import com.google.common.base.Preconditions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,9 +35,12 @@ import org.apache.cassandra.db.ClusteringBound;
 import org.apache.cassandra.db.ColumnFamilyStore;
 import org.apache.cassandra.db.DataRange;
 import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.LivenessInfo;
 import org.apache.cassandra.db.PartitionPosition;
 import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
+import org.apache.cassandra.db.ReadableView;
 import org.apache.cassandra.db.SinglePartitionReadCommand;
 import org.apache.cassandra.db.Slice;
 import org.apache.cassandra.db.Slices;
@@ -43,23 +49,112 @@ import org.apache.cassandra.db.filter.ClusteringIndexNamesFilter;
 import org.apache.cassandra.db.filter.ClusteringIndexSliceFilter;
 import org.apache.cassandra.db.filter.ColumnFilter;
 import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.lifecycle.SSTableSet;
+import org.apache.cassandra.db.lifecycle.View;
+import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.rows.Cell;
+import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.db.rows.RowIterator;
+import org.apache.cassandra.db.rows.Rows;
 import org.apache.cassandra.db.rows.UnfilteredRowIterator;
 import org.apache.cassandra.db.rows.UnfilteredRowIterators;
+import org.apache.cassandra.db.transform.Transformation;
 import org.apache.cassandra.dht.AbstractBounds;
 import org.apache.cassandra.index.Index;
 import org.apache.cassandra.index.internal.composites.CollectionValueIndex;
 import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.reads.tracked.PartialTrackedIndexRead;
+import org.apache.cassandra.service.reads.tracked.PartialTrackedRead;
+import org.apache.cassandra.utils.AbstractIterator;
+import org.apache.cassandra.utils.CloseablePeekingIterator;
 import org.apache.cassandra.utils.btree.BTreeSet;
 
-public abstract class CassandraIndexSearcher implements Index.Searcher
+public abstract class CassandraIndexSearcher<Match extends Index.IndexMatch> implements Index.MultiStepSearcher<Match>
 {
+    protected abstract class AbstractMatchIndexer<M extends Index.IndexMatch> extends CassandraIndex.AbstractIndexer implements MatchIndexer<M>
+    {
+        protected DecoratedKey key;
+        protected Consumer<M> indexTo;
+
+        @Override
+        long nowInSec()
+        {
+            return command.nowInSec();
+        }
+
+        @Override
+        CassandraIndex index()
+        {
+            return index;
+        }
+
+        @Override
+        DecoratedKey key()
+        {
+            return key;
+        }
+
+        protected abstract M createMatch(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info);
+
+        @Override
+        void insert(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info)
+        {
+            indexTo.accept(createMatch(rowKey, clustering, cell, info));
+        }
+
+        @Override
+        void delete(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, long nowInSec)
+        {
+
+        }
+
+        @Override
+        void delete(DecoratedKey rowKey, Clustering<?> clustering, DeletionTime deletion)
+        {
+
+        }
+
+        @Override
+        public void insertRow(Row row)
+        {
+            if (!expression.isSatisfiedBy(command.metadata(), key, row, nowInSec()))
+                return;
+
+            if (!command.selectsClustering(key, row.clustering()))
+                return;
+
+            super.insertRow(row);
+        }
+
+        @Override
+        public void index(PartitionUpdate update, Consumer<M> indexTo)
+        {
+            this.key = update.partitionKey();
+            this.indexTo = indexTo;
+
+            try
+            {
+                Row staticRow = update.staticRow();
+                if (staticRow != Rows.EMPTY_STATIC_ROW)
+                    insertRow(staticRow);
+
+                update.forEach(this::insertRow);
+            }
+            finally
+            {
+                this.key = null;
+                this.indexTo = null;
+            }
+        }
+    }
+
     private static final Logger logger = LoggerFactory.getLogger(CassandraIndexSearcher.class);
 
     private final RowFilter.Expression expression;
     protected final CassandraIndex index;
     protected final ReadCommand command;
+    protected final DecoratedKey indexedKey;
 
     public CassandraIndexSearcher(ReadCommand command,
                                   RowFilter.Expression expression,
@@ -68,6 +163,9 @@ public abstract class CassandraIndexSearcher implements Index.Searcher
         this.command = command;
         this.expression = expression;
         this.index = index;
+        Optional<ColumnFamilyStore> backingTable = index.getBackingTable();
+        Preconditions.checkState(backingTable.isPresent());
+        this.indexedKey = backingTable.get().decorateKey(expression.getIndexValue());
     }
 
     @Override
@@ -76,19 +174,96 @@ public abstract class CassandraIndexSearcher implements Index.Searcher
         return command;
     }
 
-    // of this method.
+    @Override
+    public PartialTrackedRead beginRead(ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
+    {
+        return PartialTrackedIndexRead.create(executionController, cfs, startTimeNanos, command, this);
+    }
+
+    @Override
+    public UnfilteredPartitionIterator filterCompletedRead(UnfilteredPartitionIterator iterator)
+    {
+        return Transformation.apply(iterator, new Transformation<UnfilteredRowIterator>()
+        {
+            DecoratedKey key = null;
+            @Override
+            protected DecoratedKey applyToPartitionKey(DecoratedKey key)
+            {
+                this.key = key;
+                return super.applyToPartitionKey(key);
+            }
+
+            @Override
+            protected Row applyToRow(Row row)
+            {
+                if (!expression.isSatisfiedBy(command.metadata(), key, row, command.nowInSec()))
+                    return null;
+                return row;
+            }
+        });
+    }
+
+    protected RowIterator queryIndex(DecoratedKey indexKey, ReadExecutionController executionController)
+    {
+        UnfilteredRowIterator indexIter = queryIndex(indexKey, command, executionController);
+        return UnfilteredRowIterators.filter(indexIter, command.nowInSec());
+    }
+
+    protected class ResultIterator extends AbstractIterator<UnfilteredRowIterator> implements UnfilteredPartitionIterator
+    {
+        private final CloseablePeekingIterator<Match> matchIterator;
+        private final ReadExecutionController executionController;
+
+        public ResultIterator(CloseablePeekingIterator<Match> matchIterator, ReadExecutionController executionController)
+        {
+            this.matchIterator = matchIterator;
+            this.executionController = executionController;
+        }
+
+        @Override
+        protected UnfilteredRowIterator computeNext()
+        {
+            while (matchIterator.hasNext())
+            {
+                DecoratedKey key = matchIterator.peek().key();
+                ReadableView view = index.baseCfs.select(View.select(SSTableSet.LIVE, key));
+                UnfilteredRowIterator partition = queryNextMatches(executionController, key, view, matchIterator);
+
+                if (partition == null)
+                    continue;
+
+                if (!partition.isEmpty())
+                    return partition;
+
+                partition.close();
+            }
+            return endOfData();
+        }
+
+        @Override
+        public TableMetadata metadata()
+        {
+            return command.metadata();
+        }
+
+        @Override
+        public void close()
+        {
+            matchIterator.close();
+        }
+    }
+
+    @Override
     public UnfilteredPartitionIterator search(ReadExecutionController executionController)
     {
-        // the value of the index expression is the partition key in the index table
-        DecoratedKey indexKey = index.getBackingTable().get().decorateKey(expression.getIndexValue());
-        UnfilteredRowIterator indexIter = queryIndex(indexKey, command, executionController);
+        CloseablePeekingIterator<Match> matchIterator = matchIterator(executionController);
         try
         {
-            return queryDataFromIndex(indexKey, UnfilteredRowIterators.filter(indexIter, command.nowInSec()), command, executionController);
+            return new ResultIterator(matchIterator, executionController);
         }
-        catch (RuntimeException | Error e)
+        catch (Throwable e)
         {
-            indexIter.close();
+            matchIterator.close();
             throw e;
         }
     }
@@ -217,9 +392,4 @@ public abstract class CassandraIndexSearcher implements Index.Searcher
     {
         return index.buildIndexClusteringPrefix(rowKey, clustering, null).build();
     }
-
-    protected abstract UnfilteredPartitionIterator queryDataFromIndex(DecoratedKey indexKey,
-                                                                      RowIterator indexHits,
-                                                                      ReadCommand command,
-                                                                      ReadExecutionController executionController);
 }

src/java/org/apache/cassandra/index/internal/IndexEntry.java

@@ -20,29 +20,29 @@
  */
 package org.apache.cassandra.index.internal;
 
-import java.nio.ByteBuffer;
-
 import org.apache.cassandra.db.Clustering;
 import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.schema.TableMetadata;
 
 /**
  * Entries in indexes on non-compact tables (tables with composite comparators)
  * can be encapsulated as IndexedEntry instances. These are not used when dealing
  * with indexes on static/compact tables (i.e. KEYS indexes).
  */
-public final class IndexEntry
+public final class IndexEntry implements Index.IndexMatch
 {
     public final DecoratedKey indexValue;
     public final Clustering<?> indexClustering;
     public final long timestamp;
 
-    public final ByteBuffer indexedKey;
+    public final DecoratedKey indexedKey;
     public final Clustering<?> indexedEntryClustering;
 
     public IndexEntry(DecoratedKey indexValue,
                       Clustering<?> indexClustering,
                       long timestamp,
-                      ByteBuffer indexedKey,
+                      DecoratedKey indexedKey,
                       Clustering<?> indexedEntryClustering)
     {
         this.indexValue = indexValue;
@@ -51,4 +51,41 @@ public final class IndexEntry
         this.indexedKey = indexedKey;
         this.indexedEntryClustering = indexedEntryClustering;
     }
+
+    @Override
+    public DecoratedKey key()
+    {
+        return indexedKey;
+    }
+
+    public static int compare(TableMetadata indexMetadata, TableMetadata baseMetadata, IndexEntry left, IndexEntry right)
+    {
+        int cmp = left.indexValue.compareTo(right.indexValue);
+        if (cmp != 0)
+            return cmp;
+
+        cmp = indexMetadata.comparator.compare(left.indexClustering, right.indexClustering);
+        if (cmp != 0)
+            return cmp;
+
+        cmp = left.indexedKey.compareTo(right.indexedKey);
+        if (cmp != 0)
+            return cmp;
+
+        // Take STATIC rows into account...
+        if (!left.indexedEntryClustering.isEmpty() || !right.indexedEntryClustering.isEmpty())
+        {
+            if (left.indexedEntryClustering.isEmpty())
+                return -1;
+
+            if (right.indexedEntryClustering.isEmpty())
+                return 1;
+
+            cmp = baseMetadata.comparator.compare(left.indexedEntryClustering, right.indexedEntryClustering);
+            if (cmp != 0)
+                return cmp;
+        }
+
+        return Long.compare(left.timestamp, right.timestamp);
+    }
 }

src/java/org/apache/cassandra/index/internal/composites/ClusteringColumnIndex.java

@@ -98,7 +98,7 @@ public class ClusteringColumnIndex extends CassandraIndex
         return new IndexEntry(indexedValue,
                               clustering,
                               indexEntry.primaryKeyLivenessInfo().timestamp(),
-                              clustering.bufferAt(0),
+                              baseCfs.decorateKey(clustering.bufferAt(0)),
                               builder.build());
     }
 

src/java/org/apache/cassandra/index/internal/composites/CollectionKeyIndexBase.java

@@ -87,7 +87,7 @@ public abstract class CollectionKeyIndexBase extends CassandraIndex
         return new IndexEntry(indexedValue,
                               clustering,
                               indexEntry.primaryKeyLivenessInfo().timestamp(),
-                              clustering.bufferAt(0),
+                              baseCfs.decorateKey(clustering.bufferAt(0)),
                               indexedEntryClustering);
     }
 }

src/java/org/apache/cassandra/index/internal/composites/CollectionValueIndex.java

@@ -97,7 +97,7 @@ public class CollectionValueIndex extends CassandraIndex
         return new IndexEntry(indexedValue,
                                 clustering,
                                 indexEntry.primaryKeyLivenessInfo().timestamp(),
-                                clustering.bufferAt(0),
+                                baseCfs.decorateKey(clustering.bufferAt(0)),
                                 indexedEntryClustering);
     }
 

src/java/org/apache/cassandra/index/internal/composites/CompositesSearcher.java

@@ -21,18 +21,23 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
 
+import com.google.common.base.Preconditions;
+import com.google.common.collect.PeekingIterator;
+
 import org.apache.cassandra.db.Clustering;
 import org.apache.cassandra.db.ClusteringComparator;
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.LivenessInfo;
 import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
+import org.apache.cassandra.db.ReadableView;
 import org.apache.cassandra.db.SinglePartitionReadCommand;
 import org.apache.cassandra.db.WriteContext;
 import org.apache.cassandra.db.filter.ClusteringIndexNamesFilter;
 import org.apache.cassandra.db.filter.DataLimits;
 import org.apache.cassandra.db.filter.RowFilter;
-import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.rows.Cell;
 import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.db.rows.RowIterator;
 import org.apache.cassandra.db.rows.Rows;
@@ -43,11 +48,12 @@ import org.apache.cassandra.index.Index;
 import org.apache.cassandra.index.internal.CassandraIndex;
 import org.apache.cassandra.index.internal.CassandraIndexSearcher;
 import org.apache.cassandra.index.internal.IndexEntry;
-import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.utils.AbstractIterator;
+import org.apache.cassandra.utils.CloseablePeekingIterator;
 import org.apache.cassandra.utils.btree.BTreeSet;
 
 
-public class CompositesSearcher extends CassandraIndexSearcher
+public class CompositesSearcher extends CassandraIndexSearcher<IndexEntry>
 {
     public CompositesSearcher(ReadCommand command,
                               RowFilter.Expression expression,
@@ -56,6 +62,25 @@ public class CompositesSearcher extends CassandraIndexSearcher
         super(command, expression, index);
     }
 
+    @Override
+    public MatchIndexer<IndexEntry> matchIndexer()
+    {
+        return new AbstractMatchIndexer<IndexEntry>()
+        {
+            @Override
+            protected IndexEntry createMatch(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info)
+            {
+                return index.createIndexEntry(rowKey, clustering, cell, info);
+            }
+        };
+    }
+
+    @Override
+    public MatchComparator<IndexEntry> matchComparator()
+    {
+        return (left, right, strict) -> IndexEntry.compare(index.getIndexCfs().metadata(), command.metadata(), left, right);
+    }
+
     private boolean isMatchingEntry(DecoratedKey partitionKey, IndexEntry entry, ReadCommand command)
     {
         return command.selectsKey(partitionKey) && command.selectsClustering(partitionKey, entry.indexedEntryClustering);
@@ -66,145 +91,105 @@ public class CompositesSearcher extends CassandraIndexSearcher
         return index.getIndexedColumn().isStatic();
     }
 
-    protected UnfilteredPartitionIterator queryDataFromIndex(final DecoratedKey indexKey,
-                                                             final RowIterator indexHits,
-                                                             final ReadCommand command,
-                                                             final ReadExecutionController executionController)
+    @Override
+    public CloseablePeekingIterator<IndexEntry> matchIterator(ReadExecutionController executionController)
     {
-        assert indexHits.staticRow() == Rows.EMPTY_STATIC_ROW;
-
-        return new UnfilteredPartitionIterator()
+        RowIterator indexHits = queryIndex(indexedKey, executionController);
+        try
         {
-            private IndexEntry nextEntry;
-
-            private UnfilteredRowIterator next;
-
-            public TableMetadata metadata()
+            Preconditions.checkState(indexHits.staticRow() == Rows.EMPTY_STATIC_ROW);
+            return new AbstractIterator<IndexEntry>()
             {
-                return command.metadata();
-            }
-
-            public boolean hasNext()
-            {
-                return prepareNext();
-            }
-
-            public UnfilteredRowIterator next()
-            {
-                if (next == null)
-                    prepareNext();
-
-                UnfilteredRowIterator toReturn = next;
-                next = null;
-                return toReturn;
-            }
-
-            private boolean prepareNext()
-            {
-                while (true)
+                @Override
+                protected IndexEntry computeNext()
                 {
-                    if (next != null)
-                        return true;
-
-                    if (nextEntry == null)
+                    while (indexHits.hasNext())
                     {
-                        if (!indexHits.hasNext())
-                            return false;
-
-                        nextEntry = index.decodeEntry(indexKey, indexHits.next());
-                    }
-
-                    SinglePartitionReadCommand dataCmd;
-                    DecoratedKey partitionKey = index.baseCfs.decorateKey(nextEntry.indexedKey);
-                    List<IndexEntry> entries = new ArrayList<>();
-                    if (isStaticColumn())
-                    {
-                        // The index hit may not match the commad key constraint
-                        if (!isMatchingEntry(partitionKey, nextEntry, command)) {
-                            nextEntry = indexHits.hasNext() ? index.decodeEntry(indexKey, indexHits.next()) : null;
-                            continue;
-                        }
-
-                        // If the index is on a static column, we just need to do a full read on the partition.
-                        // Note that we want to re-use the command.columnFilter() in case of future change.
-                        dataCmd = SinglePartitionReadCommand.create(index.baseCfs.metadata(),
-                                                                    command.nowInSec(),
-                                                                    command.columnFilter(),
-                                                                    RowFilter.none(),
-                                                                    DataLimits.NONE,
-                                                                    partitionKey,
-                                                                    command.clusteringIndexFilter(partitionKey));
-                        entries.add(nextEntry);
-                        nextEntry = indexHits.hasNext() ? index.decodeEntry(indexKey, indexHits.next()) : null;
-                    }
-                    else
-                    {
-                        // Gather all index hits belonging to the same partition and query the data for those hits.
-                        // TODO: it's much more efficient to do 1 read for all hits to the same partition than doing
-                        // 1 read per index hit. However, this basically mean materializing all hits for a partition
-                        // in memory so we should consider adding some paging mechanism. However, index hits should
-                        // be relatively small so it's much better than the previous code that was materializing all
-                        // *data* for a given partition.
-                        BTreeSet.Builder<Clustering<?>> clusterings = BTreeSet.builder(index.baseCfs.getComparator());
-                        while (nextEntry != null && partitionKey.getKey().equals(nextEntry.indexedKey))
-                        {
-                            // We're queried a slice of the index, but some hits may not match some of the clustering column constraints
-                            if (isMatchingEntry(partitionKey, nextEntry, command))
-                            {
-                                clusterings.add(nextEntry.indexedEntryClustering);
-                                entries.add(nextEntry);
-                            }
-
-                            nextEntry = indexHits.hasNext() ? index.decodeEntry(indexKey, indexHits.next()) : null;
-                        }
-
-                        // Because we've eliminated entries that don't match the clustering columns, it's possible we added nothing
-                        if (clusterings.isEmpty())
+                        IndexEntry nextEntry = index.decodeEntry(indexedKey, indexHits.next());
+                        DecoratedKey partitionKey = nextEntry.indexedKey;
+                        if (!isMatchingEntry(partitionKey, nextEntry, command))
                             continue;
 
-                        // Query the gathered index hits. We still need to filter stale hits from the resulting query.
-                        ClusteringIndexNamesFilter filter = new ClusteringIndexNamesFilter(clusterings.build(), false);
-                        dataCmd = SinglePartitionReadCommand.create(index.baseCfs.metadata(),
-                                                                    command.nowInSec(),
-                                                                    command.columnFilter(),
-                                                                    command.rowFilter(),
-                                                                    DataLimits.NONE,
-                                                                    partitionKey,
-                                                                    filter,
-                                                                    (Index.QueryPlan) null);
+                        return nextEntry;
                     }
-
-                    // by the next caller of next, or through closing this iterator is this come before.
-                    UnfilteredRowIterator dataIter =
-                        filterStaleEntries(dataCmd.queryMemtableAndDisk(index.baseCfs, executionController),
-                                           indexKey.getKey(),
-                                           entries,
-                                           executionController.getWriteContext(),
-                                           command.nowInSec());
-
-                    if (dataIter.isEmpty())
-                    {
-                        dataIter.close();
-                        continue;
-                    }
-
-                    next = dataIter;
-                    return true;
+                    return endOfData();
                 }
-            }
 
-            public void remove()
-            {
-                throw new UnsupportedOperationException();
-            }
-
-            public void close()
-            {
+                @Override
+                public void close()
+                {
+                    if (indexHits != null)
+                        indexHits.close();
+                }
+            };
+        }
+        catch (Throwable e)
+        {
+            if (indexHits != null)
                 indexHits.close();
-                if (next != null)
-                    next.close();
+            throw e;
+        }
+    }
+
+    @Override
+    public UnfilteredRowIterator queryNextMatches(ReadExecutionController executionController, DecoratedKey partitionKey, ReadableView view, PeekingIterator<IndexEntry> matches)
+    {
+        Preconditions.checkArgument(matches.hasNext());
+        SinglePartitionReadCommand dataCmd;
+        List<IndexEntry> entries = new ArrayList<>();
+        if (isStaticColumn())
+        {
+
+            // If the index is on a static column, we just need to do a full read on the partition.
+            // Note that we want to re-use the command.columnFilter() in case of future change.
+            dataCmd = SinglePartitionReadCommand.create(index.baseCfs.metadata(),
+                                                        command.nowInSec(),
+                                                        command.columnFilter(),
+                                                        RowFilter.none(),
+                                                        DataLimits.NONE,
+                                                        partitionKey,
+                                                        command.clusteringIndexFilter(partitionKey));
+            entries.add(matches.next());
+        }
+        else
+        {
+            // Gather all index hits belonging to the same partition and query the data for those hits.
+            // TODO: it's much more efficient to do 1 read for all hits to the same partition than doing
+            // 1 read per index hit. However, this basically mean materializing all hits for a partition
+            // in memory so we should consider adding some paging mechanism. However, index hits should
+            // be relatively small so it's much better than the previous code that was materializing all
+            // *data* for a given partition.
+            BTreeSet.Builder<Clustering<?>> clusterings = BTreeSet.builder(index.baseCfs.getComparator());
+            while (matches.hasNext() && partitionKey.equals(matches.peek().indexedKey))
+            {
+                // We're queried a slice of the index, and some hits may not match some of the clustering column constraints,
+                // but they will have been filtered out upstream
+                IndexEntry nextEntry = matches.next();
+                clusterings.add(nextEntry.indexedEntryClustering);
+                entries.add(nextEntry);
             }
-        };
+
+            // since non-matching entries will have been filtered out by matchIterator, it should not be possible to have empty clusterings
+            Preconditions.checkArgument(!clusterings.isEmpty());
+
+            // Query the gathered index hits. We still need to filter stale hits from the resulting query.
+            ClusteringIndexNamesFilter filter = new ClusteringIndexNamesFilter(clusterings.build(), false);
+            dataCmd = SinglePartitionReadCommand.create(index.baseCfs.metadata(),
+                                                        command.nowInSec(),
+                                                        command.columnFilter(),
+                                                        command.rowFilter(),
+                                                        DataLimits.NONE,
+                                                        partitionKey,
+                                                        filter,
+                                                        (Index.QueryPlan) null);
+        }
+
+        // by the next caller of next, or through closing this iterator is this come before.
+        return filterStaleEntries(dataCmd.queryMemtableAndDisk(view, index.baseCfs, executionController),
+                                  indexedKey.getKey(),
+                                  entries,
+                                  executionController.getWriteContext(),
+                                  command.nowInSec());
     }
 
     private void deleteAllEntries(final List<IndexEntry> entries, final WriteContext ctx, final long nowInSec)
@@ -297,8 +282,8 @@ public class CompositesSearcher extends CassandraIndexSearcher
                         // those tables do not support static columns. By consequence if a table
                         // has some static columns and all its clustering key elements are null
                         // it means that the partition exists and contains only static data
-                       if (!dataIter.metadata().hasStaticColumns() || !containsOnlyNullValues(indexedEntryClustering))
-                           staleEntries.add(entry);
+                        if (!dataIter.metadata().hasStaticColumns() || !containsOnlyNullValues(indexedEntryClustering))
+                            staleEntries.add(entry);
                     }
                     // entries correspond to the rows we've queried, so we shouldn't have a row that has no corresponding entry.
                     throw new AssertionError();

src/java/org/apache/cassandra/index/internal/composites/PartitionKeyIndex.java

@@ -90,7 +90,7 @@ public class PartitionKeyIndex extends CassandraIndex
         return new IndexEntry(indexedValue,
                               clustering,
                               indexEntry.primaryKeyLivenessInfo().timestamp(),
-                              clustering.bufferAt(0),
+                              baseCfs.decorateKey(clustering.bufferAt(0)),
                               builder.build());
     }
 

src/java/org/apache/cassandra/index/internal/composites/RegularColumnIndex.java

@@ -101,7 +101,7 @@ public class RegularColumnIndex extends CassandraIndex
         return new IndexEntry(indexedValue,
                               clustering,
                               indexEntry.primaryKeyLivenessInfo().timestamp(),
-                              clustering.bufferAt(0),
+                              baseCfs.decorateKey(clustering.bufferAt(0)),
                               indexedEntryClustering);
     }
 

src/java/org/apache/cassandra/index/internal/keys/KeysSearcher.java

@@ -19,33 +19,34 @@ package org.apache.cassandra.index.internal.keys;
 
 import java.nio.ByteBuffer;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.PeekingIterator;
 
 import org.apache.cassandra.db.Clustering;
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.LivenessInfo;
 import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
+import org.apache.cassandra.db.ReadableView;
 import org.apache.cassandra.db.SinglePartitionReadCommand;
 import org.apache.cassandra.db.WriteContext;
 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.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.rows.Cell;
 import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.db.rows.RowIterator;
 import org.apache.cassandra.db.rows.Rows;
 import org.apache.cassandra.db.rows.UnfilteredRowIterator;
-import org.apache.cassandra.index.Index;
 import org.apache.cassandra.index.internal.CassandraIndex;
 import org.apache.cassandra.index.internal.CassandraIndexSearcher;
-import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.index.internal.IndexEntry;
+import org.apache.cassandra.utils.AbstractIterator;
+import org.apache.cassandra.utils.CloseablePeekingIterator;
 
-public class KeysSearcher extends CassandraIndexSearcher
+public class KeysSearcher extends CassandraIndexSearcher<IndexEntry>
 {
-    private static final Logger logger = LoggerFactory.getLogger(KeysSearcher.class);
-
     public KeysSearcher(ReadCommand command,
                         RowFilter.Expression expression,
                         CassandraIndex indexer)
@@ -53,89 +54,91 @@ public class KeysSearcher extends CassandraIndexSearcher
         super(command, expression, indexer);
     }
 
-    protected UnfilteredPartitionIterator queryDataFromIndex(final DecoratedKey indexKey,
-                                                             final RowIterator indexHits,
-                                                             final ReadCommand command,
-                                                             final ReadExecutionController executionController)
+    @Override
+    public MatchIndexer<IndexEntry> matchIndexer()
     {
-        assert indexHits.staticRow() == Rows.EMPTY_STATIC_ROW;
-
-        return new UnfilteredPartitionIterator()
+        return new AbstractMatchIndexer<IndexEntry>()
         {
-            private UnfilteredRowIterator next;
-
-            public TableMetadata metadata()
+            @Override
+            protected IndexEntry createMatch(DecoratedKey rowKey, Clustering<?> clustering, Cell<?> cell, LivenessInfo info)
             {
-                return command.metadata();
-            }
-
-            public boolean hasNext()
-            {
-                return prepareNext();
-            }
-
-            public UnfilteredRowIterator next()
-            {
-                if (next == null)
-                    prepareNext();
-
-                UnfilteredRowIterator toReturn = next;
-                next = null;
-                return toReturn;
-            }
-
-            private boolean prepareNext()
-            {
-                while (next == null && indexHits.hasNext())
-                {
-                    Row hit = indexHits.next();
-                    DecoratedKey key = index.baseCfs.decorateKey(hit.clustering().bufferAt(0));
-                    if (!command.selectsKey(key))
-                        continue;
-
-                    ColumnFilter extendedFilter = getExtendedFilter(command.columnFilter());
-                    SinglePartitionReadCommand dataCmd = SinglePartitionReadCommand.create(index.baseCfs.metadata(),
-                                                                                           command.nowInSec(),
-                                                                                           extendedFilter,
-                                                                                           command.rowFilter(),
-                                                                                           DataLimits.NONE,
-                                                                                           key,
-                                                                                           command.clusteringIndexFilter(key),
-                                                                                           (Index.QueryPlan) null);
-
-                                                  // Otherwise, we close right away if empty, and if it's assigned to next it will be called either
-                                                  // by the next caller of next, or through closing this iterator is this come before.
-                    UnfilteredRowIterator dataIter = filterIfStale(dataCmd.queryMemtableAndDisk(index.baseCfs, executionController),
-                                                                   hit,
-                                                                   indexKey.getKey(),
-                                                                   executionController.getWriteContext(),
-                                                                   command.nowInSec());
-
-                    if (dataIter != null)
-                    {
-                        if (dataIter.isEmpty())
-                            dataIter.close();
-                        else
-                            next = dataIter;
-                    }
-                }
-                return next != null;
-            }
-
-            public void remove()
-            {
-                throw new UnsupportedOperationException();
-            }
-
-            public void close()
-            {
-                indexHits.close();
-                if (next != null)
-                    next.close();
+                return index.createIndexEntry(rowKey, clustering, cell, info);
             }
         };
     }
 
+    @Override
+    public MatchComparator<IndexEntry> matchComparator()
+    {
+        return (left, right, strict) -> IndexEntry.compare(index.getIndexCfs().metadata(), command.metadata(), left, right);
+    }
+
+    @Override
+    public CloseablePeekingIterator<IndexEntry> matchIterator(ReadExecutionController executionController)
+    {
+        RowIterator indexHits = queryIndex(indexedKey, executionController);
+        try
+        {
+            Preconditions.checkState(indexHits.staticRow() == Rows.EMPTY_STATIC_ROW);
+            return new AbstractIterator<IndexEntry>()
+            {
+                @Override
+                protected IndexEntry computeNext()
+                {
+                    while (indexHits.hasNext())
+                    {
+                        Row hit = indexHits.next();
+                        DecoratedKey key = index.baseCfs.decorateKey(hit.clustering().bufferAt(0));
+                        if (!command.selectsKey(key))
+                            continue;
+
+                        return new IndexEntry(indexedKey, hit.clustering(), hit.primaryKeyLivenessInfo().timestamp(), key, Clustering.EMPTY);
+                    }
+                    return endOfData();
+                }
+
+                @Override
+                public void close()
+                {
+                    if (indexHits != null)
+                        indexHits.close();
+                }
+            };
+
+        }
+        catch (Throwable e)
+        {
+            if (indexHits != null)
+                indexHits.close();
+            throw e;
+        }
+    }
+
+    @Override
+    public UnfilteredRowIterator queryNextMatches(ReadExecutionController executionController, DecoratedKey key, ReadableView view, PeekingIterator<IndexEntry> matches)
+    {
+        Preconditions.checkArgument(matches.hasNext());
+
+        IndexEntry entry = matches.next();
+
+        ColumnFilter extendedFilter = getExtendedFilter(command.columnFilter());
+        SinglePartitionReadCommand dataCmd = SinglePartitionReadCommand.create(index.baseCfs.metadata(),
+                                                                               command.nowInSec(),
+                                                                               extendedFilter,
+                                                                               command.rowFilter(),
+                                                                               DataLimits.NONE,
+                                                                               key,
+                                                                               command.clusteringIndexFilter(key));
+
+        // Otherwise, we close right away if empty, and if it's assigned to next it will be called either
+        // by the next caller of next, or through closing this iterator is this come before.
+        return filterIfStale(dataCmd.queryMemtableAndDisk(index.baseCfs, executionController),
+                             entry.timestamp,
+                             indexedKey.getKey(),
+                             executionController.getWriteContext(),
+                             command.nowInSec());
+    }
+
     private ColumnFilter getExtendedFilter(ColumnFilter initialFilter)
     {
         if (command.columnFilter().fetches(index.getIndexedColumn()))
@@ -148,7 +151,7 @@ public class KeysSearcher extends CassandraIndexSearcher
     }
 
     private UnfilteredRowIterator filterIfStale(UnfilteredRowIterator iterator,
-                                                Row indexHit,
+                                                long timestamp,
                                                 ByteBuffer indexedValue,
                                                 WriteContext ctx,
                                                 long nowInSec)
@@ -159,7 +162,7 @@ public class KeysSearcher extends CassandraIndexSearcher
             // Index is stale, remove the index entry and ignore
             index.deleteStaleEntry(index.getIndexCfs().decorateKey(indexedValue),
                                    makeIndexClustering(iterator.partitionKey().getKey(), Clustering.EMPTY),
-                                   DeletionTime.build(indexHit.primaryKeyLivenessInfo().timestamp(), nowInSec),
+                                   DeletionTime.build(timestamp, nowInSec),
                                    ctx);
             iterator.close();
             return null;

src/java/org/apache/cassandra/index/sai/plan/FilterTree.java

@@ -79,7 +79,12 @@ public class FilterTree
 
     public boolean isSatisfiedBy(DecoratedKey key, Row row, Row staticRow)
     {
-        boolean result = localSatisfiedBy(key, row, staticRow);
+        return isSatisfiedBy(key, row, staticRow, context.hasUnrepairedMatches);
+    }
+
+    public boolean isSatisfiedBy(DecoratedKey key, Row row, Row staticRow, boolean hasUnrepairedMatches)
+    {
+        boolean result = localSatisfiedBy(key, row, staticRow, hasUnrepairedMatches);
 
         for (FilterTree child : children)
             result = baseOperator.apply(result, child.isSatisfiedBy(key, row, staticRow));
@@ -87,19 +92,19 @@ public class FilterTree
         return result;
     }
 
-    private boolean localSatisfiedBy(DecoratedKey key, Row row, Row staticRow)
+    private boolean localSatisfiedBy(DecoratedKey key, Row row, Row staticRow, boolean hasUnrepairedMatches)
     {
         if (row == null)
             return false;
 
         final long now = FBUtilities.nowInSeconds();
         // Downgrade AND to OR unless the coordinator indicates strict filtering is safe or all matches are repaired:
-        BooleanOperator localOperator = (isStrict || !context.hasUnrepairedMatches) ? baseOperator : BooleanOperator.OR;
+        BooleanOperator localOperator = (isStrict || !hasUnrepairedMatches) ? baseOperator : BooleanOperator.OR;
         boolean result = localOperator == BooleanOperator.AND;
 
         // If all matches on indexed columns are repaired, strict filtering is not allowed, and there are multiple
         // unindexed column expressions, isolate the expressions on unindexed columns and union their results:
-        boolean isolateUnindexed = !context.hasUnrepairedMatches && !isStrict && expressions.hasMultipleUnindexedColumns();
+        boolean isolateUnindexed = !hasUnrepairedMatches && !isStrict && expressions.hasMultipleUnindexedColumns();
         boolean unindexedResult = false;
 
         Iterator<ColumnMetadata> columnIterator = expressions.columns().iterator();

src/java/org/apache/cassandra/index/sai/plan/QueryController.java

@@ -42,6 +42,7 @@ import org.apache.cassandra.db.PartitionPosition;
 import org.apache.cassandra.db.PartitionRangeReadCommand;
 import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
+import org.apache.cassandra.db.ReadableView;
 import org.apache.cassandra.db.SinglePartitionReadCommand;
 import org.apache.cassandra.db.filter.ClusteringIndexFilter;
 import org.apache.cassandra.db.filter.ClusteringIndexNamesFilter;
@@ -127,6 +128,11 @@ public class QueryController
         this.lastPrimaryKey = keyFactory.create(mergeRange.right.getToken());
         this.nextClusterings = new InsertionOrderedNavigableSet<>(cfs.metadata().comparator);
     }
+    
+    public ReadCommand command()
+    {
+        return command;
+    }
 
     public PrimaryKey.Factory primaryKeyFactory()
     {
@@ -183,7 +189,7 @@ public class QueryController
         return index != null && index.hasAnalyzer();
     }
 
-    public UnfilteredRowIterator queryStorage(List<PrimaryKey> keys, ReadExecutionController executionController)
+    public UnfilteredRowIterator queryStorage(ReadableView view, List<PrimaryKey> keys, ReadExecutionController executionController)
     {
         if (keys.isEmpty())
             throw new IllegalArgumentException("At least one primary key is required!");
@@ -196,7 +202,7 @@ public class QueryController
                                                                                  keys.get(0).partitionKey(),
                                                                                  makeFilter(keys));
 
-        return partition.queryMemtableAndDisk(cfs, executionController);
+        return partition.queryMemtableAndDisk(view, cfs, executionController);
     }
 
     private static Runnable getIndexReleaser(Set<SSTableIndex> referencedIndexes)
@@ -237,7 +243,7 @@ public class QueryController
                                                                                  key.partitionKey(),
                                                                                  makeFilter(List.of(key)));
 
-        return partition.queryMemtableAndDisk(cfs, view, SOURCE_TABLE_ROW_TRANSFORMER, executionController);
+        return partition.queryMemtableAndDisk(view, cfs, SOURCE_TABLE_ROW_TRANSFORMER, executionController);
     }
 
     /**

src/java/org/apache/cassandra/index/sai/utils/PrimaryKey.java

@@ -29,6 +29,7 @@ import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.marshal.ByteBufferAccessor;
 import org.apache.cassandra.dht.IPartitioner;
 import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.index.Index;
 import org.apache.cassandra.utils.ByteBufferUtil;
 import org.apache.cassandra.utils.bytecomparable.ByteComparable;
 import org.apache.cassandra.utils.bytecomparable.ByteSource;
@@ -40,7 +41,7 @@ import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;
  * The {@link Factory.TokenOnlyPrimaryKey} is used by the {@link org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher} to
  * position the search within the query range.
  */
-public interface PrimaryKey extends Comparable<PrimaryKey>, ByteComparable
+public interface PrimaryKey extends Comparable<PrimaryKey>, ByteComparable, Index.IndexMatch
 {
     /**
      * See the javadoc for {@link #kind()} for how this enum is used.
@@ -474,6 +475,12 @@ public interface PrimaryKey extends Comparable<PrimaryKey>, ByteComparable
      */
     DecoratedKey partitionKey();
 
+    @Override
+    default DecoratedKey key()
+    {
+        return partitionKey();
+    }
+
     /**
      * Returns the {@link Clustering} representing the clustering component of the {@link PrimaryKey}.
      * <p>

src/java/org/apache/cassandra/replication/BroadcastLogOffsets.java

@@ -21,9 +21,6 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import org.apache.cassandra.db.TypeSizes;
 import org.apache.cassandra.dht.AbstractBounds;
 import org.apache.cassandra.dht.IPartitioner;
@@ -36,8 +33,6 @@ import org.apache.cassandra.net.IVerbHandler;
 
 public class BroadcastLogOffsets
 {
-    private static final Logger logger = LoggerFactory.getLogger(BroadcastLogOffsets.class);
-
     private final String keyspace;
     private final Range<Token> range;
     private final List<Offsets.Immutable> replicatedOffsets;
@@ -59,7 +54,7 @@ public class BroadcastLogOffsets
     @Override
     public String toString()
     {
-        StringBuilder sb = new StringBuilder('[');
+        StringBuilder sb = new StringBuilder("[");
         boolean isFirst = true;
         for (Offsets.Immutable logOffsets : replicatedOffsets)
         {
@@ -73,7 +68,6 @@ public class BroadcastLogOffsets
 
     public static final IVerbHandler<BroadcastLogOffsets> verbHandler = message -> {
         BroadcastLogOffsets replicatedOffsets = message.payload;
-        logger.trace("Received replicated offsets {} from {}", replicatedOffsets, message.from());
         MutationTrackingService.instance.updateReplicatedOffsets(replicatedOffsets.keyspace,
                                                                  replicatedOffsets.range,
                                                                  replicatedOffsets.replicatedOffsets,

src/java/org/apache/cassandra/schema/ReplicationType.java

@@ -83,10 +83,4 @@ public enum ReplicationType
     {
         return this == tracked;
     }
-
-    // FIXME: used in lieu of adding support for tracked reads in parameterized tests, fix usages of this method
-    public static ReplicationType[] fixmeValues()
-    {
-        return new ReplicationType[]{ untracked };
-    }
 }

src/java/org/apache/cassandra/service/reads/tracked/AbstractPartialTrackedRead.java

@@ -1,266 +0,0 @@
-/*
- * 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.service.reads.tracked;
-
-import com.google.common.base.Preconditions;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.cassandra.db.ColumnFamilyStore;
-import org.apache.cassandra.db.Mutation;
-import org.apache.cassandra.db.ReadExecutionController;
-import org.apache.cassandra.db.partitions.PartitionUpdate;
-import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
-import org.apache.cassandra.db.transform.RTBoundValidator;
-
-
-public abstract class AbstractPartialTrackedRead implements PartialTrackedRead
-{
-    private static final Logger logger = LoggerFactory.getLogger(AbstractPartialTrackedRead.class);
-
-    protected interface Augmentable
-    {
-        State augment(PartitionUpdate update);
-    }
-
-    protected static abstract class State
-    {
-        protected static final State CLOSED = new State()
-        {
-            @Override
-            String name()
-            {
-                return "closed";
-            }
-
-            @Override
-            boolean isClosed()
-            {
-                return true;
-            }
-        };
-
-        abstract String name();
-
-        boolean isInitialized()
-        {
-            return false;
-        }
-
-        Initialized asInitialized()
-        {
-            throw new IllegalStateException("State is " + name() + ", not " + Initialized.NAME);
-        }
-
-        boolean isPrepared()
-        {
-            return false;
-        }
-
-        Prepared asPrepared()
-        {
-            throw new IllegalStateException("State is " + name() + ", not " + Prepared.NAME);
-        }
-
-        boolean isCompleted()
-        {
-            return false;
-        }
-
-        Completed asCompleted()
-        {
-            throw new IllegalStateException("State is " + name() + ", not " + Completed.NAME);
-        }
-
-        boolean isAugmentable()
-        {
-            return isPrepared() || isInitialized();
-        }
-
-        Augmentable asAugmentable()
-        {
-            if (isPrepared()) return asPrepared();
-            throw new IllegalStateException("State is " + name() + ", not augmentable");
-        }
-
-        boolean isClosed()
-        {
-            return false;
-        }
-
-        void close()
-        {
-        }
-    }
-
-    // TODO (expected): this is a redundant state, never exposed
-    protected final class Initialized extends State
-    {
-        static final String NAME = "initialized";
-
-        @Override
-        String name()
-        {
-            return NAME;
-        }
-
-        @Override
-        boolean isInitialized()
-        {
-            return true;
-        }
-
-        @Override
-        Initialized asInitialized()
-        {
-            return this;
-        }
-
-        Prepared prepare(UnfilteredPartitionIterator initialData)
-        {
-            return prepareInternal(initialData);
-        }
-    }
-
-    protected abstract Prepared prepareInternal(UnfilteredPartitionIterator initialData);
-
-    protected abstract class Prepared extends State implements Augmentable
-    {
-        private static final String NAME = "prepared";
-
-        @Override
-        String name()
-        {
-            return NAME;
-        }
-
-        @Override
-        boolean isPrepared()
-        {
-            return true;
-        }
-
-        @Override
-        Prepared asPrepared()
-        {
-            return this;
-        }
-
-        abstract Completed complete();
-
-    }
-
-    protected abstract class Completed extends State
-    {
-        private static final String NAME = "completed";
-
-        @Override
-        String name()
-        {
-            return NAME;
-        }
-
-        protected abstract UnfilteredPartitionIterator iterator();
-        protected abstract CompletedRead createResult(UnfilteredPartitionIterator iterator);
-
-        protected CompletedRead getResult()
-        {
-            UnfilteredPartitionIterator result = command().completeTrackedRead(iterator(), AbstractPartialTrackedRead.this);
-            // validate that the sequence of RT markers is correct: open is followed by close, deletion times for both
-            // ends equal, and there are no dangling RT bound in any partition.
-            result = RTBoundValidator.validate(result, RTBoundValidator.Stage.PROCESSED, true);
-            return createResult(result);
-        }
-    }
-
-    final ReadExecutionController executionController;
-    final ColumnFamilyStore cfs;
-    final long startTimeNanos;
-    private State state = new Initialized();
-
-    public AbstractPartialTrackedRead(ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
-    {
-        this.executionController = executionController;
-        this.cfs = cfs;
-        this.startTimeNanos = startTimeNanos;
-    }
-
-    @Override
-    public ReadExecutionController executionController()
-    {
-        return executionController;
-    }
-
-    @Override
-    public ColumnFamilyStore cfs()
-    {
-        return cfs;
-    }
-
-    @Override
-    public long startTimeNanos()
-    {
-        return startTimeNanos;
-    }
-
-    protected synchronized State state()
-    {
-        return state;
-    }
-
-    /**
-     * Implementors need to call this before returning this from createInProgressRead
-     * TODO (expected): this is a redundant transition from a redundant state (INITIALIZED)
-     */
-    synchronized void prepare(UnfilteredPartitionIterator initialData)
-    {
-        logger.trace("Preparing read {}", this);
-        state = state.asInitialized().prepare(initialData);
-    }
-
-    @Override
-    public synchronized void augment(Mutation mutation)
-    {
-        PartitionUpdate update = mutation.getPartitionUpdate(command().metadata());
-        if (update != null)
-            state = state.asAugmentable().augment(update);
-    }
-
-    @Override
-    public synchronized CompletedRead complete()
-    {
-        Preconditions.checkState(state.isPrepared());
-        Completed completed = state.asPrepared().complete();
-        state = completed;
-        return completed.getResult();
-    }
-
-    @Override
-    public synchronized void close()
-    {
-        if (state.isClosed())
-            return;
-
-        logger.trace("Closing read {}", this);
-        state.close();
-        executionController.close();
-        state = State.CLOSED;
-    }
-}

src/java/org/apache/cassandra/service/reads/tracked/ExtendingCompletedRead.java

@@ -18,7 +18,7 @@
 
 package org.apache.cassandra.service.reads.tracked;
 
-import org.apache.cassandra.transport.Dispatcher;
+import com.google.common.base.Preconditions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -35,48 +35,31 @@ import org.apache.cassandra.db.transform.EmptyPartitionsDiscarder;
 import org.apache.cassandra.db.transform.Transformation;
 import org.apache.cassandra.dht.AbstractBounds;
 import org.apache.cassandra.tracing.Tracing;
+import org.apache.cassandra.transport.Dispatcher;
 import org.apache.cassandra.utils.FBUtilities;
 import org.apache.cassandra.utils.concurrent.AsyncPromise;
 import org.apache.cassandra.utils.concurrent.Future;
 
-class ExtendingCompletedRead implements PartialTrackedRead.CompletedRead
+public abstract class ExtendingCompletedRead implements PartialTrackedRead.CompletedRead
 {
     private static final Logger logger = LoggerFactory.getLogger(ExtendingCompletedRead.class);
 
-    final PartitionRangeReadCommand command;
-    final UnfilteredPartitionIterator iterator;
     // merged end-result counter
     final DataLimits.Counter mergedResultCounter;
-
     private final boolean partitionsFetched;
     private final boolean initialIteratorExhausted;
-    protected final AbstractBounds<PartitionPosition> followUpBounds;
 
-    public ExtendingCompletedRead(PartitionRangeReadCommand command,
-                                  UnfilteredPartitionIterator iterator,
-                                  boolean partitionsFetched,
-                                  boolean initialIteratorExhausted,
-                                  AbstractBounds<PartitionPosition> followUpBounds)
+    public ExtendingCompletedRead(ReadCommand command, boolean partitionsFetched, boolean initialIteratorExhausted)
     {
-        this.command = command;
-        this.iterator = iterator;
-        mergedResultCounter = command.limits().newCounter(command.nowInSec(),
-                                                          true,
-                                                          command.selectsFullPartition(),
-                                                          command.metadata().enforceStrictLiveness());
+        this.mergedResultCounter = command.limits().newCounter(command.nowInSec(),
+                                                               true,
+                                                               command.selectsFullPartition(),
+                                                               command.metadata().enforceStrictLiveness());
         this.partitionsFetched = partitionsFetched;
         this.initialIteratorExhausted = initialIteratorExhausted;
-        this.followUpBounds = followUpBounds;
     }
 
-    @Override
-    public TrackedDataResponse response()
-    {
-        PartitionIterator filtered = UnfilteredPartitionIterators.filter(iterator, command.nowInSec());
-        PartitionIterator counted = Transformation.apply(filtered, mergedResultCounter);
-        PartitionIterator result = Transformation.apply(counted, new EmptyPartitionsDiscarder());
-        return TrackedDataResponse.create(result, command.columnFilter());
-    }
+    abstract ReadCommand command();
 
     static boolean followUpReadRequired(ReadCommand command, DataLimits.Counter mergedResultCounter, boolean initialIteratorExhausted, boolean partitionsFetched)
     {
@@ -116,7 +99,7 @@ class ExtendingCompletedRead implements PartialTrackedRead.CompletedRead
 
     protected boolean followUpRequired()
     {
-        return followUpReadRequired(command, mergedResultCounter, initialIteratorExhausted, partitionsFetched);
+        return followUpReadRequired(command(), mergedResultCounter, initialIteratorExhausted, partitionsFetched);
     }
 
     static int toQuery(ReadCommand command, DataLimits.Counter mergedResultCounter)
@@ -145,18 +128,21 @@ class ExtendingCompletedRead implements PartialTrackedRead.CompletedRead
          * the total # of rows remaining - if it has some. If we don't grab enough rows in some of the partitions,
          * then future ShortReadRowsProtection.moreContents() calls will fetch the missing ones.
          */
-        int toQuery = toQuery(command, mergedResultCounter);
+        int toQuery = toQuery(command(), mergedResultCounter);
 
-        ColumnFamilyStore.metricsFor(command.metadata().id).shortReadProtectionRequests.mark();
+        ColumnFamilyStore.metricsFor(command().metadata().id).shortReadProtectionRequests.mark();
         Tracing.trace("Requesting {} extra rows from {} for short read protection", toQuery, FBUtilities.getBroadcastAddressAndPort());
         logger.info("Requesting {} extra rows from {} for short read protection", toQuery, FBUtilities.getBroadcastAddressAndPort());
 
         return makeFollowupRead(initialResponse, toQuery, consistencyLevel, requestTime);
     }
 
+    protected abstract AbstractBounds<PartitionPosition> followUpBounds();
+
     protected Future<TrackedDataResponse> makeFollowupRead(TrackedDataResponse initialResponse, int toQuery, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime)
     {
-        TrackedRead.Range followUpRead = PartialTrackedRangeRead.makeFollowUpRead(command, followUpBounds, toQuery, consistencyLevel, requestTime);
+        Preconditions.checkState(command() instanceof PartitionRangeReadCommand);
+        TrackedRead.Range followUpRead = PartialTrackedRangeRead.makeFollowUpRead((PartitionRangeReadCommand) command(), followUpBounds(), toQuery, consistencyLevel, requestTime);
         logger.trace("Short read detected, starting followup read {}", followUpRead);
         followUpRead.start(requestTime);
         AsyncPromise<TrackedDataResponse> combinedRead = new AsyncPromise<>();
@@ -180,9 +166,49 @@ class ExtendingCompletedRead implements PartialTrackedRead.CompletedRead
         return combinedRead;
     }
 
-    @Override
-    public void close()
+    static class RangeRead extends ExtendingCompletedRead
     {
-        iterator.close();
+        final PartitionRangeReadCommand command;
+        final UnfilteredPartitionIterator iterator;
+        protected final AbstractBounds<PartitionPosition> followUpBounds;
+
+        public RangeRead(PartitionRangeReadCommand command,
+                         UnfilteredPartitionIterator iterator,
+                         boolean partitionsFetched,
+                         boolean initialIteratorExhausted,
+                         AbstractBounds<PartitionPosition> followUpBounds)
+        {
+            super(command, partitionsFetched, initialIteratorExhausted);
+            this.command = command;
+            this.iterator = iterator;
+            this.followUpBounds = followUpBounds;
+        }
+
+        @Override
+        ReadCommand command()
+        {
+            return command;
+        }
+
+        @Override
+        protected AbstractBounds<PartitionPosition> followUpBounds()
+        {
+            return followUpBounds;
+        }
+
+        @Override
+        public TrackedDataResponse response()
+        {
+            PartitionIterator filtered = UnfilteredPartitionIterators.filter(iterator, command.nowInSec());
+            PartitionIterator counted = Transformation.apply(filtered, mergedResultCounter);
+            PartitionIterator result = Transformation.apply(counted, new EmptyPartitionsDiscarder());
+            return TrackedDataResponse.create(result, command.columnFilter());
+        }
+
+        @Override
+        public void close()
+        {
+            iterator.close();
+        }
     }
 }

src/java/org/apache/cassandra/service/reads/tracked/FilteredFollowupRead.java

@@ -109,7 +109,7 @@ class FilteredFollowupRead extends AsyncPromise<TrackedDataResponse>
         {
             partialRead = new AtomicReference<>();
             TrackedRead.Range rangeRead = makeFollowUpRead(command, followUpBounds, remaining, consistencyLevel, requestTime);
-            rangeRead.startLocal(requestTime, partialRead::set);
+            rangeRead.startLocal(requestTime, partialRead::set, TrackedLocalReads.Completer.DEFAULT);
             futures.add(rangeRead.future());
         }
         else

src/java/org/apache/cassandra/service/reads/tracked/PartialTrackedIndexRead.java

@@ -0,0 +1,863 @@
+/*
+ * 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.service.reads.tracked;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.SortedMap;
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Iterators;
+import com.google.common.collect.PeekingIterator;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ConsistencyLevel;
+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.ReadCommand;
+import org.apache.cassandra.db.ReadExecutionController;
+import org.apache.cassandra.db.ReadableView;
+import org.apache.cassandra.db.SinglePartitionReadCommand;
+import org.apache.cassandra.db.Slices;
+import org.apache.cassandra.db.filter.ColumnFilter;
+import org.apache.cassandra.db.lifecycle.SSTableSet;
+import org.apache.cassandra.db.lifecycle.View;
+import org.apache.cassandra.db.memtable.Memtable;
+import org.apache.cassandra.db.partitions.Partition;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.db.partitions.PartitionUpdate;
+import org.apache.cassandra.db.partitions.SimpleBTreePartition;
+import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.partitions.UnfilteredPartitionIterators;
+import org.apache.cassandra.db.rows.UnfilteredRowIterator;
+import org.apache.cassandra.db.rows.UnfilteredSource;
+import org.apache.cassandra.db.transform.EmptyPartitionsDiscarder;
+import org.apache.cassandra.db.transform.Transformation;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.dht.ExcludingBounds;
+import org.apache.cassandra.dht.Range;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.index.Index.IndexMatch;
+import org.apache.cassandra.index.transactions.UpdateTransaction;
+import org.apache.cassandra.io.sstable.SSTableReadsListener;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.tcm.ClusterMetadata;
+import org.apache.cassandra.transport.Dispatcher;
+import org.apache.cassandra.utils.AbstractIterator;
+import org.apache.cassandra.utils.ByteBufferUtil;
+import org.apache.cassandra.utils.CloseablePeekingIterator;
+import org.apache.cassandra.utils.concurrent.AsyncPromise;
+import org.apache.cassandra.utils.concurrent.Future;
+import org.apache.cassandra.utils.concurrent.FutureCombiner;
+import org.apache.cassandra.utils.concurrent.UncheckedInterruptedException;
+
+public class PartialTrackedIndexRead<Match extends IndexMatch, Searcher extends Index.MultiStepSearcher<Match>> extends PartialTrackedRead
+{
+    private final ReadCommand command;
+    private final Searcher searcher;
+
+    private ConsistencyLevel consistencyLevel;
+    private Dispatcher.RequestTime requestTime;
+
+    PartialTrackedIndexRead(ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos, ReadCommand command, Searcher searcher)
+    {
+        super(executionController, cfs, startTimeNanos);
+        this.command = command;
+        this.searcher = searcher;
+    }
+
+    public static <Match extends IndexMatch, Searcher extends Index.MultiStepSearcher<Match>> PartialTrackedIndexRead<Match, Searcher> create(ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos, ReadCommand command, Searcher searcher)
+    {
+        PartialTrackedIndexRead<Match, Searcher> read = new PartialTrackedIndexRead<>(executionController, cfs, startTimeNanos, command, searcher);
+        read.prepare(null);
+        return read;
+    }
+
+    @Override
+    public ReadCommand command()
+    {
+        return command;
+    }
+
+    @Override
+    public Searcher searcher()
+    {
+        return searcher;
+    }
+
+    @Override
+    public void setFollowUpReadContext(ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime)
+    {
+        this.consistencyLevel = consistencyLevel;
+        this.requestTime = requestTime;
+    }
+
+    public interface CompletedIndexPartitionRead<Match extends IndexMatch>
+    {
+        UnfilteredRowIterator matchingRows(CloseablePeekingIterator<Match> matchIterator);
+    }
+
+    public interface CompletedIndexRead<Match extends IndexMatch> extends CompletedRead
+    {
+        CompletedIndexPartitionRead<Match> partitionRead(DecoratedKey key);
+        CloseablePeekingIterator<Match> matchIterator();
+    }
+
+    private static DecoratedKey maxKey(DecoratedKey left, DecoratedKey right)
+    {
+        if (left == null)
+            return right;
+        if (right == null)
+            return left;
+
+        return right.compareTo(left) > 0 ? right : left;
+    }
+
+    private static class FollowUpRead<Match extends IndexMatch, Searcher extends Index.MultiStepSearcher<Match>> implements CompletedIndexPartitionRead<Match>, AutoCloseable
+    {
+        private final DecoratedKey key;
+        private final PartialTrackedIndexRead<Match, Searcher> read;
+        private final CompletedIndexRead<Match> completedRead;
+        private final CompletedIndexPartitionRead<Match> partitionRead;
+
+        public FollowUpRead(DecoratedKey key, PartialTrackedIndexRead<Match, Searcher> read)
+        {
+            Preconditions.checkArgument(!read.command.isRangeRequest());
+            this.key = key;
+            this.read = read;
+            this.completedRead = (CompletedIndexRead<Match>) read.complete();
+            this.partitionRead = Preconditions.checkNotNull(completedRead.partitionRead(key));
+        }
+
+        static <Match extends IndexMatch, Searcher extends Index.MultiStepSearcher<Match>> Future<FollowUpRead<Match, Searcher>> start(ReadCommand command, DecoratedKey key, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime)
+        {
+            ClusterMetadata metadata = ClusterMetadata.current();
+            Preconditions.checkState(command instanceof PartitionRangeReadCommand, "additional reads can only be made with range reads");
+
+            PartitionRangeReadCommand rangeReadCommand = (PartitionRangeReadCommand) command;
+            SinglePartitionReadCommand partitionReadCommand = SinglePartitionReadCommand.fromRangeRead(key, rangeReadCommand, rangeReadCommand.limits());
+
+            AsyncPromise<FollowUpRead<Match, Searcher>> followUpPromise = new AsyncPromise<>();
+            TrackedRead.Partition trackedRead = TrackedRead.Partition.create(metadata, partitionReadCommand, consistencyLevel, requestTime);
+
+            trackedRead.startLocal(requestTime, null, ((promise1, read, consistencyLevel1, rt) -> {
+                try
+                {
+                    followUpPromise.trySuccess(new FollowUpRead<>(key, (PartialTrackedIndexRead<Match, Searcher>) read));
+                }
+                catch (Exception e)
+                {
+                    followUpPromise.tryFailure(e);
+                }
+            }));
+            return followUpPromise;
+        }
+
+        @Override
+        public UnfilteredRowIterator matchingRows(CloseablePeekingIterator<Match> matchIterator)
+        {
+            Preconditions.checkState(matchIterator.hasNext());
+            Preconditions.checkState(matchIterator.peek().key().equals(key));
+            return partitionRead.matchingRows(matchIterator);
+        }
+
+        @Override
+        public void close()
+        {
+            read.close();
+        }
+
+        static <Match extends IndexMatch, Searcher extends Index.MultiStepSearcher<Match>> void close(Map<DecoratedKey, Future<FollowUpRead<Match, Searcher>>> followUpReads)
+        {
+            for (Future<FollowUpRead<Match, Searcher>> future : followUpReads.values())
+            {
+                future.addCallback((followup, failure) -> {
+                    if (failure != null)
+                        followup.close();
+                });
+            }
+        }
+
+        static <Match extends IndexMatch, Searcher extends Index.MultiStepSearcher<Match>> Map<DecoratedKey, FollowUpRead<Match, Searcher>> getResults(Map<DecoratedKey, Future<FollowUpRead<Match, Searcher>>> futures, List<CloseablePeekingIterator<Match>> matchIterators)
+        {
+            Map<DecoratedKey, FollowUpRead<Match, Searcher>> followupReads = new HashMap<>();
+            for (Future<FollowUpRead<Match, Searcher>> future : futures.values())
+            {
+                try
+                {
+                    FollowUpRead<Match, Searcher> followUpRead = future.get();
+                    matchIterators.add(followUpRead.completedRead.matchIterator());
+                    followupReads.put(followUpRead.key, followUpRead);
+                }
+                catch (ExecutionException e)
+                {
+                    throw new RuntimeException(e);
+                }
+                catch (InterruptedException e)
+                {
+                    throw new UncheckedInterruptedException(e);
+                }
+            }
+            return followupReads;
+        }
+
+    }
+
+    private static class SnapshotView implements ReadableView
+    {
+        final List<SinglePartitionSource> snapshots;
+        final List<SSTableReader> sstables;
+        private AugmentedPartition augmentedPartition = null;
+
+        public SnapshotView(List<SinglePartitionSource> snapshots, List<SSTableReader> sstables)
+        {
+            this.snapshots = snapshots;
+            this.sstables = sstables;
+        }
+
+        public static SnapshotView create(DecoratedKey key, ColumnFamilyStore cfs)
+        {
+            ColumnFamilyStore.ViewFragment view = cfs.select(View.select(SSTableSet.LIVE, key));
+            return new SnapshotView(MemtableSnapshot.create(key, view.memtables), view.sstables());
+        }
+
+        @Override
+        public Iterable<? extends UnfilteredSource> memtables()
+        {
+            return snapshots;
+        }
+
+        @Override
+        public List<SSTableReader> sstables()
+        {
+            return sstables;
+        }
+
+        public void augment(PartitionUpdate update)
+        {
+            if (augmentedPartition == null)
+            {
+                augmentedPartition = new AugmentedPartition(update.partitionKey(), update.metadata());
+                snapshots.add(augmentedPartition);
+            }
+
+            augmentedPartition.augment(update);
+        }
+    }
+
+    private static abstract class SinglePartitionSource implements UnfilteredSource
+    {
+        abstract Partition partition();
+
+        @Override
+        public UnfilteredRowIterator rowIterator(DecoratedKey key, Slices slices, ColumnFilter columnFilter, boolean reversed, SSTableReadsListener listener)
+        {
+            Partition partition = partition();
+            Preconditions.checkState(key.equals(partition.partitionKey()));
+            return partition.unfilteredIterator(columnFilter, slices, reversed);
+        }
+
+        @Override
+        public UnfilteredPartitionIterator partitionIterator(ColumnFilter columnFilter, DataRange dataRange, SSTableReadsListener listener)
+        {
+            throw new IllegalStateException("Range scans not supported");
+        }
+
+        @Override
+        public long getMinTimestamp()
+        {
+            return partition().stats().minTimestamp;
+        }
+
+        @Override
+        public long getMinLocalDeletionTime()
+        {
+            return partition().stats().minLocalDeletionTime;
+        }
+    }
+
+    private static class MemtableSnapshot extends SinglePartitionSource
+    {
+        private final Partition partition;
+
+        public MemtableSnapshot(Partition partition)
+        {
+            this.partition = partition;
+        }
+
+        static List<SinglePartitionSource> create(DecoratedKey key, Iterable<Memtable> memtables)
+        {
+            List<SinglePartitionSource> snapshots = new ArrayList<>();
+            for (Memtable memtable : memtables)
+            {
+                Partition partition = memtable.snapshotPartition(key);
+                if (partition != null)
+                    snapshots.add(new MemtableSnapshot(partition));
+            }
+            return snapshots;
+        }
+
+        @Override
+        Partition partition()
+        {
+            return partition;
+        }
+    }
+
+    private static class AugmentedPartition extends SinglePartitionSource
+    {
+        private final SimpleBTreePartition data;
+
+        AugmentedPartition(DecoratedKey key, TableMetadata metadata)
+        {
+            this.data = new SimpleBTreePartition(key, metadata, UpdateTransaction.NO_OP);
+        }
+
+        void augment(PartitionUpdate update)
+        {
+            data.update(update);
+        }
+
+        @Override
+        Partition partition()
+        {
+            return data;
+        }
+    }
+
+    class AugmentableIndexPartitionRead implements CompletedIndexPartitionRead<Match>
+    {
+        private final DecoratedKey partitionKey;
+        private final SnapshotView view;
+
+        AugmentableIndexPartitionRead(DecoratedKey partitionKey, SnapshotView view)
+        {
+            this.partitionKey = partitionKey;
+            this.view = view;
+        }
+
+        void augment(PartitionUpdate update)
+        {
+            Preconditions.checkArgument(update.partitionKey().equals(partitionKey));
+            view.augment(update);
+        }
+
+        @Override
+        public UnfilteredRowIterator matchingRows(CloseablePeekingIterator<Match> matchIterator)
+        {
+            Preconditions.checkArgument(matchIterator.hasNext());
+            Preconditions.checkArgument(matchIterator.peek().key().equals(partitionKey));
+            return searcher.queryNextMatches(executionController, partitionKey, view, matchIterator);
+        }
+    }
+
+    AugmentableIndexPartitionRead createRead(DecoratedKey key, ColumnFamilyStore cfs)
+    {
+        SnapshotView view = SnapshotView.create(key, cfs);
+        return new AugmentableIndexPartitionRead(key, view);
+    }
+
+    @Override
+    protected Prepared prepareInternal(UnfilteredPartitionIterator initialData)
+    {
+        DecoratedKey maxKey = null;
+        SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads = new TreeMap<>();
+
+        if (command instanceof SinglePartitionReadCommand)
+        {
+            SinglePartitionReadCommand cmd = (SinglePartitionReadCommand) command;
+            DecoratedKey key = cmd.partitionKey();
+            AugmentableIndexPartitionRead partitionRead = createRead(key, cfs);
+            reads.put(key, partitionRead);
+            maxKey = key;
+        }
+
+        CloseablePeekingIterator<Match> matchIterator = searcher.matchIterator(executionController);
+        try
+        {
+            SortedSet<Match> materializedMatches = new TreeSet<>(searcher.matchComparator());
+            while (matchIterator.hasNext() && materializedMatches.size() < command.limits().count())
+            {
+                Match match = matchIterator.next();
+                materializedMatches.add(match);
+                if (!reads.containsKey(match.key()))
+                {
+                    DecoratedKey key = match.key();
+                    maxKey = maxKey(maxKey, key);
+                    AugmentableIndexPartitionRead partitionRead = createRead(key, cfs);
+                    reads.put(key, partitionRead);
+                }
+            }
+            return new IndexPrepared(maxKey, materializedMatches, matchIterator, reads);
+        }
+        catch (Throwable t)
+        {
+            FileUtils.closeQuietly(matchIterator);
+            throw t;
+        }
+    }
+
+    @Override
+    public synchronized void complete(AsyncPromise<TrackedDataResponse> promise, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime)
+    {
+        Preconditions.checkState(state().isPrepared());
+        IndexPrepared prepared = (IndexPrepared) state();
+
+        if (prepared.isCompletable())
+        {
+            super.complete(promise, consistencyLevel, requestTime);
+            return;
+        }
+
+        IndexPreComplete preComplete = prepared.preComplete();
+        state = preComplete;
+
+        // simple listener - completion will handle any failed futures
+        preComplete.future().addListener(() -> super.complete(promise, consistencyLevel, requestTime));
+    }
+
+    private abstract class AbstractIndexPrepared extends Prepared
+    {
+        protected DecoratedKey maxKey;
+        protected final SortedSet<Match> materializedMatches;
+        // there may be additional matches for keys we've already scanned, this allows us to read them before
+        // starting a short read
+        protected final CloseablePeekingIterator<Match> additionalMatches;
+
+        protected final SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads;
+
+        // for range scans, if we learn of new keys with matching contents as part of reconciliation, we need
+        // to do follow up reads against them since we didn't snapshot memtable contents for the keys during
+        // the prepare phase of the read. Futures for those reads are kept here
+        protected final Map<DecoratedKey, Future<FollowUpRead<Match, Searcher>>> followUpReads;
+
+        public AbstractIndexPrepared(DecoratedKey maxKey,
+                                     SortedSet<Match> materializedMatches,
+                                     CloseablePeekingIterator<Match> additionalMatches,
+                                     SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads,
+                                     Map<DecoratedKey, Future<FollowUpRead<Match, Searcher>>> followUpReads)
+        {
+            this.maxKey = maxKey;
+            this.materializedMatches = materializedMatches;
+            this.reads = reads;
+            this.additionalMatches = additionalMatches;
+            this.followUpReads = followUpReads;
+        }
+
+        boolean isCompletable()
+        {
+            return Iterables.all(followUpReads.values(), Future::isDone);
+        }
+
+        @Override
+        Completed complete()
+        {
+            Preconditions.checkState(isCompletable());
+            List<CloseablePeekingIterator<Match>> matchIterators = new ArrayList<>(followUpReads.size() + 1);
+            matchIterators.add(CloseablePeekingIterator.wrap(materializedMatches.iterator()));
+            Map<DecoratedKey, FollowUpRead<Match, Searcher>> followUpResults = FollowUpRead.getResults(followUpReads, matchIterators);
+            return new IndexCompleted(maxKey, new MergingMatchIterator(matchIterators), additionalMatches, reads, followUpResults);
+        }
+
+        @Override
+        void close()
+        {
+            FollowUpRead.close(followUpReads);
+            super.close();
+        }
+    }
+
+    private class IndexPrepared extends AbstractIndexPrepared
+    {
+        private Index.MultiStepSearcher.MatchIndexer<Match> matchIndexer = null;
+
+        public IndexPrepared(DecoratedKey maxKey, SortedSet<Match> materializedMatches, CloseablePeekingIterator<Match> additionalMatches, SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads)
+        {
+            super(maxKey, materializedMatches, additionalMatches, reads, new HashMap<>());
+        }
+
+        private Index.MultiStepSearcher.MatchIndexer<Match> matchIndexer()
+        {
+            if (matchIndexer == null)
+                matchIndexer = searcher.matchIndexer();
+            return matchIndexer;
+        }
+
+        private boolean indexNewKey(PartitionUpdate update)
+        {
+            AtomicBoolean hasMatches = new AtomicBoolean(false);
+            matchIndexer().index(update, e -> hasMatches.set(true));
+            return hasMatches.get();
+        }
+
+        private boolean indexUpdate(PartitionUpdate update)
+        {
+            int startingSize = materializedMatches.size();
+            matchIndexer().index(update, materializedMatches::add);
+            return materializedMatches.size() > startingSize;
+        }
+
+        @Override
+        public void augment(PartitionUpdate update)
+        {
+            Preconditions.checkState(consistencyLevel != null,
+                                     "PartialTrackedRead#setFollowUpReadContext needs to be called before making reads available for augmenting mutation");
+            DecoratedKey key = update.partitionKey();
+            AugmentableIndexPartitionRead read = reads.get(key);
+            if (read == null)
+            {
+                // TODO: maybe we should immediately start a follow up read if it's likely this key will be included in the response
+                if (!followUpReads.containsKey(key) && indexNewKey(update))
+                {
+                    maxKey = maxKey(maxKey, update.partitionKey());
+                    Future<FollowUpRead<Match, Searcher>> followUpRead = FollowUpRead.start(command, update.partitionKey(), consistencyLevel, requestTime);
+                    followUpReads.put(key, followUpRead);
+                }
+                return;
+            }
+
+            read.augment(update);
+            indexUpdate(update);
+        }
+
+        IndexPreComplete preComplete()
+        {
+            return new IndexPreComplete(maxKey, materializedMatches, additionalMatches, reads, followUpReads);
+        }
+    }
+
+    private class IndexPreComplete extends AbstractIndexPrepared
+    {
+        public IndexPreComplete(DecoratedKey maxKey, SortedSet<Match> materializedMatches, CloseablePeekingIterator<Match> additionalMatches, SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads, Map<DecoratedKey, Future<FollowUpRead<Match, Searcher>>> followUpReads)
+        {
+            super(maxKey, materializedMatches, additionalMatches, reads, followUpReads);
+        }
+
+        @Override
+        public void augment(PartitionUpdate update)
+        {
+            throw new IllegalStateException("Cannot augment reads pending completion");
+        }
+
+        Future<List<FollowUpRead<Match, Searcher>>> future()
+        {
+            return FutureCombiner.allOf(followUpReads.values());
+        }
+    }
+
+    private class IndexCompleted extends Completed
+    {
+        private final DecoratedKey maxKey;
+        private final CloseablePeekingIterator<Match> materializedMatchIterator;
+        private final CloseablePeekingIterator<Match> additionalMatchIterator;
+        private final SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads;
+        private final Map<DecoratedKey, FollowUpRead<Match, Searcher>> followUpReads;
+
+        public IndexCompleted(DecoratedKey maxKey, CloseablePeekingIterator<Match> materializedMatchIterator, CloseablePeekingIterator<Match> additionalMatchIterator, SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads, Map<DecoratedKey, FollowUpRead<Match, Searcher>> followUpReads)
+        {
+            this.maxKey = maxKey;
+            this.materializedMatchIterator = materializedMatchIterator;
+            this.additionalMatchIterator = additionalMatchIterator;
+            this.reads = reads;
+            this.followUpReads = followUpReads;
+        }
+
+        @Override
+        protected CompletedRead getResult()
+        {
+            return new FilteringCompletedIndexRead(maxKey, materializedMatchIterator, additionalMatchIterator, reads, followUpReads);
+        }
+    }
+
+    protected class MergingMatchIterator extends AbstractIterator<Match>
+    {
+        private final List<CloseablePeekingIterator<Match>> iterators;
+        private Match last;
+
+        public MergingMatchIterator(List<CloseablePeekingIterator<Match>> iterators)
+        {
+            this.iterators = iterators;
+        }
+
+        @Override
+        protected Match computeNext()
+        {
+            int minIdx = -1;
+            Match minMatch = null;
+            for (int i = 0, mi = iterators.size(); i < mi; i++)
+            {
+                CloseablePeekingIterator<Match> iterator = iterators.get(i);
+
+                if (last != null)
+                    searcher.matchComparator().consumeDuplicates(last, iterator);
+
+                if (!iterator.hasNext())
+                    continue;
+
+                if (minMatch == null)
+                {
+                    minMatch = iterator.peek();
+                    minIdx = i;
+                    continue;
+                }
+
+                Match thisMatch = iterator.peek();
+                int cmp = searcher.matchComparator().compare(thisMatch, minMatch);
+                if (cmp < 0)
+                {
+                    minMatch = thisMatch;
+                    minIdx = i;
+                }
+                else if (cmp == 0)
+                {
+                    // if this iterator equals the current minimum, advance the iterator - we don't merge equal matches
+                    iterator.next();
+                }
+            }
+
+            if (minMatch != null)
+            {
+                iterators.get(minIdx).next();
+                last = minMatch;
+                return minMatch;
+            }
+
+            return endOfData();
+        }
+
+        @Override
+        public void close()
+        {
+            FileUtils.closeQuietly(iterators);
+        }
+    }
+
+    /**
+     * Merges a materialized iterator and an additional iterator. The additional iterator is meant to be the initial
+     * match iterator from the searcher. If we encounter previously unseen keys from the initial match iterator, it
+     * means that we're in a short read and need to start a follow-up read, which this iterator signals to the caller
+     */
+    private class MergingStoppingMatchIterator extends AbstractIterator<Match>
+    {
+        private final DecoratedKey maxKey;
+        private final PeekingIterator<Match> materializedIterator;
+        private final CloseablePeekingIterator<Match> additionalIterator;
+        private boolean followUpRequired = false;
+
+        public MergingStoppingMatchIterator(DecoratedKey maxKey, Iterator<Match> materializedIterator, CloseablePeekingIterator<Match> additionalIterator)
+        {
+            this.maxKey = maxKey;
+            this.materializedIterator = Iterators.peekingIterator(materializedIterator);
+            this.additionalIterator = additionalIterator;
+        }
+
+        @Override
+        protected Match computeNext()
+        {
+            if (materializedIterator.hasNext() && additionalIterator.hasNext())
+            {
+                int cmp = searcher.matchComparator().compare(materializedIterator.peek(), additionalIterator.peek(), true);
+                if (cmp == 0)
+                {
+                    additionalIterator.next();
+                    return materializedIterator.next();
+                }
+                else if (cmp < 0)
+                {
+                    Match match = materializedIterator.next();
+                    searcher.matchComparator().consumeDuplicates(match, additionalIterator);
+                    return match;
+                }
+                else
+                {
+                    Match match = additionalIterator.next();
+                    searcher.matchComparator().consumeDuplicates(match, materializedIterator);
+
+                    DecoratedKey key = match.key();
+                    Preconditions.checkArgument(key.compareTo(maxKey) <= 0);
+                    return match;
+                }
+            }
+
+            if (materializedIterator.hasNext())
+                return materializedIterator.next();
+
+            if (additionalIterator.hasNext())
+            {
+                Match match = additionalIterator.next();
+                DecoratedKey key = match.key();
+                if (key.compareTo(maxKey) > 0)
+                {
+                    Preconditions.checkArgument(command.isRangeRequest());
+                    followUpRequired = true;
+                    return endOfData();
+                }
+                return match;
+            }
+
+            return endOfData();
+        }
+
+        @Override
+        public void close()
+        {
+            additionalIterator.close();
+        }
+    }
+
+    private class FilteringCompletedIndexRead extends ExtendingCompletedRead implements CompletedIndexRead<Match>
+    {
+        private final DecoratedKey maxKey;
+        private final MergingStoppingMatchIterator matchIterator;
+        private final SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads;
+
+        final Map<DecoratedKey, FollowUpRead<Match, Searcher>> followUpReads;
+
+        public FilteringCompletedIndexRead(DecoratedKey maxKey, CloseablePeekingIterator<Match> materializedMatches, CloseablePeekingIterator<Match> additionalMatches, SortedMap<DecoratedKey, AugmentableIndexPartitionRead> reads, Map<DecoratedKey, FollowUpRead<Match, Searcher>> followupReads)
+        {
+            super(command, materializedMatches.hasNext(), true);
+            this.maxKey = maxKey;
+            this.matchIterator = new MergingStoppingMatchIterator(maxKey, materializedMatches, additionalMatches);
+            this.reads = reads;
+            this.followUpReads = followupReads;
+        }
+
+        @Override
+        public CloseablePeekingIterator<Match> matchIterator()
+        {
+            return matchIterator;
+        }
+
+        @Override
+        ReadCommand command()
+        {
+            return command;
+        }
+
+        @Override
+        protected AbstractBounds<PartitionPosition> followUpBounds()
+        {
+            Preconditions.checkState(command.isRangeRequest());
+            Preconditions.checkNotNull(maxKey);
+            AbstractBounds<PartitionPosition> bounds = command.dataRange().keyRange();
+            return bounds.inclusiveRight()
+                   ? new Range<>(maxKey, bounds.right)
+                   : new ExcludingBounds<>(maxKey, bounds.right);
+        }
+
+        private class UnfilteredResultIterator extends AbstractIterator<UnfilteredRowIterator> implements UnfilteredPartitionIterator
+        {
+            private final CloseablePeekingIterator<Match> matchIter;
+
+            public UnfilteredResultIterator(CloseablePeekingIterator<Match> matchIter)
+            {
+                this.matchIter = matchIter;
+            }
+
+            @Override
+            public TableMetadata metadata()
+            {
+                return command.metadata();
+            }
+
+            @Override
+            protected UnfilteredRowIterator computeNext()
+            {
+                for (;;)
+                {
+                    if (!matchIter.hasNext())
+                        return endOfData();
+
+                    DecoratedKey nextKey = matchIter.peek().key();
+                    AugmentableIndexPartitionRead read = reads.get(nextKey);
+                    if (read != null)
+                        return read.matchingRows(matchIter);
+
+                    FollowUpRead<Match, Searcher> followUpRead = followUpReads.get(nextKey);
+                    if (followUpRead == null)
+                        throw new IllegalStateException("Received match for key without initial or followup read: " + ByteBufferUtil.bytesToHex(nextKey.getKey()));
+
+                    UnfilteredRowIterator next = followUpRead.matchingRows(matchIter);
+                    if (next != null)
+                        return next;
+                }
+            }
+
+            @Override
+            public void close()
+            {
+                matchIter.close();
+            }
+        }
+
+        private PartitionIterator filter(UnfilteredPartitionIterator iterator)
+        {
+            iterator = searcher.filterCompletedRead(iterator);
+            iterator = command.completeTrackedRead(iterator, PartialTrackedIndexRead.this);
+            PartitionIterator filtered = UnfilteredPartitionIterators.filter(iterator, command.nowInSec());
+            PartitionIterator counted = Transformation.apply(filtered, mergedResultCounter);
+            PartitionIterator result = Transformation.apply(counted, new EmptyPartitionsDiscarder());
+            return result;
+        }
+
+        @Override
+        public TrackedDataResponse response()
+        {
+            try (UnfilteredResultIterator iterator = new UnfilteredResultIterator(matchIterator))
+            {
+                PartitionIterator filtered = filter(iterator);
+                return TrackedDataResponse.create(filtered, command.columnFilter());
+            }
+        }
+
+        @Override
+        protected boolean followUpRequired()
+        {
+            if (!command.isRangeRequest())
+                return false;
+            return matchIterator.followUpRequired || super.followUpRequired();
+        }
+
+        @Override
+        public void close()
+        {
+            FileUtils.closeQuietly(matchIterator);
+            FileUtils.closeQuietly(followUpReads.values());
+        }
+
+        @Override
+        public CompletedIndexPartitionRead<Match> partitionRead(DecoratedKey key)
+        {
+            return reads.get(key);
+        }
+    }
+}

src/java/org/apache/cassandra/service/reads/tracked/PartialTrackedRangeRead.java

@@ -25,7 +25,6 @@ import java.util.SortedMap;
 import java.util.TreeMap;
 
 import com.google.common.base.Preconditions;
-import org.apache.cassandra.transport.Dispatcher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -45,6 +44,7 @@ import org.apache.cassandra.db.partitions.AbstractUnfilteredPartitionIterator;
 import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.partitions.SimpleBTreePartition;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.partitions.UnfilteredPartitionIterators;
 import org.apache.cassandra.db.rows.UnfilteredRowIterator;
 import org.apache.cassandra.db.transform.Transformation;
 import org.apache.cassandra.dht.AbstractBounds;
@@ -55,9 +55,10 @@ import org.apache.cassandra.index.transactions.UpdateTransaction;
 import org.apache.cassandra.locator.ReplicaPlan;
 import org.apache.cassandra.locator.ReplicaPlans;
 import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.transport.Dispatcher;
 import org.apache.cassandra.utils.concurrent.Future;
 
-public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
+public abstract class PartialTrackedRangeRead extends PartialTrackedRead
 {
     private static final Logger logger = LoggerFactory.getLogger(PartialTrackedRangeRead.class);
 
@@ -111,19 +112,15 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
 
     protected static class ShortReadSupport
     {
-        final DecoratedKey lastPartitionKey; // key of the last observed partition
         final boolean partitionsFetched; // whether we've seen any new partitions since iteration start or last moreContents() call
         final boolean initialIteratorExhausted;
         final AbstractBounds<PartitionPosition> followUpBounds;
-        boolean wasAugmented;
 
         ShortReadSupport(Builder builder, boolean initialIteratorExhausted, AbstractBounds<PartitionPosition> followUpBounds)
         {
-            this.lastPartitionKey = builder.lastPartitionKey;
             this.partitionsFetched = builder.partitionsFetched;
             this.initialIteratorExhausted = initialIteratorExhausted;
             this.followUpBounds = followUpBounds;
-            this.wasAugmented = false;
         }
 
         protected static class Builder
@@ -175,34 +172,19 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
 
         RangePrepared materialize(UnfilteredPartitionIterator inputIterator)
         {
-            try
+            try (inputIterator)
             {
-                UnfilteredPartitionIterator materialized = Transformation.apply(inputIterator, new Transformation<UnfilteredRowIterator>()
-                {
-                    @Override
-                    protected UnfilteredRowIterator applyToPartition(UnfilteredRowIterator partition)
-                    {
-                        SimpleBTreePartition materialized = data.computeIfAbsent(partition.partitionKey(), key -> new SimpleBTreePartition(key, partition.metadata(), UpdateTransaction.NO_OP));
-                        materialized.update(PartitionUpdate.fromIterator(partition, command.columnFilter()));
-                        shortReadSupport.lastPartitionKey = partition.partitionKey();
-                        shortReadSupport.partitionsFetched = true;
-                        return queryPartition(materialized);
-                    }
-                });
+                UnfilteredPartitionIterator materialized = Transformation.apply(inputIterator, this);
 
                 UnfilteredPartitionIterator filtered = filter(materialized);
 
                 try (UnfilteredPartitionIterator iterator = shortReadSupport.counter.applyTo(filtered))
                 {
-                    consume(iterator);
+                    UnfilteredPartitionIterators.consume(iterator);
                 }
 
                 return createRangePrepared();
             }
-            finally
-            {
-                inputIterator.close();
-            }
         }
 
         @Override
@@ -241,7 +223,7 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
         }
 
         @Override
-        public State augment(PartitionUpdate update)
+        public void augment(PartitionUpdate update)
         {
             // if the input iterator reached the row limit, then we can't apply any augmenting mutations that are past
             // the last materialized key. Since we wouldn't have materialized the local data for that key, applying an
@@ -256,11 +238,10 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
                 logger.trace("Ignoring unacceptable update from key {} on read {}", update.partitionKey(), PartialTrackedRangeRead.this);
             }
             wasAugmented = true;
-            return this;
         }
     }
 
-    protected abstract class RangeCompleted extends Completed
+    protected abstract class RangeCompleted extends AbstractCompleted
     {
         protected final SortedMap<DecoratedKey, SimpleBTreePartition> data;
         protected final ShortReadSupport shortReadSupport;
@@ -337,18 +318,6 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
                                             command.clusteringIndexFilter(partition.partitionKey()).isReversed());
     }
 
-    private static void consume(UnfilteredPartitionIterator iterator)
-    {
-        while (iterator.hasNext())
-        {
-            try (UnfilteredRowIterator partition = iterator.next())
-            {
-                while (partition.hasNext())
-                    partition.next();
-            }
-        }
-    }
-
     public AbstractBounds<PartitionPosition> followUpBounds()
     {
         RangeCompleted completed = (RangeCompleted) state().asCompleted();
@@ -401,7 +370,7 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
             @Override
             protected CompletedRead extendRead(UnfilteredPartitionIterator iterator)
             {
-                return new ExtendingCompletedRead(command, iterator, shortReadSupport.partitionsFetched, shortReadSupport.initialIteratorExhausted, shortReadSupport.followUpBounds);
+                return new ExtendingCompletedRead.RangeRead(command, iterator, shortReadSupport.partitionsFetched, shortReadSupport.initialIteratorExhausted, shortReadSupport.followUpBounds);
             }
         }
 
@@ -523,7 +492,7 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
             }
         }
 
-        static class FilteredCompletedRead extends ExtendingCompletedRead
+        static class FilteredCompletedRead extends ExtendingCompletedRead.RangeRead
         {
             private final DecoratedKey lastMatchingKey;
             private final SortedMap<DecoratedKey, FollowUpReadInfo> followUpReadInfo;
@@ -537,7 +506,6 @@ public abstract class PartialTrackedRangeRead extends AbstractPartialTrackedRead
             /**
              * Even if we reached the limit during materialization, if there are keys ahead of the first materialized key
              * or interleaved with them, then we need to read them
-             * @return
              */
             private boolean hasInterleavedFollowupKeys()
             {

src/java/org/apache/cassandra/service/reads/tracked/PartialTrackedRead.java

@@ -15,10 +15,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package org.apache.cassandra.service.reads.tracked;
 
 import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.cassandra.db.ColumnFamilyStore;
 import org.apache.cassandra.db.ConsistencyLevel;
@@ -27,22 +28,257 @@ import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
 import org.apache.cassandra.db.filter.DataLimits;
 import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterators;
+import org.apache.cassandra.db.transform.RTBoundValidator;
 import org.apache.cassandra.index.Index;
+import org.apache.cassandra.utils.concurrent.AsyncPromise;
 import org.apache.cassandra.utils.concurrent.Future;
 import org.apache.cassandra.replication.Log2OffsetsMap;
 import org.apache.cassandra.replication.MutationJournal;
 import org.apache.cassandra.replication.ShortMutationId;
 import org.apache.cassandra.transport.Dispatcher;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
-public interface PartialTrackedRead
+public abstract class PartialTrackedRead
 {
-    Logger logger = LoggerFactory.getLogger(PartialTrackedRead.class);
-    
-    interface CompletedRead extends AutoCloseable
+    private static final Logger logger = LoggerFactory.getLogger(PartialTrackedRead.class);
+
+    final ReadExecutionController executionController;
+    final ColumnFamilyStore cfs;
+    final long startTimeNanos;
+
+    public PartialTrackedRead(ReadExecutionController executionController, ColumnFamilyStore cfs, long startTimeNanos)
+    {
+        this.executionController = executionController;
+        this.cfs = cfs;
+        this.startTimeNanos = startTimeNanos;
+    }
+
+    public ReadExecutionController executionController()
+    {
+        return executionController;
+    }
+
+    public ColumnFamilyStore cfs()
+    {
+        return cfs;
+    }
+
+    public long startTimeNanos()
+    {
+        return startTimeNanos;
+    }
+
+    abstract ReadCommand command();
+
+    public abstract Index.Searcher searcher();
+
+    protected interface Augmentable
+    {
+        void augment(PartitionUpdate update);
+    }
+
+    protected static abstract class State
+    {
+        protected static final State CLOSED = new State()
+        {
+            @Override
+            String name()
+            {
+                return "closed";
+            }
+
+            @Override
+            boolean isClosed()
+            {
+                return true;
+            }
+        };
+
+        abstract String name();
+
+        Initialized asInitialized()
+        {
+            throw new IllegalStateException("State is " + name() + ", not " + Initialized.NAME);
+        }
+
+        boolean isPrepared()
+        {
+            return false;
+        }
+
+        Prepared asPrepared()
+        {
+            throw new IllegalStateException("State is " + name() + ", not " + Prepared.NAME);
+        }
+
+        Completed asCompleted()
+        {
+            throw new IllegalStateException("State is " + name() + ", not " + Completed.NAME);
+        }
+
+        Augmentable asAugmentable()
+        {
+            if (isPrepared()) return asPrepared();
+            throw new IllegalStateException("State is " + name() + ", not augmentable");
+        }
+
+        boolean isClosed()
+        {
+            return false;
+        }
+
+        void close()
+        {
+        }
+    }
+
+    // TODO (expected): this is a redundant state, never exposed
+    protected final class Initialized extends State
+    {
+        static final String NAME = "initialized";
+
+        @Override
+        String name()
+        {
+            return NAME;
+        }
+
+        @Override
+        Initialized asInitialized()
+        {
+            return this;
+        }
+
+        Prepared prepare(UnfilteredPartitionIterator initialData)
+        {
+            return prepareInternal(initialData);
+        }
+    }
+
+    protected abstract Prepared prepareInternal(UnfilteredPartitionIterator initialData);
+
+    protected static abstract class Prepared extends State implements Augmentable
+    {
+        private static final String NAME = "prepared";
+
+        @Override
+        String name()
+        {
+            return NAME;
+        }
+
+        @Override
+        boolean isPrepared()
+        {
+            return true;
+        }
+
+        @Override
+        Prepared asPrepared()
+        {
+            return this;
+        }
+
+        abstract Completed complete();
+    }
+
+    protected static abstract class Completed extends State
+    {
+        private static final String NAME = "completed";
+
+        @Override
+        String name()
+        {
+            return NAME;
+        }
+
+        protected abstract CompletedRead getResult();
+    }
+
+    protected abstract class AbstractCompleted extends Completed
+    {
+        protected abstract UnfilteredPartitionIterator iterator();
+        protected abstract CompletedRead createResult(UnfilteredPartitionIterator iterator);
+
+        @Override
+        protected CompletedRead getResult()
+        {
+            UnfilteredPartitionIterator result = command().completeTrackedRead(iterator(), PartialTrackedRead.this);
+            // validate that the sequence of RT markers is correct: open is followed by close, deletion times for both
+            // ends equal, and there are no dangling RT bound in any partition.
+            result = RTBoundValidator.validate(result, RTBoundValidator.Stage.PROCESSED, true);
+            return createResult(result);
+        }
+    }
+
+    protected State state = new Initialized();
+
+    protected synchronized State state()
+    {
+        return state;
+    }
+
+    /**
+     * Implementors need to call this before returning this from createInProgressRead
+     * TODO (expected): this is a redundant transition from a redundant state (INITIALIZED)
+     */
+    synchronized void prepare(UnfilteredPartitionIterator initialData)
+    {
+        logger.trace("Preparing read {}", this);
+        state = state.asInitialized().prepare(initialData);
+    }
+
+    void augment(PartitionUpdate update)
+    {
+        state.asAugmentable().augment(update);
+    }
+
+    public synchronized void augment(Mutation mutation)
+    {
+        PartitionUpdate update = mutation.getPartitionUpdate(command().metadata());
+        if (update != null)
+            augment(update);
+    }
+
+    void augment(Log2OffsetsMap<?> augmentingOffsets)
+    {
+        augmentingOffsets.forEach(this::augment);
+    }
+
+    void augment(ShortMutationId mutationId)
+    {
+        Mutation mutation = MutationJournal.instance.read(mutationId);
+        Preconditions.checkNotNull(mutation, "Missing mutation %s", mutationId);
+        if (!command().selectsKey(mutation.key()))
+        {
+            logger.trace("Skipping mutation {} - {} not in read range", mutationId, mutation.key());
+            return;
+        }
+        augment(mutation);
+    }
+
+    public synchronized CompletedRead complete()
+    {
+        Preconditions.checkState(state.isPrepared());
+        Completed completed = state.asPrepared().complete();
+        state = completed;
+        return completed.getResult();
+    }
+
+    public synchronized void close()
+    {
+        if (state.isClosed())
+            return;
+
+        logger.trace("Closing read {}", this);
+        state.close();
+        executionController.close();
+        state = State.CLOSED;
+    }
+
+    public interface CompletedRead extends AutoCloseable
     {
         TrackedDataResponse response(); // must be called from the read stage
         Future<TrackedDataResponse> followupRead(TrackedDataResponse initialResponse, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime);
@@ -86,36 +322,48 @@ public interface PartialTrackedRead
         }
     }
 
-    CompletedRead complete();
+    /**
+     * Sets consistency level and expiration info to be used for follow up reads. Needs to be called before making the
+     * read available for receiving augmenting mutations
+     */
+    void setFollowUpReadContext(ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime) {}
 
-    void augment(Mutation mutation);
-
-    default void augment(Log2OffsetsMap<?> augmentingOffsets)
+    void complete(AsyncPromise<TrackedDataResponse> promise, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime)
     {
-        augmentingOffsets.forEach(this::augment);
+        complete(promise, this, consistencyLevel, requestTime);
     }
 
-    default void augment(ShortMutationId mutationId)
+    static void complete(AsyncPromise<TrackedDataResponse> promise, PartialTrackedRead read, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime)
     {
-        Mutation mutation = MutationJournal.instance.read(mutationId);
-        Preconditions.checkNotNull(mutation, "Missing mutation %s", mutationId);
-        if (!command().selectsKey(mutation.key()))
+        try (CompletedRead completedRead = read.complete())
         {
-            logger.trace("Skipping mutation {} - {} not in read range", mutationId, mutation.key());
-            return;
+            TrackedDataResponse response = completedRead.response();
+            Future<TrackedDataResponse> followUp = completedRead.followupRead(response, consistencyLevel, requestTime);
+
+            if (followUp != null)
+            {
+                followUp.addCallback((newResponse, error) -> {
+                    if (error != null)
+                    {
+                        promise.tryFailure(error);
+                        return;
+                    }
+                    promise.trySuccess(newResponse);
+                });
+            }
+            else
+            {
+                promise.trySuccess(response);
+            }
+        }
+        catch (Exception e)
+        {
+            promise.tryFailure(e);
+            throw e;
+        }
+        finally
+        {
+            read.close();
         }
-        augment(mutation);
     }
-
-    ReadExecutionController executionController();
-
-    Index.Searcher searcher();
-
-    ColumnFamilyStore cfs();
-
-    long startTimeNanos();
-
-    ReadCommand command();
-
-    void close();
 }

src/java/org/apache/cassandra/service/reads/tracked/PartialTrackedSinglePartitionRead.java

@@ -36,7 +36,7 @@ import org.apache.cassandra.index.transactions.UpdateTransaction;
 
 import static org.apache.cassandra.db.partitions.UnfilteredPartitionIterators.MergeListener.NOOP;
 
-public class PartialTrackedSinglePartitionRead extends AbstractPartialTrackedRead
+public class PartialTrackedSinglePartitionRead extends PartialTrackedRead
 {
     private final Index.Searcher searcher;
     private final SinglePartitionReadCommand command;
@@ -74,7 +74,7 @@ public class PartialTrackedSinglePartitionRead extends AbstractPartialTrackedRea
         }
 
         @Override
-        public State augment(PartitionUpdate update)
+        public void augment(PartitionUpdate update)
         {
             if (!update.partitionKey().equals(command.partitionKey()))
                 throw new IllegalArgumentException(String.format("Received update for partition key %s but command was for %s",
@@ -85,7 +85,6 @@ public class PartialTrackedSinglePartitionRead extends AbstractPartialTrackedRea
                 augmentedData = new SimpleBTreePartition(command.partitionKey(), command.metadata(), UpdateTransaction.NO_OP);
 
             augmentedData.update(update);
-            return this;
         }
 
         @Override
@@ -95,7 +94,7 @@ public class PartialTrackedSinglePartitionRead extends AbstractPartialTrackedRea
         }
     }
 
-    private class SinglePartitionCompleted extends Completed
+    private class SinglePartitionCompleted extends AbstractCompleted
     {
         private final UnfilteredPartitionIterator initialData;
         private final SimpleBTreePartition augmentedData;

src/java/org/apache/cassandra/service/reads/tracked/TrackedLocalReads.java

@@ -22,6 +22,8 @@ import java.util.Map;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.cassandra.concurrent.Stage;
 import org.apache.cassandra.db.*;
@@ -37,12 +39,8 @@ import org.apache.cassandra.service.reads.SpeculativeRetryPolicy;
 import org.apache.cassandra.tcm.ClusterMetadata;
 import org.apache.cassandra.transport.Dispatcher;
 import org.apache.cassandra.utils.concurrent.AsyncPromise;
-import org.apache.cassandra.utils.concurrent.Future;
 import org.jctools.maps.NonBlockingHashMap;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 /**
  * Since the read reconciliations don't use 2 way callbacks, maps of active reads and reconciliations
  * are maintained and expired here.
@@ -53,6 +51,12 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
 {
     private static final Logger logger = LoggerFactory.getLogger(TrackedLocalReads.class);
 
+    public interface Completer
+    {
+        void complete(AsyncPromise<TrackedDataResponse> promise, PartialTrackedRead read, ConsistencyLevel consistencyLevel, Dispatcher.RequestTime requestTime);
+        TrackedLocalReads.Completer DEFAULT = (promise, read, consistencyLevel, requestTime) -> read.complete(promise, consistencyLevel, requestTime);
+    }
+
     private final NonBlockingHashMap<TrackedRead.Id, Coordinator> coordinators = new NonBlockingHashMap<>();
 
     public TrackedLocalReads()
@@ -66,7 +70,8 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
         ReadCommand command,
         ConsistencyLevel consistencyLevel,
         int[] summaryNodes,
-        Dispatcher.RequestTime requestTime)
+        Dispatcher.RequestTime requestTime,
+        TrackedLocalReads.Completer completer)
     {
         Keyspace keyspace = Keyspace.open(command.metadata().keyspace);
         ColumnFamilyStore cfs = keyspace.getColumnFamilyStore(command.metadata().id);
@@ -96,18 +101,18 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
         }
         // TODO: confirm all summaryNodes are present in the replica plan
         AsyncPromise<TrackedDataResponse> promise = new AsyncPromise<>();
-        beginReadInternal(readId, command, replicaPlan, summaryNodes, requestTime, promise);
+        beginReadInternal(readId, command, replicaPlan, summaryNodes, requestTime, promise, completer);
         return promise;
     }
 
     // TODO (expected): skip local summaries and reconcile when summaryNodes is empty (e.g. for CL.ONE)
-    private void beginReadInternal(
-        TrackedRead.Id readId,
-        ReadCommand command,
-        ReplicaPlan.AbstractForRead<?, ?> replicaPlan,
-        int[] summaryNodes,
-        Dispatcher.RequestTime requestTime,
-        AsyncPromise<TrackedDataResponse> promise)
+    private void beginReadInternal(TrackedRead.Id readId,
+                                   ReadCommand command,
+                                   ReplicaPlan.AbstractForRead<?, ?> replicaPlan,
+                                   int[] summaryNodes,
+                                   Dispatcher.RequestTime requestTime,
+                                   AsyncPromise<TrackedDataResponse> promise,
+                                   TrackedLocalReads.Completer completer)
     {
         PartialTrackedRead read = null;
         MutationSummary secondarySummary;
@@ -118,6 +123,7 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
         try
         {
             read = command.beginTrackedRead(controller);
+            read.setFollowUpReadContext(replicaPlan.consistencyLevel(), requestTime);
             // Create another summary once initial data has been read fully. We do this to catch
             // any mutations that may have arrived during initial read execution.
             secondarySummary = command.createMutationSummary(true);
@@ -131,8 +137,7 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
             throw e;
         }
 
-        Coordinator coordinator =
-                new Coordinator(readId, promise, read, replicaPlan.consistencyLevel(), requestTime);
+        Coordinator coordinator = new Coordinator(readId, promise, read, replicaPlan.consistencyLevel(), requestTime, completer);
         coordinators.put(readId, coordinator);
 
         // TODO (expected): reconsider the approach to tracked mutation metrics
@@ -185,19 +190,21 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
         private final PartialTrackedRead read;
         private final ConsistencyLevel consistencyLevel;
         private final Dispatcher.RequestTime requestTime;
+        private final Completer completer;
 
-        Coordinator(
-                TrackedRead.Id readId,
-                AsyncPromise<TrackedDataResponse> promise,
-                PartialTrackedRead read,
-                ConsistencyLevel consistencyLevel,
-                Dispatcher.RequestTime requestTime)
+        Coordinator(TrackedRead.Id readId,
+                    AsyncPromise<TrackedDataResponse> promise,
+                    PartialTrackedRead read,
+                    ConsistencyLevel consistencyLevel,
+                    Dispatcher.RequestTime requestTime,
+                    Completer completer)
         {
             this.readId = readId;
             this.promise = promise;
             this.read = Preconditions.checkNotNull(read);
             this.consistencyLevel = consistencyLevel;
             this.requestTime = requestTime;
+            this.completer = completer;
         }
 
         boolean isPurgeable(long nanoTime)
@@ -215,52 +222,20 @@ public class TrackedLocalReads implements ExpiredStatePurger.Expireable
         {
             logger.trace("Reconciliation completed for {}, missing {}", readId, augmentingOffsets);
 
-            Stage.READ.submit(() -> { 
-
+            Stage.READ.submit(() -> {
                 try
                 {
                     read.augment(augmentingOffsets);
-                    complete();
-                } catch (Throwable t) {
+                    // TODO: Does the implementation of this in FollowUpRead in PartialTrackedIndexRead need to be on a READ Stage thread?
+                    completer.complete(promise, read, consistencyLevel, requestTime);
+                }
+                catch (Throwable t)
+                {
                     logger.error("Exception thrown during read completion", t);
                     promise.tryFailure(t);
                     throw t;
                 }
             });
         }
-
-        private void complete()
-        {
-            try (PartialTrackedRead.CompletedRead completedRead = read.complete())
-            {
-                TrackedDataResponse response = completedRead.response();
-                Future<TrackedDataResponse> followUp = completedRead.followupRead(response, consistencyLevel, requestTime);
-
-                if (followUp != null)
-                {
-                    followUp.addCallback((newResponse, error) -> {
-                        if (error != null)
-                        {
-                            promise.tryFailure(error);
-                            return;
-                        }
-                        promise.trySuccess(newResponse);
-                    });
-                }
-                else
-                {
-                    promise.trySuccess(response);
-                }
-            }
-            catch (Exception e)
-            {
-                promise.tryFailure(e);
-                throw e;
-            }
-            finally
-            {
-                read.close();
-            }
-        }
     }
 }

src/java/org/apache/cassandra/service/reads/tracked/TrackedRead.java

@@ -18,8 +18,19 @@
 
 package org.apache.cassandra.service.reads.tracked;
 
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.cassandra.concurrent.Stage;
 import org.apache.cassandra.db.*;
@@ -47,18 +58,6 @@ import org.apache.cassandra.utils.concurrent.AsyncPromise;
 import org.apache.cassandra.utils.concurrent.Future;
 import org.apache.cassandra.utils.concurrent.UncheckedInterruptedException;
 
-import java.io.IOException;
-import java.util.Collections;
-import java.util.Map;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicLong;
-import java.util.function.Consumer;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 import static org.apache.cassandra.metrics.ClientRequestsMetricsHolder.readMetrics;
 
@@ -249,7 +248,7 @@ public abstract class TrackedRead<E extends Endpoints<E>, P extends ReplicaPlan.
         return metadata.directory.peerId(replica.endpoint()).id();
     }
 
-    private void start(Dispatcher.RequestTime requestTime, Consumer<PartialTrackedRead> partialReadConsumer)
+    private void start(Dispatcher.RequestTime requestTime, Consumer<PartialTrackedRead> partialReadConsumer, TrackedLocalReads.Completer completer)
     {
         // TODO: skip local coordination if this node knows its recovering from an outage
         // TODO: read speculation
@@ -279,7 +278,7 @@ public abstract class TrackedRead<E extends Endpoints<E>, P extends ReplicaPlan.
             logger.trace("Locally coordinating {}", readId);
             Stage.READ.submit(() -> {
                 AsyncPromise<TrackedDataResponse> promise =
-                    MutationTrackingService.instance.localReads().beginRead(readId, ClusterMetadata.current(), command, consistencyLevel, summaryNodes, requestTime);
+                    MutationTrackingService.instance.localReads().beginRead(readId, ClusterMetadata.current(), command, consistencyLevel, summaryNodes, requestTime, completer);
                 promise.addCallback((response, error) -> {
                     if (error != null)
                     {
@@ -326,12 +325,12 @@ public abstract class TrackedRead<E extends Endpoints<E>, P extends ReplicaPlan.
 
     public void start(Dispatcher.RequestTime requestTime)
     {
-        start(requestTime, null);
+        start(requestTime, null, TrackedLocalReads.Completer.DEFAULT);
     }
 
-    public void startLocal(Dispatcher.RequestTime requestTime, Consumer<PartialTrackedRead> partialReadConsumer)
+    public void startLocal(Dispatcher.RequestTime requestTime, Consumer<PartialTrackedRead> partialReadConsumer, TrackedLocalReads.Completer completer)
     {
-        start(requestTime, partialReadConsumer);
+        start(requestTime, partialReadConsumer, completer);
     }
 
     private void onResponse(TrackedDataResponse response)
@@ -450,7 +449,7 @@ public abstract class TrackedRead<E extends Endpoints<E>, P extends ReplicaPlan.
             AsyncPromise<TrackedDataResponse> promise =
                 MutationTrackingService.instance
                                        .localReads()
-                                       .beginRead(readId, metadata, command, consistencyLevel, summaryNodes, requestTime);
+                                       .beginRead(readId, metadata, command, consistencyLevel, summaryNodes, requestTime, TrackedLocalReads.Completer.DEFAULT);
             promise.addCallback((response, error) -> {
                 if (error != null)
                 {

src/java/org/apache/cassandra/utils/AbstractIterator.java

@@ -18,12 +18,9 @@
 */
 package org.apache.cassandra.utils;
 
-import java.util.Iterator;
 import java.util.NoSuchElementException;
 
-import com.google.common.collect.PeekingIterator;
-
-public abstract class AbstractIterator<V> implements Iterator<V>, PeekingIterator<V>, CloseableIterator<V>
+public abstract class AbstractIterator<V> implements CloseablePeekingIterator<V>
 {
 
     private static enum State { MUST_FETCH, HAS_NEXT, DONE, FAILED }

src/java/org/apache/cassandra/utils/CloseablePeekingIterator.java

@@ -0,0 +1,40 @@
+/*
+ * 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.utils;
+
+import java.util.Iterator;
+
+import com.google.common.collect.PeekingIterator;
+
+public interface CloseablePeekingIterator<V> extends PeekingIterator<V>, CloseableIterator<V>
+{
+    static <V> CloseablePeekingIterator<V> wrap(Iterator<V> iterator)
+    {
+        return new AbstractIterator<>()
+        {
+            @Override
+            protected V computeNext()
+            {
+                if (!iterator.hasNext())
+                    return endOfData();
+                return iterator.next();
+            }
+        };
+    }
+}

test/distributed/org/apache/cassandra/distributed/test/ReadRepairEmptyRangeTombstonesTestBase.java

@@ -81,9 +81,12 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
         for (int coordinator = 1; coordinator <= NUM_NODES; coordinator++)
             for (boolean paging : BOOLEANS)
                 for (boolean reverse : BOOLEANS)
-                    for (ReplicationType replication : ReplicationType.fixmeValues())
-                        result.add(new Object[]{ ReadRepairStrategy.BLOCKING, coordinator, paging, reverse, replication });
-        result.add(new Object[]{ ReadRepairStrategy.NONE, 1, false, false, ReplicationType.untracked });
+                    for (ReplicationType replication : ReplicationType.values())
+                    {
+                        ReadRepairStrategy rrStrategy = replication == ReplicationType.untracked ? ReadRepairStrategy.BLOCKING : ReadRepairStrategy.NONE;
+                        result.add(new Object[]{ rrStrategy, coordinator, paging, reverse, replication });
+                    }
+
         return result;
     }
 
@@ -170,6 +173,11 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
     @Test
     public void testRangeQueriesWithRowsOvetrlappingWithTombstoneRangeStart()
     {
+        // Read-repair covers only the queried range, while mutation tracking pulls all missing mutations.
+        Object[][] postRepairInternalRows = replicationType == ReplicationType.tracked
+                                            ? new Object[][] { row(1), row(2), row(3), row(4), row(5), row(6) }
+                                            : new Object[][] { row(1), row(2), row(3), row(4), row(5) };
+
         tester().createTable("CREATE TABLE %s(k int, c int, PRIMARY KEY (k, c)) " +
                              "WITH CLUSTERING ORDER BY (c %s) AND read_repair='%s'")
                 .mutate(1, "DELETE FROM %s USING TIMESTAMP 1 WHERE k=0 AND c>=3 AND c<=6")
@@ -185,7 +193,7 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
                 .assertRowsDistributed("SELECT c FROM %s WHERE k=0 AND c>=2 AND c<=5",
                                        1,
                                        row(2), row(3), row(4), row(5))
-                .assertRowsInternal("SELECT c FROM %s", row(1), row(2), row(3), row(4), row(5))
+                .assertRowsInternal("SELECT c FROM %s", postRepairInternalRows)
                 .mutate(2, "DELETE FROM %s WHERE k=0 AND c>=1 AND c<=6")
                 .assertRowsInternal("SELECT * FROM %s");
     }
@@ -196,6 +204,11 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
     @Test
     public void testRangeQueriesWithRowsOverlappingWithTombstoneRangeEnd()
     {
+        // Read-repair covers only the queried range, while mutation tracking pulls all missing mutations.
+        Object[][] postRepairInternalRows = replicationType == ReplicationType.tracked 
+                                            ? new Object[][] { row(1), row(2), row(3), row(4), row(5), row(6) } 
+                                            : new Object[][] { row(2), row(3), row(4), row(5), row(6) };
+
         tester().createTable("CREATE TABLE %s(k int, c int, PRIMARY KEY (k, c)) " +
                              "WITH CLUSTERING ORDER BY (c %s) AND read_repair='%s'")
                 .mutate(1, "DELETE FROM %s USING TIMESTAMP 1 WHERE k=0 AND c>=1 AND c<=4")
@@ -211,7 +224,7 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
                 .assertRowsDistributed("SELECT c FROM %s WHERE k=0 AND c>=3 AND c<=6",
                                        1,
                                        row(3), row(4), row(5), row(6))
-                .assertRowsInternal("SELECT c FROM %s", row(2), row(3), row(4), row(5), row(6))
+                .assertRowsInternal("SELECT c FROM %s", postRepairInternalRows)
                 .mutate(2, "DELETE FROM %s WHERE k=0 AND c>=1 AND c<=6")
                 .assertRowsInternal("SELECT * FROM %s");
     }
@@ -222,6 +235,15 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
     @Test
     public void testPointQueriesWithRowsContainedInTombstoneRange()
     {
+        // Read-repair covers only the queried range, while mutation tracking pulls all missing mutations.
+        Object[][] postRepairInternalRows = replicationType == ReplicationType.tracked 
+                                            ? new Object[][] { row(0, 0), row(0, 1), row(0, 2) }
+                                            : new Object[][] { row(0, 1), row(0, 2) };
+
+        @SuppressWarnings("ZeroLengthArrayAllocation") Object[][] postDeleteInternalRows = replicationType == ReplicationType.tracked 
+                                                                                           ? new Object[][] { row(0, 0) }
+                                                                                           : new Object[0][0];
+
         tester().createTable("CREATE TABLE %s(k int, c int, PRIMARY KEY (k, c)) " +
                              "WITH CLUSTERING ORDER BY (c %s) AND read_repair='%s'")
                 .mutate(1, "DELETE FROM %s USING TIMESTAMP 1 WHERE k=0 AND c>0 AND c<3")
@@ -231,9 +253,9 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
                 .assertRowsDistributed("SELECT * FROM %s WHERE k=0 AND c=1", 1, row(0, 1))
                 .assertRowsDistributed("SELECT * FROM %s WHERE k=0 AND c=2", 1, row(0, 2))
                 .assertRowsDistributed("SELECT * FROM %s WHERE k=0 AND c=3", 0)
-                .assertRowsInternal("SELECT * FROM %s", row(0, 1), row(0, 2))
+                .assertRowsInternal("SELECT * FROM %s", postRepairInternalRows)
                 .mutate(2, "DELETE FROM %s WHERE k=0 AND c>0 AND c<3")
-                .assertRowsInternal("SELECT * FROM %s");
+                .assertRowsInternal("SELECT * FROM %s", postDeleteInternalRows);
     }
 
     /**
@@ -275,7 +297,7 @@ public abstract class ReadRepairEmptyRangeTombstonesTestBase extends TestBaseImp
         {
             String formattedQuery = String.format(query, qualifiedTableName);
 
-            if (strategy == ReadRepairStrategy.NONE)
+            if (strategy == ReadRepairStrategy.NONE && replicationType == ReplicationType.untracked)
                 expectedRows = EMPTY_ROWS;
             else if (reverse)
                 expectedRows = reverse(expectedRows);

test/distributed/org/apache/cassandra/distributed/test/ReadRepairIndexTest.java

@@ -0,0 +1,227 @@
+/*
+ * 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.distributed.test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+import org.junit.AfterClass;
+import org.junit.Assume;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.index.internal.CassandraIndex;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.schema.ReplicationType;
+import org.apache.cassandra.service.reads.repair.ReadRepairStrategy;
+
+import static org.apache.cassandra.distributed.shared.AssertUtils.row;
+
+@RunWith(Parameterized.class)
+public class ReadRepairIndexTest extends TestBaseImpl
+{
+    private static final int NUM_NODES = 2;
+
+    enum IndexType
+    {
+        SECONDARY(CassandraIndex.NAME),
+        SAI(StorageAttachedIndex.NAME);
+
+        final String name;
+
+        IndexType(String name)
+        {
+            this.name = name;
+        }
+
+        @Override
+        public String toString()
+        {
+            return name;
+        }
+    }
+
+    /**
+     * The read repair strategy to be used
+     */
+    @Parameterized.Parameter
+    public ReadRepairStrategy strategy;
+
+    /**
+     * The node to be used as coordinator
+     */
+    @Parameterized.Parameter(1)
+    public int coordinator;
+
+    /**
+     * Whether to flush data after mutations
+     */
+    @Parameterized.Parameter(2)
+    public boolean flush;
+
+    /**
+     * Whether paging is used for the distributed queries
+     */
+    @Parameterized.Parameter(3)
+    public boolean paging;
+
+    @Parameterized.Parameter(4)
+    public ReplicationType replicationType;
+
+    @SuppressWarnings("ClassEscapesDefinedScope")
+    @Parameterized.Parameter(5)
+    public IndexType indexType;
+
+
+    @Parameterized.Parameters(name = "{index}: strategy={0} coordinator={1} flush={2} paging={3} replication={4} index={5}")
+    public static Collection<Object[]> data()
+    {
+        List<Object[]> result = new ArrayList<>();
+        for (int coordinator = 1; coordinator <= NUM_NODES; coordinator++)
+            for (boolean flush : BOOLEANS)
+                for (boolean paging : BOOLEANS)
+                    for (ReplicationType replication : ReplicationType.values())
+                        for (IndexType indexType : IndexType.values())
+                            result.add(new Object[]{ ReadRepairStrategy.BLOCKING, coordinator, flush, paging, replication, indexType});
+        return result;
+    }
+
+    private static Cluster cluster;
+
+    @BeforeClass
+    public static void setupCluster() throws IOException
+    {
+        cluster = Cluster.build(NUM_NODES)
+                              .withConfig(config -> config.set("read_request_timeout", "1m")
+                                                          .set("write_request_timeout", "1m"))
+                              .start();
+    }
+
+    @AfterClass
+    public static void teardownCluster()
+    {
+        if (cluster != null)
+            cluster.close();
+    }
+
+    protected Tester tester(String restriction)
+    {
+        return new Tester(restriction, cluster, strategy, coordinator, flush, paging, replicationType, indexType);
+    }
+
+    protected static class Tester extends ReadRepairQueryTester.AbstractTester<Tester>
+    {
+        private int nameSeq = 0;
+        private final IndexType indexType;
+
+        @SuppressWarnings("ClassEscapesDefinedScope")
+        public Tester(String restriction, Cluster cluster, ReadRepairStrategy strategy, int coordinator, boolean flush, boolean paging, ReplicationType replicationType, IndexType indexType)
+        {
+            super(restriction, cluster, strategy, coordinator, flush, paging, replicationType);
+            this.indexType = indexType;
+        }
+
+        @Override
+        Tester self()
+        {
+            return this;
+        }
+
+        Tester createIndex(String column)
+        {
+            String query = String.format("CREATE INDEX %s_index_%d ON %s(%s) USING '%s'", tableName, nameSeq++, qualifiedTableName, column, indexType);
+            cluster.schemaChange(query);
+            return this;
+        }
+    }
+
+    /**
+     * A partition that would not be an index hit on one node would be on the other
+     */
+    @Test
+    public void singlePartitionUpdatedPartition()
+    {
+        tester("WHERE k=1 AND v=2")
+        .createTable("CREATE TABLE %s (k int, c int, v int, PRIMARY KEY (k, c))")
+        .createIndex("v")
+        .mutate(2, "INSERT INTO %s (k, c, v) VALUES (1, 2, 2)")
+        .mutate(1, "INSERT INTO %s (k, c, v) VALUES (1, 1, 1)")
+        .queryColumns("k, c, v", 1, 0,
+                      rows(row(1, 2, 2)),
+                      rows(row(1, 2, 2)),
+                      rows(row(1, 2, 2)))
+        .tearDown(1,
+                  rows(row(1, 1, 1), row(1, 2, 2)),
+                  rows(row(1, 2, 2)));
+
+    }
+    @Test
+    public void rangeReadTest()
+    {
+        tester("WHERE v=2")
+        .createTable("CREATE TABLE %s (k int, v int, PRIMARY KEY (k))")
+        .createIndex("v")
+        .mutate(2, "INSERT INTO %s (k, v) VALUES (1, 2)")
+        .mutate(1, "INSERT INTO %s (k, v) VALUES (2, 1)")
+        .mutate(2, "INSERT INTO %s (k, v) VALUES (3, 1)")
+        .mutate(1, "INSERT INTO %s (k, v) VALUES (4, 2)")
+        .queryColumns("k, v", 2, 0,
+                      rows(row(1, 2), row(4, 2)),
+                      rows(row(1, 2), row(4, 2)),
+                      rows(row(1, 2), row(4, 2)))
+        .tearDown(2,
+                  rows(row(1, 2), row(2, 1), row(4, 2), row(3, 1)),
+                  (replicationType.isTracked()
+                   ? rows(row(1, 2), row(2, 1), row(4, 2), row(3, 1))
+                   : rows(row(1, 2), row(2, 1), row(4, 2))),
+                  (replicationType.isTracked()
+                   ? rows(row(1, 2), row(2, 1), row(4, 2), row(3, 1))
+                   : rows(row(1, 2), row(4, 2), row(3, 1))));
+    }
+
+    @Test
+    public void sortedRangeRead()
+    {
+        Assume.assumeTrue("CassandraIndex doesn't support numerical ranges", indexType == IndexType.SAI);
+
+        tester("WHERE v>2")
+        .createTable("CREATE TABLE %s (k int, c int, v int, PRIMARY KEY (k, c))")
+        .createIndex("v")
+        .mutate(2, "INSERT INTO %s (k, c, v) VALUES (1, 2, 2)")
+        .mutate(1, "INSERT INTO %s (k, c, v) VALUES (1, 4, 4)")
+        .mutate(2, "INSERT INTO %s (k, c, v) VALUES (5, 2, 1)")
+        .mutate(1, "INSERT INTO %s (k, c, v) VALUES (8, 4, 3)")
+        .queryColumns("k, c, v", 2, 0,
+                      rows(row(1, 4, 4), row(8, 4, 3)),
+                      rows(row(1, 4, 4), row(8, 4, 3)),
+                      rows(row(1, 4, 4), row(8, 4, 3)))
+        .tearDown(2,
+                  rows(row(5, 2, 1), row(1, 2, 2), row(1, 4, 4), row(8, 4, 3)),
+                  replicationType.isTracked()
+                    ? rows(row(5, 2, 1), row(1, 2, 2), row(1, 4, 4), row(8, 4, 3))
+                    : rows(row(1, 4, 4), row(8, 4, 3)),
+                  rows(row(5, 2, 1), row(1, 2, 2), row(1, 4, 4), row(8, 4, 3)));
+
+    }
+}

test/distributed/org/apache/cassandra/distributed/test/ReadRepairQueryTester.java

@@ -34,6 +34,7 @@ import org.apache.cassandra.service.reads.repair.ReadRepairStrategy;
 
 import static org.apache.cassandra.distributed.shared.AssertUtils.assertEquals;
 import static org.apache.cassandra.distributed.shared.AssertUtils.assertRows;
+import static org.apache.cassandra.service.reads.repair.ReadRepairStrategy.BLOCKING;
 import static org.apache.cassandra.service.reads.repair.ReadRepairStrategy.NONE;
 
 /**
@@ -109,8 +110,8 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
             for (boolean flush : BOOLEANS)
                 for (boolean paging : BOOLEANS)
                     for (ReplicationType replication : ReplicationType.values())
-                        result.add(new Object[]{ ReadRepairStrategy.BLOCKING, coordinator, flush, paging, replication });
-        result.add(new Object[]{ ReadRepairStrategy.NONE, 1, false, false, ReplicationType.untracked });
+                        result.add(new Object[] { replication.isTracked() ? NONE : BLOCKING, coordinator, flush, paging, replication });
+
         return result;
     }
 
@@ -137,12 +138,12 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
         return new Tester(restriction, cluster, strategy, coordinator, flush, paging, replicationType);
     }
 
-    protected static class Tester extends ReadRepairTester<Tester>
+    static abstract class AbstractTester<T extends AbstractTester<T>> extends ReadRepairTester<T>
     {
         private final String restriction; // the tested CQL query WHERE restriction
         private final String allColumnsQuery; // a SELECT * query for the table using the tested restriction
 
-        Tester(String restriction, Cluster cluster, ReadRepairStrategy strategy, int coordinator, boolean flush, boolean paging, ReplicationType replicationType)
+        AbstractTester(String restriction, Cluster cluster, ReadRepairStrategy strategy, int coordinator, boolean flush, boolean paging, ReplicationType replicationType)
         {
             super(cluster, strategy, coordinator, flush, paging, false, replicationType);
             this.restriction = restriction;
@@ -150,12 +151,6 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
             allColumnsQuery = String.format("SELECT * FROM %s %s", qualifiedTableName, restriction);
         }
 
-        @Override
-        Tester self()
-        {
-            return this;
-        }
-
         /**
          * Runs the tested query with CL=ALL selectig only the specified columns and verifies that it returns the
          * specified rows. Then, it runs the query again selecting all the columns, and verifies that the first query
@@ -168,12 +163,12 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
          * @param node1Rows                the rows in the first node, which is the one with the most updated data
          * @param node2Rows                the rows in the second node, which is the one meant to receive the RR writes
          */
-        Tester queryColumns(String columns,
-                            long columnsQueryRepairedRows,
-                            long rowsQueryRepairedRows,
-                            Object[][] columnsQueryResults,
-                            Object[][] node1Rows,
-                            Object[][] node2Rows)
+        T queryColumns(String columns,
+                       long columnsQueryRepairedRows,
+                       long rowsQueryRepairedRows,
+                       Object[][] columnsQueryResults,
+                       Object[][] node1Rows,
+                       Object[][] node2Rows)
         {
             // query only the selected columns with CL=ALL to trigger partial read repair on that column
             String columnsQuery = String.format("SELECT %s FROM %s %s", columns, qualifiedTableName, restriction);
@@ -206,7 +201,7 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
             assertRowsDistributed(columnsQuery, columnsQueryRepairedRows, columnsQueryResults);
 
             // query entire rows to repair the rest of the columns, that might trigger new repairs for those columns
-            return verifyQuery(allColumnsQuery, rowsQueryRepairedRows, node1Rows, node2Rows);
+            return verifyQuery(allColumnsQuery, rowsQueryRepairedRows, node1Rows, node1Rows, node2Rows);
         }
 
         /**
@@ -223,13 +218,13 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
          * @param node1Rows                the rows in the first node, which is the one with the most updated data
          * @param node2Rows                the rows in the second node, which is the one meant to receive the RR writes
          */
-        Tester deleteColumn(String columnDeletion,
-                            String columns,
-                            long columnsQueryRepairedRows,
-                            long rowsQueryRepairedRows,
-                            Object[][] columnsQueryResults,
-                            Object[][] node1Rows,
-                            Object[][] node2Rows)
+        T deleteColumn(String columnDeletion,
+                       String columns,
+                       long columnsQueryRepairedRows,
+                       long rowsQueryRepairedRows,
+                       Object[][] columnsQueryResults,
+                       Object[][] node1Rows,
+                       Object[][] node2Rows)
         {
             assert restriction != null;
 
@@ -250,7 +245,7 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
          * Executes the specified row deletion on just one node and verifies the tested query, to ensure that the tested
          * query propagates the row deletion.
          */
-        Tester deleteRows(String rowDeletion, long repairedRows, Object[][] node1Rows, Object[][] node2Rows)
+        T deleteRows(String rowDeletion, long repairedRows, Object[][] node1Rows, Object[][] node2Rows)
         {
             mutate(1, rowDeletion);
 
@@ -259,28 +254,28 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
             if (replicationType.isTracked())
                 repairedRows = Math.min(repairedRows, 1);
 
-            return verifyQuery(allColumnsQuery, repairedRows, node1Rows, node2Rows);
+            return verifyQuery(allColumnsQuery, repairedRows, node1Rows, node1Rows, node2Rows);
         }
 
-        Tester mutate(String... queries)
+        T mutate(String... queries)
         {
             return mutate(1, queries);
         }
 
-        private Tester verifyQuery(String query, long expectedRepairedRows, Object[][] node1Rows, Object[][] node2Rows)
+        private T verifyQuery(String query, long expectedRepairedRows, Object[][] allRows, Object[][] node1Rows, Object[][] node2Rows)
         {
             // verify the per-replica status before running the query distributedly
             assertRows(cluster.get(1).executeInternal(query), node1Rows);
-            assertRows(cluster.get(2).executeInternal(query), strategy == NONE ? EMPTY_ROWS : node2Rows);
+            assertRows(cluster.get(2).executeInternal(query), strategy == NONE && !replicationType.isTracked() ? EMPTY_ROWS : node2Rows);
 
             // now, run the query with CL=ALL to reconcile and repair the replicas
-            assertRowsDistributed(query, expectedRepairedRows, node1Rows);
+            assertRowsDistributed(query, expectedRepairedRows, allRows);
 
             // run the query locally again to verify that the distributed query has repaired everything
-            assertRows(cluster.get(1).executeInternal(query), node1Rows);
-            assertRows(cluster.get(2).executeInternal(query), strategy == NONE ? EMPTY_ROWS : node1Rows);
+            assertRows(cluster.get(1).executeInternal(query), allRows);
+            assertRows(cluster.get(2).executeInternal(query), strategy == NONE && !replicationType.isTracked() ? EMPTY_ROWS : allRows);
 
-            return this;
+            return self();
         }
 
         /**
@@ -295,11 +290,11 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
          * Verifies the final status of the nodes with an unrestricted query, to ensure that the main tested query
          * hasn't triggered any unexpected repairs. Then, it verifies that the node that hasn't been used as coordinator
          * hasn't triggered any unexpected repairs. Finally, it drops the table.
-         *
+         * <p>
          * The expectUnrepaired flag is meant for range query tests where logged replication table special casing
          * doesn't apply since we do expect the final query to find and repair missing mutations
          */
-        void tearDown(long repairedRows, Object[][] node1Rows, Object[][] node2Rows, boolean expectUnrepaired)
+        void tearDown(long repairedRows, Object[][] allRows, Object[][] node1Rows, Object[][] node2Rows, boolean expectUnrepaired)
         {
             if (replicationType.isTracked() && !expectUnrepaired)
             {
@@ -321,7 +316,7 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
                 // we also expect all pending mutations to be reconciled in the initial read, and none to be reconciled on the verification step
                 repairedRows = 0;
             }
-            verifyQuery("SELECT * FROM " + qualifiedTableName, repairedRows, node1Rows, node2Rows);
+            verifyQuery("SELECT * FROM " + qualifiedTableName, repairedRows, allRows, node1Rows, node2Rows);
             for (int n = 1; n <= cluster.size(); n++)
             {
                 if (n == coordinator)
@@ -337,7 +332,37 @@ public abstract class ReadRepairQueryTester extends TestBaseImpl
 
         void tearDown(long repairedRows, Object[][] node1Rows, Object[][] node2Rows)
         {
-            tearDown(repairedRows, node1Rows, node2Rows, false);
+            tearDown(repairedRows, node1Rows, node1Rows, node2Rows, false);
+        }
+
+        void tearDown(long repairedRows, Object[][] allRows, Object[][] node1Rows, Object[][] node2Rows)
+        {
+            tearDown(repairedRows, allRows, node1Rows, node2Rows, false);
+        }
+
+        void tearDown(long repairedRows, Object[][] node1Rows, Object[][] node2Rows, boolean expectUnrepaired)
+        {
+            tearDown(repairedRows, node1Rows, node1Rows, node2Rows, expectUnrepaired);
+        }
+    }
+
+    void tearDown(long repairedRows, Object[][] node1Rows, Object[][] node2Rows, boolean expectUnrepaired)
+    {
+
+    }
+
+
+    protected static class Tester extends AbstractTester<Tester>
+    {
+        public Tester(String restriction, Cluster cluster, ReadRepairStrategy strategy, int coordinator, boolean flush, boolean paging, ReplicationType replicationType)
+        {
+            super(restriction, cluster, strategy, coordinator, flush, paging, replicationType);
+        }
+
+        @Override
+        Tester self()
+        {
+            return this;
         }
     }
 }

test/distributed/org/apache/cassandra/distributed/test/ReadRepairTester.java

@@ -36,7 +36,6 @@ import org.apache.cassandra.service.reads.repair.ReadRepairStrategy;
 
 import static org.apache.cassandra.distributed.api.ConsistencyLevel.ALL;
 import static org.apache.cassandra.distributed.shared.AssertUtils.assertEquals;
-import static org.apache.cassandra.distributed.test.TestBaseImpl.KEYSPACE;
 
 /**
  * Extensible helper class for read repair tests.
@@ -46,7 +45,7 @@ public abstract class ReadRepairTester<T extends ReadRepairTester<T>>
     private static final AtomicInteger seqNumber = new AtomicInteger();
 
     private final String keyspaceName = "ks_" + seqNumber.getAndIncrement();
-    private static final String tableName = "tbl";
+    static final String tableName = "tbl";
     final String qualifiedTableName = keyspaceName + '.' + tableName;
 
     protected final Cluster cluster;
@@ -117,7 +116,7 @@ public abstract class ReadRepairTester<T extends ReadRepairTester<T>>
 
         // flush the update node to ensure reads come from sstables
         if (flush)
-            cluster.get(node).flush(KEYSPACE);
+            cluster.get(node).flush(keyspaceName);
 
         return self();
     }

test/distributed/org/apache/cassandra/distributed/test/cql3/MultiNodeTableWalkWithMutationTrackingTest.java

@@ -38,6 +38,7 @@ import org.apache.cassandra.utils.LoggingCommand;
 
 import static accord.utils.Property.commands;
 import static accord.utils.Property.stateful;
+import static org.apache.cassandra.cql3.KnownIssue.AF_MULTI_NODE_MULTI_COLUMN_AND_NODE_LOCAL_WRITES;
 
 public class MultiNodeTableWalkWithMutationTrackingTest extends MultiNodeTableWalkBase
 {
@@ -74,23 +75,27 @@ public class MultiNodeTableWalkWithMutationTrackingTest extends MultiNodeTableWa
         // if a failing seed is detected, populate here
         // Example: builder.withSeed(42L);
         // builder.withExamples(10);
+        // if a failing seed is detected, populate here
+        // Example: builder.withSeed(42L);
         // CQL operations may have opertors such as +, -, and / (example 4 + 4), to "apply" them to get a constant value
         // CQL_DEBUG_APPLY_OPERATOR = true;
         // When mutations look to be lost as seen by more complex SELECTs, it can be useful to just SELECT the partition/row right after to write to see if it was safe at the time.
         // READ_AFTER_WRITE = true;
     }
 
+    // TODO: Remove this override entirely when range reads and indexing are working properly together.
     @Override
     protected List<CreateIndexDDL.Indexer> supportedIndexers()
     {
-        return Collections.emptyList();
+        return Collections.singletonList(CreateIndexDDL.SAI);
     }
 
     @Override
     protected void clusterConfig(IInstanceConfig c)
     {
         super.clusterConfig(c);
-        c.set("mutation_tracking_enabled", "true");
+        c.set("mutation_tracking_enabled", true);
+        IGNORED_ISSUES.remove(AF_MULTI_NODE_MULTI_COLUMN_AND_NODE_LOCAL_WRITES);
     }
 
     @Test
@@ -100,17 +105,20 @@ public class MultiNodeTableWalkWithMutationTrackingTest extends MultiNodeTableWa
         {
             Property.StatefulBuilder statefulBuilder = stateful().withExamples(10).withSteps(400);
             preCheck(cluster, statefulBuilder);
+
+            // TODO: Uncomment the commented bits below to test range queries w/ the seeds above.
             statefulBuilder.check(commands(() -> rs -> createState(rs, cluster))
                                   .add(StatefulASTBase::insert)
-                                  .add(StatefulASTBase::fullTableScan)
-                                  .addIf(State::allowUsingTimestamp, StatefulASTBase::validateUsingTimestamp)
+//                                  .add(StatefulASTBase::fullTableScan)
+//                                  .addIf(State::allowUsingTimestamp, StatefulASTBase::validateUsingTimestamp)
                                   .addIf(State::hasPartitions, this::selectExisting)
-                                  .addAllIf(State::supportTokens, this::selectToken, this::selectTokenRange, StatefulASTBase::selectMinTokenRange)
+//                                  .addAllIf(State::supportTokens, this::selectToken, this::selectTokenRange, StatefulASTBase::selectMinTokenRange)
                                   .addIf(State::hasEnoughMemtable, StatefulASTBase::flushTable)
                                   .addIf(State::hasEnoughSSTables, StatefulASTBase::compactTable)
-                                  .addIf(State::allowNonPartitionQuery, this::nonPartitionQuery)
-                                  .addIf(State::allowNonPartitionMultiColumnQuery, this::multiColumnQuery)
+//                                  .addIf(State::allowNonPartitionQuery, this::nonPartitionQuery)
+//                                  .addIf(State::allowNonPartitionMultiColumnQuery, this::multiColumnQuery)
                                   .addIf(State::allowPartitionQuery, this::partitionRestrictedQuery)
+                                  .addIf(State::allowPartitionMultiColumnQuery, this::multiColumnPartitionQuery)
                                   .destroyState(State::close)
                                   .commandsTransformer(LoggingCommand.factory())
                                   .onSuccess(onSuccess(logger))

test/distributed/org/apache/cassandra/distributed/test/cql3/SingleNodeTableWalkTest.java

@@ -135,16 +135,10 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
 
     public Property.Command<State, Void, ?> selectExisting(RandomSource rs, State state)
     {
-        NavigableSet<BytesPartitionState.Ref> keys = state.model.partitionKeys();
-        BytesPartitionState.Ref ref = rs.pickOrderedSet(keys);
-        Clustering<ByteBuffer> key = ref.key;
-
         Select.Builder builder = Select.builder().table(state.metadata);
-        ImmutableUniqueList<Symbol> pks = state.model.factory.partitionColumns;
-        ImmutableUniqueList<Symbol> cks = state.model.factory.clusteringColumns;
-        for (Symbol pk : pks)
-            builder.value(pk, key.bufferAt(pks.indexOf(pk)));
+        BytesPartitionState.Ref ref = restrictPartition(rs, state, builder);
 
+        ImmutableUniqueList<Symbol> cks = state.model.factory.clusteringColumns;
         boolean wholePartition = cks.isEmpty() || rs.nextBoolean();
         if (!wholePartition)
         {
@@ -228,16 +222,8 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
 
     public Property.Command<State, Void, ?> partitionRestrictedQuery(RandomSource rs, State state)
     {
-        //TODO (now): remove duplicate logic
-        NavigableSet<BytesPartitionState.Ref> keys = state.model.partitionKeys();
-        BytesPartitionState.Ref ref = rs.pickOrderedSet(keys);
-        Clustering<ByteBuffer> key = ref.key;
-
         Select.Builder builder = Select.builder().table(state.metadata);
-        ImmutableUniqueList<Symbol> pks = state.model.factory.partitionColumns;
-        for (Symbol pk : pks)
-            builder.value(pk, key.bufferAt(pks.indexOf(pk)));
-
+        BytesPartitionState.Ref ref = restrictPartition(rs, state, builder);
 
         List<Symbol> searchableColumns = state.searchableNonPartitionColumns;
         Symbol symbol = rs.pick(searchableColumns);
@@ -269,6 +255,18 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
         return eqSearch(rs, state, symbol, value, builder);
     }
 
+    private static BytesPartitionState.Ref restrictPartition(RandomSource rs, State state, Select.Builder builder)
+    {
+        NavigableSet<BytesPartitionState.Ref> keys = state.model.partitionKeys();
+        BytesPartitionState.Ref ref = rs.pickOrderedSet(keys);
+        Clustering<ByteBuffer> key = ref.key;
+
+        ImmutableUniqueList<Symbol> pks = state.model.factory.partitionColumns;
+        for (Symbol pk : pks)
+            builder.value(pk, key.bufferAt(pks.indexOf(pk)));
+        return ref;
+    }
+
     public Property.Command<State, Void, ?> nonPartitionQuery(RandomSource rs, State state)
     {
         Symbol symbol = rs.pick(state.searchableColumns);
@@ -300,22 +298,31 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
 
     public Property.Command<State, Void, ?> multiColumnQuery(RandomSource rs, State state)
     {
+        Select.Builder builder = Select.builder().table(state.metadata).allowFiltering();
         List<Symbol> allowedColumns = state.multiColumnQueryColumns();
+        return multiColumnQuery(rs, state, builder, null, allowedColumns);
+    }
 
+    public Property.Command<State, Void, ?> multiColumnPartitionQuery(RandomSource rs, State state)
+    {
+        Select.Builder builder = Select.builder().table(state.metadata).allowFiltering();
+        BytesPartitionState.Ref ref = restrictPartition(rs, state, builder);
+        List<Symbol> allowedColumns = state.multiColumnPartitionQueryColumns();
+        return multiColumnQuery(rs, state, builder, ref, allowedColumns);
+    }
+
+    private static Property.Command<State, Void, ?> multiColumnQuery(RandomSource rs, State state, Select.Builder builder, BytesPartitionState.Ref ref, List<Symbol> allowedColumns)
+    {
         if (allowedColumns.size() <= 1)
             throw new IllegalArgumentException("Unable to do multiple column query when there is only a single column");
 
         int numColumns = rs.nextInt(1, allowedColumns.size()) + 1;
-
         List<Symbol> cols = Gens.lists(Gens.pick(allowedColumns)).unique().ofSize(numColumns).next(rs);
 
-        Select.Builder builder = Select.builder().table(state.metadata).allowFiltering();
-
         for (Symbol symbol : cols)
         {
-            TreeMap<ByteBuffer, List<BytesPartitionState.PrimaryKey>> universe = state.model.index(symbol);
+            TreeMap<ByteBuffer, List<BytesPartitionState.PrimaryKey>> universe = ref == null ? state.model.index(symbol) : state.model.index(ref, symbol);
             NavigableSet<ByteBuffer> allowed = Sets.filter(universe.navigableKeySet(), b -> !ByteBufferUtil.EMPTY_BYTE_BUFFER.equals(b));
-            //TODO (now): support
             if (allowed.isEmpty())
                 return Property.ignoreCommand();
             ByteBuffer value = rs.pickOrderedSet(allowed);
@@ -456,6 +463,7 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
                                   .addIf(State::allowNonPartitionMultiColumnQuery, this::multiColumnQuery)
                                   .addIf(State::allowPartitionQuery, this::partitionRestrictedQuery)
                                   .addIf(State::allowClusteringBetweenQuery, this::clusteringBetweenQuery)
+                                  .addIf(State::allowPartitionMultiColumnQuery, this::multiColumnPartitionQuery)
                                   .destroyState(State::close)
                                   .commandsTransformer(LoggingCommand.factory())
                                   .onSuccess(onSuccess(logger))
@@ -661,6 +669,11 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
             return allowNonPartitionQuery() && multiColumnQueryColumns().size() > 1;
         }
 
+        public boolean allowPartitionMultiColumnQuery()
+        {
+            return allowPartitionQuery() && multiColumnPartitionQueryColumns().size() > 1;
+        }
+
         private List<Symbol> multiColumnQueryColumns()
         {
             List<Symbol> allowedColumns = searchableColumns;
@@ -671,6 +684,16 @@ public class SingleNodeTableWalkTest extends StatefulASTBase
             return allowedColumns;
         }
 
+        private List<Symbol> multiColumnPartitionQueryColumns()
+        {
+            List<Symbol> allowedColumns = searchableNonPartitionColumns;
+            if (hasMultiNodeMultiColumnAllowFilteringWithLocalWritesIssue())
+                allowedColumns = nonPkIndexedColumns;
+            if (IGNORED_ISSUES.contains(KnownIssue.SAI_AND_VECTOR_COLUMNS) && !indexes.isEmpty())
+                allowedColumns = allowedColumns.stream().filter(s -> !s.type().isVector()).collect(Collectors.toList());
+            return allowedColumns;
+        }
+
         private boolean hasMultiNodeMultiColumnAllowFilteringWithLocalWritesIssue()
         {
             return isMultiNode() && IGNORED_ISSUES.contains(KnownIssue.AF_MULTI_NODE_MULTI_COLUMN_AND_NODE_LOCAL_WRITES);

test/distributed/org/apache/cassandra/distributed/test/sai/PartialUpdateHandlingTest.java

@@ -42,6 +42,7 @@ import org.apache.cassandra.distributed.Cluster;
 import org.apache.cassandra.distributed.api.ConsistencyLevel;
 import org.apache.cassandra.distributed.test.TestBaseImpl;
 import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.schema.ReplicationType;
 
 import static org.apache.cassandra.distributed.api.ConsistencyLevel.ALL;
 import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
@@ -70,7 +71,7 @@ import static org.junit.Assert.assertEquals;
  *     across different partitions).
  * 5.) Whether data resides in SSTables or Memtables. The latter is implicitly unrepaired.
  * 6.) Interaction w/ existing mechanisms on the distributed read path that deal with short reads, replica filtering
- *     protection, etc.
+ *     protection, mutation tracking, etc.
  * 7.) The relationship between columns selected and columns restricted by queries. (If coordinator filtering is
  *     involved at the implementation level, retrieving enough information to do that filtering is important.)
  * 8.) The timestamps of partial updates and deletes, especially for single-column queries that might produce
@@ -88,22 +89,28 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
     @BeforeClass
     public static void setUpCluster() throws IOException
     {
-        CLUSTER = init(Cluster.build(NODES).withConfig(config -> config.set("hinted_handoff_enabled", false).with(GOSSIP).with(NETWORK)).start());
+        CLUSTER = Cluster.build(NODES).withConfig(config -> config.set("hinted_handoff_enabled", false).with(GOSSIP).with(NETWORK)).start();
 
-        // All parameterized test scenarios share the same table and attached indexes, but write to different partitions
-        // that are deleted after each scenario completes.
-        String createTableDDL = String.format("CREATE TABLE %s.%s (pk int, pk2 int, ck int, s int static, y int static, a int, b int, x int, PRIMARY KEY ((pk, pk2), ck)) WITH read_repair = 'NONE'",
-                                              KEYSPACE, TEST_TABLE_NAME);
-        CLUSTER.schemaChange(createTableDDL);
-        CLUSTER.disableAutoCompaction(KEYSPACE);
+        for (ReplicationType replicationType : ReplicationType.values())
+        {
+            String keyspace = KEYSPACE + '_' + replicationType;
+            CLUSTER.schemaChange("CREATE KEYSPACE " + keyspace + " WITH replication = {'class': 'SimpleStrategy', 'replication_factor': " + CLUSTER.size() + "} AND replication_type = '" + replicationType + '\'');
 
-        CLUSTER.schemaChange(String.format("CREATE INDEX pk2_idx ON %s.%s(pk2) USING 'sai'", KEYSPACE, TEST_TABLE_NAME));
-        CLUSTER.schemaChange(String.format("CREATE INDEX ck_idx ON %s.%s(ck) USING 'sai'", KEYSPACE, TEST_TABLE_NAME));
-        CLUSTER.schemaChange(String.format("CREATE INDEX s_idx ON %s.%s(s) USING 'sai'", KEYSPACE, TEST_TABLE_NAME));
-        CLUSTER.schemaChange(String.format("CREATE INDEX a_idx ON %s.%s(a) USING 'sai'", KEYSPACE, TEST_TABLE_NAME));
-        CLUSTER.schemaChange(String.format("CREATE INDEX b_idx ON %s.%s(b) USING 'sai'", KEYSPACE, TEST_TABLE_NAME));
+            // All parameterized test scenarios share the same table and attached indexes, but write to different partitions
+            // that are deleted after each scenario completes.
+            String createTableDDL = String.format("CREATE TABLE %s.%s (pk int, pk2 int, ck int, s int static, y int static, a int, b int, x int, PRIMARY KEY ((pk, pk2), ck)) WITH read_repair = 'NONE'",
+                                                  keyspace, TEST_TABLE_NAME);
+            CLUSTER.schemaChange(createTableDDL);
+            CLUSTER.disableAutoCompaction(keyspace);
 
-        SAIUtil.waitForIndexQueryable(CLUSTER, KEYSPACE);
+            CLUSTER.schemaChange(String.format("CREATE INDEX pk2_idx ON %s.%s(pk2) USING 'sai'", keyspace, TEST_TABLE_NAME));
+            CLUSTER.schemaChange(String.format("CREATE INDEX ck_idx ON %s.%s(ck) USING 'sai'", keyspace, TEST_TABLE_NAME));
+            CLUSTER.schemaChange(String.format("CREATE INDEX s_idx ON %s.%s(s) USING 'sai'", keyspace, TEST_TABLE_NAME));
+            CLUSTER.schemaChange(String.format("CREATE INDEX a_idx ON %s.%s(a) USING 'sai'", keyspace, TEST_TABLE_NAME));
+            CLUSTER.schemaChange(String.format("CREATE INDEX b_idx ON %s.%s(b) USING 'sai'", keyspace, TEST_TABLE_NAME));
+
+            SAIUtil.waitForIndexQueryable(CLUSTER, keyspace);
+        }
     }
 
     static class Specification
@@ -115,6 +122,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
         final int partitionKey;
         final boolean flushPartials;
         final Expression.IndexOperator validationMode;
+        final ReplicationType replicationType;
 
         Specification(boolean restrictPartitionKey,
                       String[] columns,
@@ -122,7 +130,8 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
                       StatementType partialUpdateType,
                       int partitionKey,
                       boolean flushPartials,
-                      Expression.IndexOperator validationMode)
+                      Expression.IndexOperator validationMode,
+                      ReplicationType replicationType)
         {
             this.restrictPartitionKey = restrictPartitionKey;
             this.columns = columns;
@@ -131,12 +140,18 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
             this.partitionKey = partitionKey;
             this.flushPartials = flushPartials;
             this.validationMode = validationMode;
+            this.replicationType = replicationType;
         }
 
         public String[] nonKeyColumns()
         {
             return Arrays.stream(columns).filter(c -> !c.equals("ck") && !c.equals("pk") && !c.equals("pk2")).toArray(String[]::new);
         }
+        
+        public String keyspaceName()
+        {
+            return KEYSPACE + '_' + replicationType;
+        }
 
         public String tableName()
         {
@@ -152,7 +167,8 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
                    ", partialUpdateType=" + partialUpdateType +
                    ", partitionKey=" + partitionKey +
                    ", flushPartials=" + flushPartials +
-                   ", validationMode=" + validationMode;
+                   ", validationMode=" + validationMode +
+                   ", replicationType=" + replicationType;
         }
 
         @Override
@@ -163,13 +179,14 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
             Specification that = (Specification) o;
             return Arrays.equals(columns, that.columns) 
                    && existing == that.existing && restrictPartitionKey == that.restrictPartitionKey 
-                   && partialUpdateType == that.partialUpdateType && partitionKey == that.partitionKey && flushPartials == that.flushPartials;
+                   && partialUpdateType == that.partialUpdateType && partitionKey == that.partitionKey 
+                   && flushPartials == that.flushPartials && replicationType == that.replicationType;
         }
 
         @Override
         public int hashCode()
         {
-            int result = Objects.hash(existing, restrictPartitionKey, partialUpdateType, partitionKey, flushPartials);
+            int result = Objects.hash(existing, restrictPartitionKey, partialUpdateType, partitionKey, flushPartials, replicationType);
             result = 31 * result + Arrays.hashCode(columns);
             return result < 0 ? -result : result;
         }
@@ -194,7 +211,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
         {
             for (int i = 0; i < PARTITIONS_PER_TEST; i++)
             {
-                StringBuilder insert = new StringBuilder("INSERT INTO ").append(KEYSPACE).append('.').append(specification.tableName());
+                StringBuilder insert = new StringBuilder("INSERT INTO ").append(specification.keyspaceName()).append('.').append(specification.tableName());
                 insert.append("(pk, pk2, ck");
 
                 for (Object column : specification.nonKeyColumns())
@@ -220,7 +237,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
                 CLUSTER.coordinator(1).execute(insert.toString(), ConsistencyLevel.ALL);
             }
 
-            CLUSTER.get(1).nodetoolResult("repair", KEYSPACE).asserts().success();
+            CLUSTER.get(1).nodetoolResult("repair", specification.keyspaceName()).asserts().success();
         }
 
         public void writeUnrepairedRows()
@@ -269,7 +286,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
             }
 
             String dml = String.format("INSERT INTO %s.%s(pk, pk2, ck, %s) VALUES (?, ?, 0, ?) USING TIMESTAMP %d",
-                                       KEYSPACE, specification.tableName(), column, nextTimestamp++);
+                                       specification.keyspaceName(), specification.tableName(), column, nextTimestamp++);
             CLUSTER.get(node).executeInternal(dml, partitionKey, partitionKey, value);
             node = nextNode(node);
             return node;
@@ -283,11 +300,11 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
 
             int partitionKey = specification.partitionKey + partitionIndex;
             String dml = String.format("DELETE %s FROM %s.%s USING TIMESTAMP %d WHERE pk = %d AND pk2 = %d AND ck = 0",
-                                       column, KEYSPACE, specification.tableName(), nextTimestamp++, partitionKey, partitionKey);
+                                       column, specification.keyspaceName(), specification.tableName(), nextTimestamp++, partitionKey, partitionKey);
 
             if (isStatic((String) column))
                 dml = String.format("DELETE %s FROM %s.%s USING TIMESTAMP %d WHERE pk = %d AND pk2 = %d",
-                                    column, KEYSPACE, specification.tableName(), nextTimestamp++, partitionKey, partitionKey);
+                                    column, specification.keyspaceName(), specification.tableName(), nextTimestamp++, partitionKey, partitionKey);
 
             CLUSTER.get(node).executeInternal(dml);
             node = nextNode(node);
@@ -341,7 +358,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
             for (Object column : specification.nonKeyColumns())
                 select.append(", ").append(column);
 
-            select.append(" FROM ").append(KEYSPACE).append('.').append(specification.tableName()).append(" WHERE ");
+            select.append(" FROM ").append(specification.keyspaceName()).append('.').append(specification.tableName()).append(" WHERE ");
 
             ArrayList<String> restricted = Lists.newArrayList(specification.columns);
 
@@ -400,6 +417,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
         }
     }
 
+    @SuppressWarnings("ClassEscapesDefinedScope")
     @Parameterized.Parameter
     public Specification specification;
 
@@ -417,33 +435,49 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
             for (boolean restrictPartitionKey : new boolean[] { false, true })
             {
                 for (String[] columns : new String[][] { { "ck", "a" }, { "ck", "s" }, { "s", "a" }, { "a", "b" }, { "s", "x" }, { "s", "y" }, { "a", "x" }, { "a", "y" }, { "a" }, { "s" } })
-                    for (boolean existing : new boolean[] { false, true })
+                {
+                    for (boolean existing : new boolean[]{ false, true })
                     {
-                        parameters.add(new Object[] { new Specification(restrictPartitionKey, columns, existing, StatementType.INSERT, nextPartitionKey, flushPartials, EQ) });
+                        parameters.add(new Object[]{ new Specification(restrictPartitionKey, columns, existing, StatementType.INSERT, nextPartitionKey, flushPartials, EQ, ReplicationType.untracked) });
                         nextPartitionKey += PARTITIONS_PER_TEST;
+
+                        if (restrictPartitionKey)
+                        {
+                            parameters.add(new Object[]{ new Specification(true, columns, existing, StatementType.INSERT, nextPartitionKey, flushPartials, EQ, ReplicationType.tracked) });
+                            nextPartitionKey += PARTITIONS_PER_TEST;
+                        }
                     }
+                }
 
                 // Deletion scenarios assume existing data.
                 for (String[] columns : new String[][] { { "s", "a" }, { "a", "b" }, { "s", "x" }, { "a", "x" }, { "a", "y" }, { "a" }, { "s" } })
                 {
-                    parameters.add(new Object[] { new Specification(restrictPartitionKey, columns, true, StatementType.DELETE, nextPartitionKey, flushPartials, EQ) });
+                    parameters.add(new Object[] { new Specification(restrictPartitionKey, columns, true, StatementType.DELETE, nextPartitionKey, flushPartials, EQ, ReplicationType.untracked) });
                     nextPartitionKey += PARTITIONS_PER_TEST;
+
+                    if (restrictPartitionKey)
+                    {
+                        parameters.add(new Object[] { new Specification(true, columns, true, StatementType.DELETE, nextPartitionKey, flushPartials, EQ, ReplicationType.tracked) });
+                        nextPartitionKey += PARTITIONS_PER_TEST;
+                    }
                 }
             }
 
             // Note that scenarios around indexes on a partition key element only appear here where we neither
             // delete nor restrict on partition, as both would be nonsensical.
             for (String[] columns : new String[][] { { "pk2", "a" }, { "s", "a" }, { "a", "b" }, { "s", "x" }, { "a", "x" }, { "a", "y" }, { "a" }, { "s" } })
-                for (boolean existing : new boolean[] { false, true })
+            {
+                for (boolean existing : new boolean[]{ false, true })
                 {
-                    parameters.add(new Object[]{ new Specification(false, columns, existing, StatementType.INSERT, nextPartitionKey, flushPartials, RANGE) });
+                    parameters.add(new Object[]{ new Specification(false, columns, existing, StatementType.INSERT, nextPartitionKey, flushPartials, RANGE, ReplicationType.untracked) });
                     nextPartitionKey += PARTITIONS_PER_TEST;
                 }
+            }
 
             // Deletion scenarios assume existing data.
             for (String[] columns : new String[][] { { "s", "a" }, { "a", "b" }, { "s", "x" }, { "a", "x" }, { "a", "y" }, { "a" }, { "s" } })
             {
-                parameters.add(new Object[]{ new Specification(false, columns, true, StatementType.DELETE, nextPartitionKey, flushPartials, RANGE) });
+                parameters.add(new Object[]{ new Specification(false, columns, true, StatementType.DELETE, nextPartitionKey, flushPartials, RANGE, ReplicationType.untracked) });
                 nextPartitionKey += PARTITIONS_PER_TEST;
             }
         }
@@ -468,7 +502,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
 
         if (specification.flushPartials)
             // Flushg partial rows from Memtable-attached indexes to SSTable indexes:
-            CLUSTER.stream().forEach(i -> i.flush(KEYSPACE));
+            CLUSTER.stream().forEach(i -> i.flush(specification.keyspaceName()));
 
         // If we wrote an initial (repaired) version of the row, do negative validation.
         // (i.e. Ensure queries that would have initially produced matches no longer do.) 
@@ -483,7 +517,7 @@ public class PartialUpdateHandlingTest extends TestBaseImpl
     @After
     public void truncateTable()
     {
-        CLUSTER.coordinator(1).execute(String.format("TRUNCATE TABLE %s.%s ", KEYSPACE, specification.tableName()), ALL);
+        CLUSTER.coordinator(1).execute(String.format("TRUNCATE TABLE %s.%s ", specification.keyspaceName(), specification.tableName()), ALL);
     }
 
     @AfterClass

test/distributed/org/apache/cassandra/distributed/test/sai/StrictFilteringTest.java

@@ -19,21 +19,29 @@
 package org.apache.cassandra.distributed.test.sai;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Iterator;
+import java.util.List;
 
 import org.junit.AfterClass;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
 import org.apache.cassandra.db.Keyspace;
 import org.apache.cassandra.distributed.Cluster;
 import org.apache.cassandra.distributed.api.ConsistencyLevel;
 import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.schema.ReplicationType;
 
 import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
 import static org.apache.cassandra.distributed.api.Feature.NETWORK;
 import static org.apache.cassandra.distributed.shared.AssertUtils.assertRows;
 import static org.apache.cassandra.distributed.shared.AssertUtils.row;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 
 /**
@@ -42,16 +50,38 @@ import static org.junit.Assert.assertEquals;
  * 
  * @see <a href="https://issues.apache.org/jira/browse/CASSANDRA-19018">CASSANDRA-19018</a>
  */
+@RunWith(Parameterized.class)
 public class StrictFilteringTest extends TestBaseImpl
 {
     private static Cluster CLUSTER;
 
+    private static int keyspaceIdx;
+
+    @Parameterized.Parameter
+    public ReplicationType replicationType;
+
     @BeforeClass
     public static void setUpCluster() throws IOException
     {
         CLUSTER = init(Cluster.build(2).withConfig(config -> config.set("hinted_handoff_enabled", false).with(GOSSIP).with(NETWORK)).start());
     }
 
+    @Parameterized.Parameters(name = "{index}: replication={0}")
+    public static Collection<Object[]> data()
+    {
+        List<Object[]> result = new ArrayList<>();
+        for (ReplicationType replication : ReplicationType.values())
+            result.add(new Object[]{replication});
+        return result;
+    }
+
+    @Before
+    public void setup()
+    {
+        KEYSPACE = "ks_" + keyspaceIdx++;
+        CLUSTER.schemaChange("CREATE KEYSPACE " + KEYSPACE + " WITH replication = {'class': 'SimpleStrategy', 'replication_factor': " + CLUSTER.size() + "} AND replication_type='" + replicationType + "';");
+    }
+
     @Test
     public void shouldDegradeToUnionOnSingleStatic()
     {
@@ -260,7 +290,11 @@ public class StrictFilteringTest extends TestBaseImpl
         assertRows(initialRows, row(0, 4, 1));
 
         Long srpRequestsAfter = CLUSTER.get(1).callOnInstance(() -> Keyspace.open(KEYSPACE).getColumnFamilyStore("necessary_short_read").metric.shortReadProtectionRequests.getCount());
-        assertEquals(srpRequestsBefore + 2L, srpRequestsAfter.longValue());
+
+        if (replicationType == ReplicationType.untracked)
+            assertEquals(srpRequestsBefore + 2L, srpRequestsAfter.longValue());
+        else if (replicationType == ReplicationType.tracked)
+            assertThat(srpRequestsAfter).isGreaterThanOrEqualTo(0L).isLessThanOrEqualTo(2L);
     }
 
     @Test

test/distributed/org/apache/cassandra/distributed/test/tracking/MutationTrackingPartitionReadTest.java

@@ -0,0 +1,218 @@
+/*
+ * 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.distributed.test.tracking;
+
+import java.io.IOException;
+import java.math.BigInteger;
+import java.util.Map;
+import java.util.UUID;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.api.Feature;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.distributed.test.sai.SAIUtil;
+
+import static org.apache.cassandra.distributed.shared.AssertUtils.assertRows;
+import static org.apache.cassandra.distributed.shared.AssertUtils.row;
+import static org.junit.Assert.assertEquals;
+
+public class MutationTrackingPartitionReadTest extends TestBaseImpl
+{
+    private static final int REPLICAS = 3;
+
+    private static Cluster cluster;
+
+    @BeforeClass
+    public static void setup() throws IOException
+    {
+        cluster = Cluster.build()
+                         .withNodes(REPLICAS)
+                         .withConfig(cfg -> cfg.with(Feature.NETWORK, Feature.GOSSIP)
+                                                             .set("mutation_tracking_enabled", true)
+                                                             .set("hinted_handoff_enabled", false))
+                         .start();
+    }
+
+    @AfterClass
+    public static void teardown()
+    {
+        if (cluster != null)
+            cluster.close();
+    }
+
+    @Test
+    public void testEqQueryOnStaticColumn()
+    {
+        String keyspace = "test_eq_query_on_static_column";
+        cluster.schemaChange(withKeyspace("CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 3} AND replication_type='tracked'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk0 varint, pk1 uuid, ck0 time, s0 ascii static, v1 double, PRIMARY KEY ((pk0, pk1), ck0)) " +
+                                          "WITH CLUSTERING ORDER BY (ck0 DESC) AND read_repair = 'NONE'", keyspace));
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_s0 ON %s.tbl(s0) USING 'sai'", keyspace));
+        SAIUtil.waitForIndexQueryable(cluster, keyspace);
+        cluster.forEach(i -> i.nodetoolResult("disableautocompaction", keyspace, "tbl").asserts().success());
+
+        cluster.get(3).executeInternal(withKeyspace("INSERT INTO %s.tbl (pk0, pk1, s0) VALUES (-58, 00000000-0000-4d00-8600-000000000000, 'foo') USING TIMESTAMP 5", keyspace));
+        cluster.get(1).executeInternal(withKeyspace("DELETE s0, s0 FROM %s.tbl USING TIMESTAMP 13 WHERE  pk0 = 7 AND  pk1 = 00000000-0000-4e00-9600-000000000000", keyspace));
+        cluster.get(1).executeInternal(withKeyspace("INSERT INTO %s.tbl (pk0, pk1, ck0, v1) VALUES (-58, 00000000-0000-4d00-8600-000000000000, '16:40:27.677919817', 1.6896613611522374E184) USING TIMESTAMP 14", keyspace));
+        cluster.get(1).executeInternal(withKeyspace("UPDATE %s.tbl USING TIMESTAMP 15 SET v1=8.05223257349057E-164 WHERE  pk0 = -58 AND  pk1 = 00000000-0000-4d00-8600-000000000000 AND  ck0 = '20:02:33.822429155'", keyspace));
+        cluster.get(2).executeInternal(withKeyspace("INSERT INTO %s.tbl (pk0, pk1, s0) VALUES (-58, 00000000-0000-4d00-8600-000000000000, 'bar') USING TIMESTAMP 18", keyspace));
+
+        String select = withKeyspace("SELECT pk0, pk1, ck0 FROM %s.tbl WHERE pk0 = -58 AND pk1 = 00000000-0000-4d00-8600-000000000000 AND s0 = 'bar'", keyspace);
+        Object[][] result = cluster.coordinator(1).execute(select, ConsistencyLevel.ALL);
+        assertRows(result, row(BigInteger.valueOf(-58), UUID.fromString("00000000-0000-4d00-8600-000000000000"), 72153822429155L),
+                           row(BigInteger.valueOf(-58), UUID.fromString("00000000-0000-4d00-8600-000000000000"), 60027677919817L));
+    }
+
+    @Test
+    public void testMissingPartitionDelete()
+    {
+        String keyspace = "test_missing_partition_delete";
+        cluster.schemaChange(withKeyspace("CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 3} AND replication_type='tracked'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk0 text, pk1 bigint, ck0 smallint, v0 timestamp, v1 int, PRIMARY KEY ((pk0, pk1), ck0)) " +
+                                          "WITH CLUSTERING ORDER BY (ck0 ASC) AND read_repair = 'NONE'", keyspace));
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_v1_idx ON %s.tbl(v1) USING 'sai'", keyspace));
+        SAIUtil.waitForIndexQueryable(cluster, keyspace);
+        cluster.forEach(i -> i.nodetoolResult("disableautocompaction", keyspace, "tbl").asserts().success());
+
+        cluster.get(2).executeInternal(withKeyspace("UPDATE %s.tbl USING TIMESTAMP 8 SET v1=1778069545 WHERE pk0 = 'ad1b:bbdc:e712:574:e7ca:104e:5abb:d9e1' AND pk1 = -5572993830691022649 AND ck0 = 32379", keyspace));
+        cluster.get(1).executeInternal(withKeyspace("DELETE FROM %s.tbl USING TIMESTAMP 12 WHERE pk0 = 'ad1b:bbdc:e712:574:e7ca:104e:5abb:d9e1' AND pk1 = -5572993830691022649", keyspace));
+        cluster.get(2).executeInternal(withKeyspace("UPDATE %s.tbl USING TIMESTAMP 14 SET v0=null, v1=1353378764 WHERE pk0 = 'ad1b:bbdc:e712:574:e7ca:104e:5abb:d9e1' AND pk1 = -5572993830691022649 AND ck0 = 29521", keyspace));
+
+        String select = withKeyspace("SELECT pk0, pk1, ck0, v1 FROM %s.tbl WHERE pk0 = 'ad1b:bbdc:e712:574:e7ca:104e:5abb:d9e1' AND pk1 = -5572993830691022649 AND v1 <= 1353378764 LIMIT 136", keyspace);
+        cluster.coordinator(1).execute(select, ConsistencyLevel.ALL);
+
+        select = withKeyspace("SELECT pk0, pk1, ck0, v1 FROM %s.tbl WHERE pk0 = 'ad1b:bbdc:e712:574:e7ca:104e:5abb:d9e1' AND pk1 = -5572993830691022649 AND v1 >= 1353378764 LIMIT 116", keyspace);
+        Object[][] result = cluster.coordinator(1).execute(select, ConsistencyLevel.ALL);
+
+        assertRows(result, row("ad1b:bbdc:e712:574:e7ca:104e:5abb:d9e1", -5572993830691022649L, (short) 29521, 1353378764));
+    }
+
+    @Test
+    public void testMultiColumnPartitionRestrictedQuery()
+    {
+        String keyspace = "test_multi_column_partition_restricted";
+        cluster.schemaChange(withKeyspace("CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '3'} AND replication_type='tracked'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk0 varint, ck0 timeuuid,ck1 uuid, s0 int static, s1 vector<timeuuid, 2> static, v0 frozen<set<vector<time, 1>>>, v2 date, v3 int, v1 set<frozen<set<boolean>>>, PRIMARY KEY (pk0, ck0, ck1)) " +
+                                          "WITH CLUSTERING ORDER BY (ck0 ASC, ck1 DESC) AND read_repair = 'NONE'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_ck0 ON %s.tbl(ck0) USING 'sai'", keyspace));
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_ck1 ON %s.tbl(ck1) USING 'sai'", keyspace));
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_v3 ON %s.tbl(v3) USING 'sai'", keyspace));
+        SAIUtil.waitForIndexQueryable(cluster, keyspace);
+        cluster.forEach(i -> i.nodetoolResult("disableautocompaction", keyspace, "tbl").asserts().success());
+
+        cluster.get(1).executeInternal(withKeyspace("DELETE FROM %s.tbl USING TIMESTAMP 160 WHERE pk0 = -320778557", keyspace));
+        cluster.get(3).executeInternal(withKeyspace("INSERT INTO %s.tbl (pk0, s0, s1) VALUES (-320778557, 1363549784, [00000000-0000-1d00-8300-000000000000, 00000000-0000-1a00-8900-000000000000]) USING TIMESTAMP 166", keyspace));
+        cluster.get(1).executeInternal(withKeyspace("UPDATE %s.tbl USING TIMESTAMP 168 SET v1 = {{false}, {false, true}}, v3 = -1669443995, s0 = 1234171012 " +
+                                                    "WHERE pk0 = -320778557 AND ck0 = 00000000-0000-1400-9600-000000000000 " +
+                                                    "AND ck1 IN (00000000-0000-4e00-bf00-000000000000, 00000000-0000-4500-a100-000000000000, 00000000-0000-4000-9100-000000000000)", keyspace));
+
+        cluster.forEach(i -> i.nodetoolResult("flush", keyspace, "tbl").asserts().success());
+
+        String select = withKeyspace("SELECT * FROM %s.tbl WHERE pk0 = -320778557 AND ck1 = 00000000-0000-4c00-a700-000000000000 AND v3 = -1669443995 AND ck0 = 00000000-0000-1300-b800-000000000000 ALLOW FILTERING", keyspace);
+        Object[][] result = cluster.coordinator(2).execute(select, ConsistencyLevel.ALL);
+        assertRows(result);
+    }
+
+    @Test
+    public void testMissingRowWithMultipleIndexedColumns()
+    {
+        String keyspace = "test_missing_row";
+        cluster.schemaChange(withKeyspace("CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '3'} AND replication_type='tracked'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk0 timeuuid, pk1 timestamp, ck0 boolean, ck1 varint, v0 frozen<set<int>>, v1 text, v2 frozen<tuple<frozen<set<tinyint>>, frozen<list<boolean>>>>, v3 ascii, PRIMARY KEY ((pk0, pk1), ck0, ck1)) " +
+                                          "WITH CLUSTERING ORDER BY (ck0 DESC, ck1 ASC) AND read_repair = 'NONE'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_v2 ON %s.tbl(v2) USING 'sai'", keyspace));
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_v3 ON %s.tbl(v3) USING 'sai'", keyspace));
+        SAIUtil.waitForIndexQueryable(cluster, keyspace);
+
+        cluster.forEach(i -> i.nodetoolResult("disableautocompaction", keyspace, "tbl").asserts().success());
+
+        // Insert row on node3 at ts=11 with v2 and v3
+        cluster.get(3).executeInternal(withKeyspace("INSERT INTO %s.tbl (pk0, pk1, ck0, ck1, v0, v1, v2, v3) " +
+                                                    "VALUES (00000000-0000-1200-8700-000000000000, '2028-05-17T03:51:34.765Z', true, 0, {1, 2, 3}, 'test', ({13, 55}, [true]), 'original_value') " +
+                                                    "USING TIMESTAMP 11", keyspace));
+
+        cluster.get(3).executeInternal(withKeyspace("DELETE FROM %s.tbl USING TIMESTAMP 12 " +
+                                                    "WHERE pk0 = 00000000-0000-1200-8700-000000000000 AND pk1 = '2028-05-17T03:51:34.765Z' AND ck0 = false AND ck1 = 0", keyspace));
+
+        cluster.forEach(i -> i.nodetoolResult("flush", keyspace, "tbl").asserts().success());
+
+        cluster.get(3).executeInternal(withKeyspace("UPDATE %s.tbl USING TIMESTAMP 15 SET v0 = {10, 20}, v2 = ({-13, 44}, [false]), v1 = 'updated' " +
+                                                    "WHERE pk0 = 00000000-0000-1200-8700-000000000000 AND pk1 = '2028-05-17T03:51:34.765Z' AND ck0 = true AND ck1 = 0", keyspace));
+
+        cluster.forEach(i -> i.nodetoolResult("flush", keyspace, "tbl").asserts().success());
+
+        String select = withKeyspace("SELECT pk0, pk1, ck0, ck1, v3 FROM %s.tbl " +
+                                     "WHERE pk0 = 00000000-0000-1200-8700-000000000000 AND pk1 = '2028-05-17T03:51:34.765Z' AND v2 = ({-13, 44}, [false]) AND v3 = 'original_value' " +
+                                     "LIMIT 47 ALLOW FILTERING", keyspace);
+
+        Object[][] result = cluster.coordinator(1).execute(select, ConsistencyLevel.ALL);
+
+        assertEquals("Query should return exactly 1 row", 1, result.length);
+        assertEquals("ck0 should be true", true, result[0][2]);
+        assertEquals("ck1 should be 0", BigInteger.ZERO, result[0][3]);
+        assertEquals("v3 should be 'original_value'", "original_value", result[0][4]);
+    }
+
+    @Test
+    public void testStaticColumnUpdateWithRowQuery()
+    {
+        String keyspace = "test_static_column_update";
+        cluster.schemaChange(withKeyspace("CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 3} AND replication_type='tracked'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk0 int, pk1 text, ck0 int, s0 frozen<map<text, int>> static, v0 int, v1 text, v3 int, PRIMARY KEY ((pk0, pk1), ck0)) " +
+                                          "WITH CLUSTERING ORDER BY (ck0 ASC) AND read_repair = 'NONE'", keyspace));
+
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_ck0 ON %s.tbl(ck0) USING 'sai'", keyspace));
+        cluster.schemaChange(withKeyspace("CREATE INDEX tbl_s0 ON %s.tbl(FULL(s0)) USING 'sai'", keyspace));
+        SAIUtil.waitForIndexQueryable(cluster, keyspace);
+        cluster.forEach(i -> i.nodetoolResult("disableautocompaction", keyspace, "tbl").asserts().success());
+
+        // Step 1: UPDATE on node2 @ ts=6 - sets s0=null and regular columns for ck0=100
+        cluster.get(2).executeInternal(withKeyspace("UPDATE %s.tbl USING TIMESTAMP 6 SET s0=null, v0=42, v1='value_from_ts6', v3=999 WHERE pk0 = 1 AND pk1 = 'partition1' AND ck0 = 100", keyspace));
+
+        // Flush to ensure data is in SSTables
+        cluster.forEach(i -> i.nodetoolResult("flush", keyspace, "tbl").asserts().success());
+
+        // Step 2: INSERT on node1 @ ts=23 - updates static column s0 only
+        cluster.get(1).executeInternal(withKeyspace("INSERT INTO %s.tbl (pk0, pk1, s0) VALUES (1, 'partition1', {'key1': 10, 'key2': 20, 'key3': 30}) USING TIMESTAMP 23", keyspace));
+        cluster.forEach(i -> i.nodetoolResult("flush", keyspace, "tbl").asserts().success());
+
+        String select = withKeyspace("SELECT pk0, pk1, ck0, s0, v0, v1, v3 FROM %s.tbl " +
+                                     "WHERE pk0 = 1 AND pk1 = 'partition1' AND ck0 = 100 AND v0 = 42 AND v3 = 999 AND s0 = {'key1': 10, 'key2': 20, 'key3': 30} ALLOW FILTERING", keyspace);
+
+        Object[][] result = cluster.coordinator(2).execute(select, ConsistencyLevel.ALL);
+        assertRows(result, row(1, "partition1", 100, Map.of("key1", 10, "key2", 20, "key3", 30), 42, "value_from_ts6", 999));
+    }
+
+    public static String withKeyspace(String replaceIn, String keyspace)
+    {
+        return String.format(replaceIn, keyspace);
+    }
+}

test/distributed/org/apache/cassandra/fuzz/sai/MultiNodeSAIMutationTrackingTest.java

@@ -0,0 +1,49 @@
+/*
+ * 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.fuzz.sai;
+
+import org.apache.cassandra.harry.SchemaSpec;
+import org.apache.cassandra.harry.gen.Generator;
+
+public class MultiNodeSAIMutationTrackingTest extends MultiNodeSAITestBase
+{
+    public MultiNodeSAIMutationTrackingTest()
+    {
+        super(null);
+    }
+
+    @Override
+    protected String replicationType()
+    {
+        return "tracked";
+    }
+
+    @Override
+    protected Generator<SchemaSpec> schemaGenerator(boolean disableReadRepair)
+    {
+        // Mutation tracking does not require read-repair
+        return super.schemaGenerator(true);
+    }
+
+    @Override
+    protected void repair(SchemaSpec schema)
+    {
+        // Mutation tracking should not require normal repairs
+    }
+}

test/distributed/org/apache/cassandra/fuzz/sai/SingleNodeSAITestBase.java

@@ -67,8 +67,6 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
     private static final int COMPACTION_SKIP = 4435;
     private static final int DEFAULT_REPAIR_SKIP = 8869;
 
-    private static final int OPERATIONS_PER_RUN = 30_000;
-
     private static final int NUM_PARTITIONS = 64;
     private static final int NUM_VISITED_PARTITIONS = 16;
     protected static final int MAX_PARTITION_SIZE = 2000;
@@ -90,6 +88,11 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
         return 5;
     }
 
+    protected int operationsPerRun()
+    {
+        return 30_000;
+    }
+
     @BeforeClass
     public static void before() throws Throwable
     {
@@ -113,6 +116,7 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
         cluster.startup();
         cluster = init(cluster);
     }
+
     @AfterClass
     public static void afterClass()
     {
@@ -122,8 +126,7 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
     @Before
     public void beforeEach()
     {
-        cluster.schemaChange("DROP KEYSPACE IF EXISTS " + KEYSPACE);
-        cluster.schemaChange("CREATE KEYSPACE " + KEYSPACE + " WITH replication = {'class': 'SimpleStrategy', 'replication_factor': " + rf() + "};");
+        cluster.schemaChange("CREATE KEYSPACE IF NOT EXISTS " + KEYSPACE + " WITH replication = {'class': 'SimpleStrategy', 'replication_factor': " + rf() + "} AND replication_type = '" + replicationType() + "';");
     }
 
     protected int rf()
@@ -131,6 +134,11 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
         return 1;
     }
 
+    protected String replicationType()
+    {
+        return "untracked";
+    }
+
     @Test
     public void simplifiedSaiTest()
     {
@@ -235,7 +243,7 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
                 if (IndexTermType.isEqOnlyType(schema.clusteringKeys.get(i).type.asServerType()))
                     eqOnlyClusteringColumns.add(i);
 
-            for (int i = 0; i < OPERATIONS_PER_RUN; i++)
+            for (int i = 0; i < operationsPerRun(); i++)
             {
                 int partitionIndex = pkGen.generate(rng);
                 HistoryBuilderHelper.insertRandomData(schema, partitionIndex, ckGen.generate(rng), rng, 0.5d, history);
@@ -307,8 +315,7 @@ public abstract class SingleNodeSAITestBase extends TestBaseImpl
 
     protected Generator<SchemaSpec> schemaGenerator(boolean disableReadRepair)
     {
-        SchemaSpec.OptionsBuilder builder = SchemaSpec.optionsBuilder().disableReadRepair(disableReadRepair)
-                                                                       .compactionStrategy("LeveledCompactionStrategy");
+        SchemaSpec.OptionsBuilder builder = SchemaSpec.optionsBuilder().disableReadRepair(disableReadRepair).compactionStrategy("LeveledCompactionStrategy");
         return SchemaGenerators.schemaSpecGen(KEYSPACE, "basic_sai", MAX_PARTITION_SIZE, builder);
     }
 

test/unit/org/apache/cassandra/db/ColumnFamilyStoreTest.java

@@ -59,6 +59,7 @@ import org.apache.cassandra.db.lifecycle.SSTableSet;
 import org.apache.cassandra.db.memtable.AbstractMemtable;
 import org.apache.cassandra.db.memtable.Memtable;
 import org.apache.cassandra.db.partitions.FilteredPartition;
+import org.apache.cassandra.db.partitions.Partition;
 import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
 import org.apache.cassandra.db.rows.Cell;
@@ -968,6 +969,12 @@ public class ColumnFamilyStoreTest
             {
             }
 
+            @Override
+            public Partition snapshotPartition(DecoratedKey partitionKey)
+            {
+                return null;
+            }
+
             @Override
             public UnfilteredRowIterator rowIterator(DecoratedKey key,
                                                      Slices slices,