Adding Mutual exclusion logic for SAI index rebuilding and ZCS strea…

src/java/org/apache/cassandra/db/compaction/CompactionManager.java

@@ -2121,7 +2121,14 @@ public void run()
             }
         };
 
-        return secondaryIndexExecutor.submitIfRunning(runnable, "index build");
+        Future<?> future = secondaryIndexExecutor.submitIfRunning(runnable, "index build");
+        if (future.isCancelled())
+        {
+            // Submission was rejected (e.g. the executor is shutting down), so build() will not run. Let the builder
+            // release any resources it reserved at construction time (e.g. the SAI rebuild status, CASSANDRA-21520).
+            builder.onNotExecuted();
+        }
+        return future;
     }
 
     /**

src/java/org/apache/cassandra/db/streaming/CassandraOutgoingFile.java

@@ -50,6 +50,8 @@ public class CassandraOutgoingFile implements OutgoingStream
     private final StreamOperation operation;
     private final CassandraStreamHeader header;
     private final List<Range<Token>> ranges;
+    // guards against releasing the entire-sstable streaming status more than once (CASSANDRA-21520)
+    private boolean streamRebuildStatusReleased = false;
 
     public CassandraOutgoingFile(StreamOperation operation, Ref<SSTableReader> ref,
                                  List<SSTableReader.PartitionPositionBounds> sections, List<Range<Token>> normalizedRanges,
@@ -66,9 +68,24 @@ public CassandraOutgoingFile(StreamOperation operation, Ref<SSTableReader> ref,
         SSTableReader sstable = ref.get();
 
         this.filename = sstable.getFilename();
-        this.shouldStreamEntireSSTable = computeShouldStreamEntireSSTables();
-        ComponentManifest manifest = ComponentManifest.create(sstable);
-        this.header = makeHeader(sstable, operation, sections, estimatedKeys, shouldStreamEntireSSTable, manifest);
+        // Decide the streaming mode here (not at write() time): getNumFiles()/the manifest are advertised to the
+        // receiver at plan time, and the receiver only completes once it has received exactly that many files. If
+        // an SAI rebuild is in progress we must fall back to legacy streaming now, so the advertised count matches
+        // what is actually sent. See CASSANDRA-21520.
+        boolean entire = computeShouldStreamEntireSSTables() && sstable.streamRebuildState().tryBeginStreaming();
+        this.shouldStreamEntireSSTable = entire;
+        try
+        {
+            ComponentManifest manifest = ComponentManifest.create(sstable);
+            this.header = makeHeader(sstable, operation, sections, estimatedKeys, entire, manifest);
+        }
+        catch (RuntimeException | Error e)
+        {
+            // The stream will never be finish()ed if construction fails, so release the status we just acquired.
+            if (entire)
+                sstable.streamRebuildState().endStreaming();
+            throw e;
+        }
     }
 
     private static CassandraStreamHeader makeHeader(SSTableReader sstable,
@@ -212,9 +229,25 @@ public boolean contained(List<SSTableReader.PartitionPositionBounds> sections, S
     @Override
     public void finish()
     {
+        releaseStreamRebuildStatus();
         ref.release();
     }
 
+    /**
+     * Releases the entire-sstable streaming status this stream reserved at construction, exactly once, without
+     * releasing the sstable reference. Used both by {@link #finish()} and by callers that must clean up a
+     * constructed-but-never-transferred stream (e.g. an error while planning outgoing streams) so the status
+     * cannot leak. See CASSANDRA-21520.
+     */
+    public void releaseStreamRebuildStatus()
+    {
+        if (shouldStreamEntireSSTable && !streamRebuildStatusReleased)
+        {
+            streamRebuildStatusReleased = true;
+            ref.get().streamRebuildState().endStreaming();
+        }
+    }
+
     public boolean equals(Object o)
     {
         if (this == o) return true;

src/java/org/apache/cassandra/db/streaming/CassandraStreamManager.java

@@ -91,6 +91,9 @@ public StreamReceiver createStreamReceiver(StreamSession session, List<Range<Tok
     public Collection<OutgoingStream> createOutgoingStreams(StreamSession session, RangesAtEndpoint replicas, TimeUUID pendingRepair, PreviewKind previewKind)
     {
         Refs<SSTableReader> refs = new Refs<>();
+        // Declared outside the try so the catch can release the entire-sstable streaming status reserved by any
+        // stream already constructed before a failure, without leaking it (CASSANDRA-21520).
+        List<OutgoingStream> streams = new ArrayList<>();
         try
         {
             final List<Range<PartitionPosition>> keyRanges = new ArrayList<>(replicas.size());
@@ -142,7 +145,6 @@ else if (pendingRepair == ActiveRepairService.NO_PENDING_REPAIR)
             List<Range<Token>> normalizedFullRanges = Range.normalize(replicas.onlyFull().ranges());
             List<Range<Token>> normalizedAllRanges = Range.normalize(replicas.ranges());
             //Create outgoing file streams for ranges possibly skipping repaired ranges in sstables
-            List<OutgoingStream> streams = new ArrayList<>(refs.size());
             for (SSTableReader sstable : refs)
             {
                 List<Range<Token>> ranges = sstable.isRepaired() ? normalizedFullRanges : normalizedAllRanges;
@@ -162,6 +164,10 @@ else if (pendingRepair == ActiveRepairService.NO_PENDING_REPAIR)
         }
         catch (Throwable t)
         {
+            // Release the entire-sstable streaming status held by any already-constructed stream (their refs are
+            // released below via refs.release()), so a planning failure cannot leak the status (CASSANDRA-21520).
+            for (OutgoingStream stream : streams)
+                ((CassandraOutgoingFile) stream).releaseStreamRebuildStatus();
             refs.release();
             throw t;
         }

src/java/org/apache/cassandra/db/streaming/ComponentContext.java

@@ -78,8 +78,13 @@ public FileChannel channel(Descriptor descriptor, Component component, long size
         @SuppressWarnings("resource") // file channel will be closed by Caller
         FileChannel channel = toTransfer.newReadChannel();
 
-        assert size == channel.size() : String.format("Entire sstable streaming expects %s file size to be %s but got %s.",
-                                                      component, size, channel.size());
+        if (size != channel.size())
+        {
+            long actual = channel.size();
+            FileUtils.closeQuietly(channel);
+            throw new IOException(String.format("Entire sstable streaming expects %s file size to be %s but got %s.",
+                                                component, size, actual));
+        }
         return channel;
     }
 

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

@@ -30,4 +30,13 @@ public boolean isGlobal()
     {
         return false;
     }
+
+    /**
+     * Invoked when this builder was created but its {@link #build()} will never run, e.g. because the executor
+     * rejected the submission (typically during shutdown). Implementations must release any resources they
+     * reserved at construction time. Default is a no-op.
+     */
+    public void onNotExecuted()
+    {
+    }
 }

src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuilder.java

@@ -78,6 +78,11 @@ public class StorageAttachedIndexBuilder extends SecondaryIndexBuilder
     private final boolean isFullRebuild;
     private final boolean isInitialBuild;
 
+    // True when the caller reserved the per-sstable rebuild status (via SSTableReader#streamRebuildState) before
+    // constructing this builder, e.g. a full/partial rebuild that must exclude concurrent entire-sstable streaming
+    // (CASSANDRA-21520). When true, this builder is responsible for releasing that status when it finishes.
+    private final boolean ownsStreamRebuildStatus;
+
     private final SortedMap<SSTableReader, Set<StorageAttachedIndex>> sstables;
 
     private long bytesProcessed = 0;
@@ -87,13 +92,23 @@ public class StorageAttachedIndexBuilder extends SecondaryIndexBuilder
                                 SortedMap<SSTableReader, Set<StorageAttachedIndex>> sstables,
                                 boolean isFullRebuild,
                                 boolean isInitialBuild)
+    {
+        this(group, sstables, isFullRebuild, isInitialBuild, false);
+    }
+
+    StorageAttachedIndexBuilder(StorageAttachedIndexGroup group,
+                                SortedMap<SSTableReader, Set<StorageAttachedIndex>> sstables,
+                                boolean isFullRebuild,
+                                boolean isInitialBuild,
+                                boolean ownsStreamRebuildStatus)
     {
         this.group = group;
         this.metadata = group.metadata();
         this.sstables = sstables;
         this.tracker = group.table().getTracker();
         this.isFullRebuild = isFullRebuild;
         this.isInitialBuild = isInitialBuild;
+        this.ownsStreamRebuildStatus = ownsStreamRebuildStatus;
         this.totalSizeInBytes = sstables.keySet().stream().mapToLong(SSTableReader::uncompressedLength).sum();
     }
 
@@ -104,23 +119,46 @@ public void build()
                                               isInitialBuild ? "initial" : "non-initial",
                                               isFullRebuild ? "full" : "partial")));
 
-        for (Map.Entry<SSTableReader, Set<StorageAttachedIndex>> e : sstables.entrySet())
+        try
         {
-            SSTableReader sstable = e.getKey();
-            Set<StorageAttachedIndex> indexes = e.getValue();
-
-            Set<StorageAttachedIndex> existing = validateIndexes(indexes, sstable.descriptor);
-            if (existing.isEmpty())
+            for (Map.Entry<SSTableReader, Set<StorageAttachedIndex>> e : sstables.entrySet())
             {
-                logger.debug(logMessage("{} dropped during index build"), indexes);
-                continue;
-            }
+                SSTableReader sstable = e.getKey();
+                Set<StorageAttachedIndex> indexes = e.getValue();
+
+                Set<StorageAttachedIndex> existing = validateIndexes(indexes, sstable.descriptor);
+                if (existing.isEmpty())
+                {
+                    logger.debug(logMessage("{} dropped during index build"), indexes);
+                    continue;
+                }
 
-            if (indexSSTable(sstable, existing))
-                return;
+                if (indexSSTable(sstable, existing))
+                    return;
+            }
+        }
+        finally
+        {
+            // Release the rebuild status reserved by the caller (see ownsStreamRebuildStatus), so entire-sstable
+            // streaming of these sstables is unblocked once the rebuild completes or aborts (CASSANDRA-21520).
+            releaseStreamRebuildStatus();
         }
     }
 
+    @Override
+    public void onNotExecuted()
+    {
+        // build() will never run (e.g. the executor rejected submission on shutdown), so release here instead to
+        // avoid leaving the reserved sstables stuck in the REBUILDING state (CASSANDRA-21520).
+        releaseStreamRebuildStatus();
+    }
+
+    private void releaseStreamRebuildStatus()
+    {
+        if (ownsStreamRebuildStatus)
+            sstables.keySet().forEach(sstable -> sstable.streamRebuildState().endRebuild());
+    }
+
     private String logMessage(String message)
     {
         return String.format("[%s.%s.*] %s", metadata.keyspace, metadata.name, message);

src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuildingSupport.java

@@ -21,6 +21,9 @@
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
 import java.util.NavigableMap;
 import java.util.Set;
 import java.util.TreeMap;
@@ -46,6 +49,8 @@ public SecondaryIndexBuilder getIndexBuildTask(ColumnFamilyStore cfs,
 
         assert group != null : "Index group does not exist for table " + cfs.keyspace + '.' + cfs.name;
 
+        // First resolve, without mutating anything, which sstables each index will rebuild.
+        Map<StorageAttachedIndex, Collection<SSTableReader>> targets = new LinkedHashMap<>();
         indexes.stream()
                .filter((i) -> i instanceof StorageAttachedIndex)
                .forEach((i) ->
@@ -62,11 +67,50 @@ public SecondaryIndexBuilder getIndexBuildTask(ColumnFamilyStore cfs,
                                                       .collect(Collectors.toList());
                             }
 
-                            group.dropIndexSSTables(ss, sai);
-
-                            ss.forEach(sstable -> sstables.computeIfAbsent(sstable, ignore -> new HashSet<>()).add(sai));
+                            targets.put(sai, ss);
                         });
 
-        return new StorageAttachedIndexBuilder(group, sstables, isFullRebuild, false);
+        // Reserve the per-sstable rebuild status for every unique target sstable BEFORE deleting any index
+        // components. An entire-sstable (zero-copy) stream reserves the same status when its outgoing file is
+        // constructed, so this ensures a rebuild cannot delete/rewrite SAI components underneath an in-flight
+        // stream (and vice versa the stream degrades to legacy). See CASSANDRA-21520. The returned builder owns
+        // these reservations and releases them when it finishes.
+        Set<SSTableReader> reserved = new LinkedHashSet<>();
+        for (Collection<SSTableReader> ss : targets.values())
+        {
+            for (SSTableReader sstable : ss)
+            {
+                if (reserved.contains(sstable))
+                    continue;
+                if (sstable.streamRebuildState().tryBeginRebuild())
+                {
+                    reserved.add(sstable);
+                }
+                else
+                {
+                    reserved.forEach(s -> s.streamRebuildState().endRebuild());
+                    throw new RuntimeException(String.format(
+                        "Cannot build SAI index on %s while entire-sstable (zero-copy) streaming is in progress.",
+                        sstable.descriptor));
+                }
+            }
+        }
+
+        // Now it is safe to drop existing components and assemble the build map.
+        try
+        {
+            targets.forEach((sai, ss) ->
+                            {
+                                group.dropIndexSSTables(ss, sai);
+                                ss.forEach(sstable -> sstables.computeIfAbsent(sstable, ignore -> new HashSet<>()).add(sai));
+                            });
+        }
+        catch (RuntimeException | Error e)
+        {
+            reserved.forEach(s -> s.streamRebuildState().endRebuild());
+            throw e;
+        }
+
+        return new StorageAttachedIndexBuilder(group, sstables, isFullRebuild, false, true);
     }
 }

src/java/org/apache/cassandra/io/sstable/format/SSTableReader.java

@@ -580,6 +580,16 @@ public <R, E extends Exception> R runWithLock(CheckedFunction<Descriptor, R, E>
         }
     }
 
+    /**
+     * @return the shared, per-descriptor coordination status between SAI index rebuilds and entire-sstable
+     * streaming. Lock-free to read (the field is {@code final} on the shared {@code GlobalTidy}); the returned
+     * object guards its own transitions.
+     */
+    public SSTableStreamRebuildState streamRebuildState()
+    {
+        return tidy.global.streamRebuildState;
+    }
+
     public void setReplaced()
     {
         synchronized (tidy.global)
@@ -1707,6 +1717,9 @@ static final class GlobalTidy implements Tidy
         private WeakReference<ScheduledFuture<?>> readMeterSyncFuture = NULL;
         // shared state managing if the logical sstable has been compacted; this is used in cleanup
         private volatile Runnable obsoletion;
+        // in-memory coordination between SAI index rebuilds and entire-sstable streaming for this sstable.
+        // final -> safely published; the object carries its own monitor for transitions.
+        final SSTableStreamRebuildState streamRebuildState = new SSTableStreamRebuildState();
 
         GlobalTidy(final SSTableReader reader)
         {