From 5e7b788fa0098cc7fee821f4da8e4693d28b0907 Mon Sep 17 00:00:00 2001 From: Pranav Shenoy Date: Wed, 26 Aug 2026 21:13:20 -0700 Subject: [PATCH] Adding Mutual exclusion logic for SAI index rebuilding and ZCS streaming --- .../db/compaction/CompactionManager.java | 9 +- .../db/streaming/CassandraOutgoingFile.java | 39 ++- .../db/streaming/CassandraStreamManager.java | 8 +- .../db/streaming/ComponentContext.java | 9 +- .../index/SecondaryIndexBuilder.java | 9 + .../sai/StorageAttachedIndexBuilder.java | 60 +++- .../StorageAttachedIndexBuildingSupport.java | 52 +++- .../io/sstable/format/SSTableReader.java | 13 + .../format/SSTableStreamRebuildState.java | 103 +++++++ ...eleteDuringEntireSSTableStreamingTest.java | 258 ++++++++++++++++ ...buildDuringEntireSSTableStreamingTest.java | 278 ++++++++++++++++++ .../sai/StreamingDuringIndexRebuildTest.java | 167 +++++++++++ .../format/SSTableStreamRebuildStateTest.java | 103 +++++++ 13 files changed, 1086 insertions(+), 22 deletions(-) create mode 100644 src/java/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildState.java create mode 100644 test/distributed/org/apache/cassandra/distributed/test/sai/IndexDeleteDuringEntireSSTableStreamingTest.java create mode 100644 test/distributed/org/apache/cassandra/distributed/test/sai/IndexRebuildDuringEntireSSTableStreamingTest.java create mode 100644 test/distributed/org/apache/cassandra/distributed/test/sai/StreamingDuringIndexRebuildTest.java create mode 100644 test/unit/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildStateTest.java diff --git a/src/java/org/apache/cassandra/db/compaction/CompactionManager.java b/src/java/org/apache/cassandra/db/compaction/CompactionManager.java index 67fa14f2a842..5f33cb0f7a30 100644 --- a/src/java/org/apache/cassandra/db/compaction/CompactionManager.java +++ b/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; } /** diff --git a/src/java/org/apache/cassandra/db/streaming/CassandraOutgoingFile.java b/src/java/org/apache/cassandra/db/streaming/CassandraOutgoingFile.java index f4f79a780cfb..0733022303cc 100644 --- a/src/java/org/apache/cassandra/db/streaming/CassandraOutgoingFile.java +++ b/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> ranges; + // guards against releasing the entire-sstable streaming status more than once (CASSANDRA-21520) + private boolean streamRebuildStatusReleased = false; public CassandraOutgoingFile(StreamOperation operation, Ref ref, List sections, List> normalizedRanges, @@ -66,9 +68,24 @@ public CassandraOutgoingFile(StreamOperation operation, Ref 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 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; diff --git a/src/java/org/apache/cassandra/db/streaming/CassandraStreamManager.java b/src/java/org/apache/cassandra/db/streaming/CassandraStreamManager.java index d3cee3d296b4..179b32c60aa3 100644 --- a/src/java/org/apache/cassandra/db/streaming/CassandraStreamManager.java +++ b/src/java/org/apache/cassandra/db/streaming/CassandraStreamManager.java @@ -91,6 +91,9 @@ public StreamReceiver createStreamReceiver(StreamSession session, List createOutgoingStreams(StreamSession session, RangesAtEndpoint replicas, TimeUUID pendingRepair, PreviewKind previewKind) { Refs 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 streams = new ArrayList<>(); try { final List> keyRanges = new ArrayList<>(replicas.size()); @@ -142,7 +145,6 @@ else if (pendingRepair == ActiveRepairService.NO_PENDING_REPAIR) List> normalizedFullRanges = Range.normalize(replicas.onlyFull().ranges()); List> normalizedAllRanges = Range.normalize(replicas.ranges()); //Create outgoing file streams for ranges possibly skipping repaired ranges in sstables - List streams = new ArrayList<>(refs.size()); for (SSTableReader sstable : refs) { List> 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; } diff --git a/src/java/org/apache/cassandra/db/streaming/ComponentContext.java b/src/java/org/apache/cassandra/db/streaming/ComponentContext.java index c03e7b4c3436..14f9a2139026 100644 --- a/src/java/org/apache/cassandra/db/streaming/ComponentContext.java +++ b/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; } diff --git a/src/java/org/apache/cassandra/index/SecondaryIndexBuilder.java b/src/java/org/apache/cassandra/index/SecondaryIndexBuilder.java index 73dc3345a250..221a33c07909 100644 --- a/src/java/org/apache/cassandra/index/SecondaryIndexBuilder.java +++ b/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() + { + } } diff --git a/src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuilder.java b/src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuilder.java index 4abd4e21bf15..16bd9e2f51b4 100644 --- a/src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuilder.java +++ b/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> sstables; private long bytesProcessed = 0; @@ -87,6 +92,15 @@ public class StorageAttachedIndexBuilder extends SecondaryIndexBuilder SortedMap> sstables, boolean isFullRebuild, boolean isInitialBuild) + { + this(group, sstables, isFullRebuild, isInitialBuild, false); + } + + StorageAttachedIndexBuilder(StorageAttachedIndexGroup group, + SortedMap> sstables, + boolean isFullRebuild, + boolean isInitialBuild, + boolean ownsStreamRebuildStatus) { this.group = group; this.metadata = group.metadata(); @@ -94,6 +108,7 @@ public class StorageAttachedIndexBuilder extends SecondaryIndexBuilder 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> e : sstables.entrySet()) + try { - SSTableReader sstable = e.getKey(); - Set indexes = e.getValue(); - - Set existing = validateIndexes(indexes, sstable.descriptor); - if (existing.isEmpty()) + for (Map.Entry> e : sstables.entrySet()) { - logger.debug(logMessage("{} dropped during index build"), indexes); - continue; - } + SSTableReader sstable = e.getKey(); + Set indexes = e.getValue(); + + Set 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); diff --git a/src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuildingSupport.java b/src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuildingSupport.java index 7a13b4b186ab..ddc6e4720ca0 100644 --- a/src/java/org/apache/cassandra/index/sai/StorageAttachedIndexBuildingSupport.java +++ b/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> 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 reserved = new LinkedHashSet<>(); + for (Collection 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); } } diff --git a/src/java/org/apache/cassandra/io/sstable/format/SSTableReader.java b/src/java/org/apache/cassandra/io/sstable/format/SSTableReader.java index ab3171200f17..3adedb7b620c 100644 --- a/src/java/org/apache/cassandra/io/sstable/format/SSTableReader.java +++ b/src/java/org/apache/cassandra/io/sstable/format/SSTableReader.java @@ -580,6 +580,16 @@ public R runWithLock(CheckedFunction } } + /** + * @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> 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) { diff --git a/src/java/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildState.java b/src/java/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildState.java new file mode 100644 index 000000000000..bd2190d8d386 --- /dev/null +++ b/src/java/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildState.java @@ -0,0 +1,103 @@ +/* + * 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.io.sstable.format; + +import com.google.common.annotations.VisibleForTesting; + +/** + * In-memory coordination between entire-sstable (zero-copy) streaming and SAI index rebuilds for a single + * logical sstable. Held on the shared, per-{@code Descriptor} {@code GlobalTidy}, so all reader instances of + * the same sstable observe one authoritative status, and a crash naturally resets it on restart. + * + *

Multiple entire-sstable streams of the same sstable can run concurrently ({@code ZCS_STREAMING} with a + * reference count), but a rebuild is exclusive ({@code REBUILDING}) and mutually exclusive with streaming. All + * transitions are guarded by this object's own monitor, held only for the brief check-and-set; the status flag + * itself is what persists for the duration of an operation and enforces the exclusion. + */ +public class SSTableStreamRebuildState +{ + public enum State + { + NORMAL, + REBUILDING, + ZCS_STREAMING + } + + private State state = State.NORMAL; + private int zcsStreamCount = 0; + + /** + * Attempt to begin an entire-sstable stream. Fails only if a rebuild is in progress. + * + * @return true if streaming may proceed (caller must later call {@link #endStreaming()}), false if a rebuild + * is active and the caller should fall back to legacy streaming. + */ + public synchronized boolean tryBeginStreaming() + { + if (state == State.REBUILDING) + return false; + state = State.ZCS_STREAMING; + zcsStreamCount++; + return true; + } + + /** + * Release one entire-sstable stream. Returns to {@code NORMAL} when the last stream ends. Defensive against + * over-release so cleanup paths cannot corrupt the state. + */ + public synchronized void endStreaming() + { + if (zcsStreamCount > 0 && --zcsStreamCount == 0) + state = State.NORMAL; + } + + /** + * Attempt to begin a rebuild. Fails if any stream is in progress or another rebuild is already running. + * + * @return true if the rebuild may proceed (caller must later call {@link #endRebuild()}), false otherwise. + */ + public synchronized boolean tryBeginRebuild() + { + if (state != State.NORMAL) + return false; + state = State.REBUILDING; + return true; + } + + /** + * Release the rebuild. Defensive: only resets if currently {@code REBUILDING}. + */ + public synchronized void endRebuild() + { + if (state == State.REBUILDING) + state = State.NORMAL; + } + + @VisibleForTesting + public synchronized State state() + { + return state; + } + + @VisibleForTesting + public synchronized int zcsStreamCount() + { + return zcsStreamCount; + } +} diff --git a/test/distributed/org/apache/cassandra/distributed/test/sai/IndexDeleteDuringEntireSSTableStreamingTest.java b/test/distributed/org/apache/cassandra/distributed/test/sai/IndexDeleteDuringEntireSSTableStreamingTest.java new file mode 100644 index 000000000000..f7984a21e8a5 --- /dev/null +++ b/test/distributed/org/apache/cassandra/distributed/test/sai/IndexDeleteDuringEntireSSTableStreamingTest.java @@ -0,0 +1,258 @@ +/* + * 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.sai; + +import java.nio.channels.FileChannel; +import java.util.StringJoiner; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.dynamic.loading.ClassLoadingStrategy; +import net.bytebuddy.implementation.MethodDelegation; +import net.bytebuddy.implementation.bind.annotation.AllArguments; +import net.bytebuddy.implementation.bind.annotation.FieldValue; +import net.bytebuddy.implementation.bind.annotation.SuperCall; + +import org.junit.Test; + +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.streaming.CassandraEntireSSTableStreamWriter; +import org.apache.cassandra.db.streaming.ComponentContext; +import org.apache.cassandra.db.streaming.ComponentManifest; +import org.apache.cassandra.distributed.Cluster; +import org.apache.cassandra.distributed.api.IInvokableInstance; +import org.apache.cassandra.distributed.api.NodeToolResult; +import org.apache.cassandra.distributed.test.TestBaseImpl; +import org.apache.cassandra.io.sstable.Component; +import org.apache.cassandra.io.sstable.format.SSTableReader; +import org.apache.cassandra.io.util.File; +import org.apache.cassandra.utils.concurrent.CountDownLatch; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.apache.cassandra.distributed.api.Feature.GOSSIP; +import static org.apache.cassandra.distributed.api.Feature.NETWORK; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Verifies the CASSANDRA-21520 safety net for entire-sstable (zero-copy) streaming: if a streamed SAI component + * changes size underneath an in-flight stream, the sender must fail the stream rather than ship bytes that + * disagree with the already-sent {@link ComponentManifest}. + * + *

While the sender is paused after sending the manifest (component -> advertised size) but before opening the + * component channels, the test directly truncates one of the streamed SAI component files. SAI components are not + * part of {@code mutableComponents()}, so they are not hard-linked and are streamed from the live file. When the + * writer opens the channel, the size check in {@link ComponentContext#channel} detects the mismatch and fails the + * stream. This check is a real exception (not an {@code assert}), so it is effective in production where assertions + * are disabled.

+ * + *

The failed streaming transaction must be rolled back on the receiver, which must not expose any partial or + * size-mismatched data.

+ */ +public class IndexDeleteDuringEntireSSTableStreamingTest extends TestBaseImpl +{ + private static final String TABLE = "tbl"; + private static final String INDEX = "sai_idx"; + private static final int ROWS = 200; + + @Test + public void testIndexDeleteWhileEntireSSTableStreaming() throws Exception + { + try (Cluster cluster = init(Cluster.build(2) + .withDataDirCount(1) + .withConfig(c -> c.with(NETWORK, GOSSIP) + .set("stream_entire_sstables", true) + .set("autocompaction_on_startup_enabled", false)) + // Only the sender (node1) needs the streaming pause installed. + .withInstanceInitializer((cl, num) -> { + if (num == 1) + BBHelper.install(cl); + }) + .start())) + { + cluster.disableAutoCompaction(KEYSPACE); + cluster.schemaChange(withKeyspace("CREATE TABLE %s." + TABLE + " (pk int PRIMARY KEY, v text)")); + cluster.schemaChange(withKeyspace("CREATE INDEX " + INDEX + " ON %s." + TABLE + "(v) USING 'sai'")); + SAIUtil.waitForIndexQueryable(cluster, KEYSPACE, INDEX); + + IInvokableInstance node1 = cluster.get(1); // sender + IInvokableInstance node2 = cluster.get(2); // receiver + + // Populate the sender only (executeInternal bypasses replication) and flush to a single sstable that + // carries the SAI components. node2 owns the same range (RF == cluster size) but has no data yet. + for (int i = 0; i < ROWS; i++) + node1.executeInternal(withKeyspace("INSERT INTO %s." + TABLE + "(pk, v) VALUES (?, ?)"), i, "v" + i); + node1.flush(KEYSPACE); + + // Arm the sender-side pause so only the upcoming ZCS stream is intercepted. + node1.runOnInstance(() -> BBHelper.armed.set(true)); + + // Kick off the streaming asynchronously: node2 rebuilds its data from node1. The call blocks until + // streaming completes, and streaming is about to block on the sender, so run it off-thread. + ExecutorService executor = Executors.newSingleThreadExecutor(); + try + { + Future streaming = + executor.submit(() -> node2.nodetoolResult("rebuild", "--keyspace", KEYSPACE)); + + // Wait until the sender has serialized and flushed the manifest and is paused before opening the + // outgoing component channels. + node1.runOnInstance(() -> BBHelper.manifestSent.awaitUninterruptibly(2, TimeUnit.MINUTES)); + + // While the stream is paused mid-flight (status is already ZCS_STREAMING for this sstable), a + // concurrent SAI rebuild must fail fast rather than mutate components underneath the in-flight + // stream. CASSANDRA-21520. + // While the stream is paused mid-flight (after the manifest advertising each component's size has + // been sent, before the component channels are opened), directly truncate one of the streamed SAI + // component files on the sender. SAI components are not part of mutableComponents(), so they are not + // hard-linked and are streamed from the live file. When the writer opens the channel the on-disk + // size no longer matches the size advertised in the manifest, and the production size check in + // ComponentContext.channel() must fail the stream rather than ship corrupt bytes. This check must be + // effective even with assertions disabled (it is a real exception, not an assert). CASSANDRA-21520. + node1.runOnInstance(IndexDeleteDuringEntireSSTableStreamingTest::truncateStreamedSaiComponent); + + // Let the paused stream continue now that a component has been mutated underneath it. + node1.runOnInstance(() -> BBHelper.proceed.decrement()); + + NodeToolResult result = streaming.get(3, TimeUnit.MINUTES); + result.asserts().failure(); + assertThat(node1.logs() + .grep("Entire sstable streaming expects .* file size to be .* but got") + .getResult()) + .describedAs("Expected the ComponentContext size check to fail the stream") + .isNotEmpty(); + } + finally + { + executor.shutdownNow(); + } + + for (int i = 0; i < ROWS; i++) + { + assertThat(node2.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE pk = ?"), i)) + .describedAs("Receiver must not retain data from the failed stream (pk=%d)", i) + .isEmpty(); + } + } + } + + private static void truncateStreamedSaiComponent() + { + ColumnFamilyStore cfs = Keyspace.open(KEYSPACE).getColumnFamilyStore(TABLE); + SSTableReader sstable = cfs.getLiveSSTables().iterator().next(); + for (Component component : sstable.getStreamingComponents()) + { + // Only mutate a component that is streamed from the live file (not a hard-linked mutable component and + // not a primary component), i.e. a SAI index component, so the size divergence hits ComponentContext. + if (sstable.descriptor.getFormat().primaryComponents().contains(component)) + continue; + if (sstable.descriptor.getFormat().mutableComponents().contains(component)) + continue; + File file = sstable.descriptor.fileFor(component); + if (!file.exists() || file.length() == 0) + continue; + try (FileChannel channel = file.newReadWriteChannel()) + { + channel.truncate(Math.max(0, channel.size() - 1)); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + return; + } + throw new IllegalStateException("No streamed SAI component available to mutate"); + } + + public static class BBHelper + { + // Per-instance state, resolved in the sender's classloader. + static final AtomicBoolean armed = new AtomicBoolean(false); + static final AtomicBoolean paused = new AtomicBoolean(false); + static final CountDownLatch manifestSent = CountDownLatch.newCountDownLatch(1); + static final CountDownLatch proceed = CountDownLatch.newCountDownLatch(1); + + /** + * Intercepts {@link CassandraEntireSSTableStreamWriter#write}. On entry the caller has already serialized and + * flushed the component manifest, so this is the "manifest sent, channels not yet opened" point. Print the + * manifest that was just sent (component -> advertised size), then pause once so the test can truncate a + * streamed SAI component before the component channels are created. + */ + @SuppressWarnings("unused") + public static void write(@FieldValue("manifest") ComponentManifest manifest, @SuperCall Callable zuper) throws Exception + { + if (armed.get() && paused.compareAndSet(false, true)) + { + StringJoiner sj = new StringJoiner(", ", "{", "}"); + for (Component component : manifest.components()) + sj.add(component.name + '=' + manifest.sizeOf(component)); + System.out.println("[ZCS-VERIFY] manifest sent (component=advertisedSize): " + sj); + + manifestSent.decrement(); + proceed.awaitUninterruptibly(2, TimeUnit.MINUTES); + } + zuper.call(); + } + + /** + * Intercepts {@link ComponentContext#channel}. This is invoked for every component right before it is copied + * onto the wire. Print the component name, the size advertised in the manifest, and the actual on-disk size + * at stream time. + */ + @SuppressWarnings("unused") + public static FileChannel channel(@AllArguments Object[] args, @SuperCall Callable zuper) throws Exception + { + Component component = (Component) args[1]; + long advertised = (Long) args[2]; + try + { + FileChannel channel = zuper.call(); + System.out.println("[ZCS-VERIFY] streaming component " + component.name + + ": advertisedSize=" + advertised + ", actualSize=" + channel.size()); + return channel; + } + catch (Throwable t) + { + System.out.println("[ZCS-VERIFY] streaming component " + component.name + + ": advertisedSize=" + advertised + " -> FAILED TO OPEN: " + t); + throw t; + } + } + + public static void install(ClassLoader classLoader) + { + new ByteBuddy().rebase(CassandraEntireSSTableStreamWriter.class) + .method(named("write").and(takesArguments(1))) + .intercept(MethodDelegation.to(BBHelper.class)) + .make() + .load(classLoader, ClassLoadingStrategy.Default.INJECTION); + new ByteBuddy().rebase(ComponentContext.class) + .method(named("channel").and(takesArguments(3))) + .intercept(MethodDelegation.to(BBHelper.class)) + .make() + .load(classLoader, ClassLoadingStrategy.Default.INJECTION); + } + } +} diff --git a/test/distributed/org/apache/cassandra/distributed/test/sai/IndexRebuildDuringEntireSSTableStreamingTest.java b/test/distributed/org/apache/cassandra/distributed/test/sai/IndexRebuildDuringEntireSSTableStreamingTest.java new file mode 100644 index 000000000000..73def08c8a45 --- /dev/null +++ b/test/distributed/org/apache/cassandra/distributed/test/sai/IndexRebuildDuringEntireSSTableStreamingTest.java @@ -0,0 +1,278 @@ +/* + * 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.sai; + +import java.nio.channels.FileChannel; +import java.util.Collections; +import java.util.StringJoiner; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.dynamic.loading.ClassLoadingStrategy; +import net.bytebuddy.implementation.MethodDelegation; +import net.bytebuddy.implementation.bind.annotation.AllArguments; +import net.bytebuddy.implementation.bind.annotation.FieldValue; +import net.bytebuddy.implementation.bind.annotation.SuperCall; + +import org.junit.Test; + +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.streaming.CassandraEntireSSTableStreamWriter; +import org.apache.cassandra.db.streaming.ComponentContext; +import org.apache.cassandra.db.streaming.ComponentManifest; +import org.apache.cassandra.distributed.Cluster; +import org.apache.cassandra.distributed.api.IInvokableInstance; +import org.apache.cassandra.distributed.api.NodeToolResult; +import org.apache.cassandra.distributed.test.TestBaseImpl; +import org.apache.cassandra.io.sstable.Component; +import org.apache.cassandra.utils.concurrent.CountDownLatch; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.apache.cassandra.distributed.api.Feature.GOSSIP; +import static org.apache.cassandra.distributed.api.Feature.NETWORK; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Verifies the CASSANDRA-21520 contract: a concurrent SAI index rebuild fails fast while an + * entire-sstable (zero-copy) stream is in flight, and the stream completes with correct data on the receiver. + * + *

Entire-sstable-streaming (ZCS) first records every streamable component and its size in a + * {@link org.apache.cassandra.db.streaming.ComponentManifest}, serializes and flushes that manifest to the + * peer, and only then copies each component file verbatim onto the wire. The sender now acquires the per-sstable + * ZCS streaming status when the outgoing file is constructed, before the manifest is advertised, so a concurrent + * SAI rebuild must be rejected rather than deleting and rewriting index components underneath the in-flight stream.

+ * + *

This test forces exactly that interleaving: + *

    + *
  1. creates an SAI index and flushes a single sstable on the sender (node1),
    2. + *
  2. triggers ZCS from the receiver (node2) and pauses the sender right after the manifest has been sent,
    3. + *
  3. attempts a blocking SAI index rebuild on the sender while the stream is paused,
    4. + *
  4. asserts that the rebuild is rejected, then resumes streaming and verifies the receiver has correct data.
    5. + *
+ */ +public class IndexRebuildDuringEntireSSTableStreamingTest extends TestBaseImpl +{ + private static final String TABLE = "tbl"; + private static final String INDEX = "sai_idx"; + private static final int ROWS = 200; + + @Test + public void testIndexRebuildWhileEntireSSTableStreaming() throws Exception + { + try (Cluster cluster = init(Cluster.build(2) + .withDataDirCount(1) + .withConfig(c -> c.with(NETWORK, GOSSIP) + .set("stream_entire_sstables", true) + .set("autocompaction_on_startup_enabled", false)) + // Only the sender (node1) needs the streaming pause installed. + .withInstanceInitializer((cl, num) -> { + if (num == 1) + BBHelper.install(cl); + }) + .start())) + { + cluster.disableAutoCompaction(KEYSPACE); + cluster.schemaChange(withKeyspace("CREATE TABLE %s." + TABLE + " (pk int PRIMARY KEY, v text)")); + cluster.schemaChange(withKeyspace("CREATE INDEX " + INDEX + " ON %s." + TABLE + "(v) USING 'sai'")); + SAIUtil.waitForIndexQueryable(cluster, KEYSPACE, INDEX); + + IInvokableInstance node1 = cluster.get(1); // sender + IInvokableInstance node2 = cluster.get(2); // receiver + + // Populate the sender only (executeInternal bypasses replication) and flush to a single sstable that + // carries the SAI components. node2 owns the same range (RF == cluster size) but has no data yet. + for (int i = 0; i < ROWS; i++) + node1.executeInternal(withKeyspace("INSERT INTO %s." + TABLE + "(pk, v) VALUES (?, ?)"), i, "v" + i); + node1.flush(KEYSPACE); + + // Arm the sender-side pause so only the upcoming ZCS stream is intercepted. + node1.runOnInstance(() -> BBHelper.armed.set(true)); + + // Kick off the streaming asynchronously: node2 rebuilds its data from node1. The call blocks until + // streaming completes, and streaming is about to block on the sender, so run it off-thread. + ExecutorService executor = Executors.newSingleThreadExecutor(); + try + { + Future streaming = + executor.submit(() -> node2.nodetoolResult("rebuild", "--keyspace", KEYSPACE)); + + // Wait until the sender has serialized and flushed the manifest and is paused before opening the + // outgoing component channels. + node1.runOnInstance(() -> BBHelper.manifestSent.awaitUninterruptibly(2, TimeUnit.MINUTES)); + + // While the stream is paused mid-flight (status is already ZCS_STREAMING for this sstable), a + // concurrent SAI rebuild must fail fast rather than mutate components underneath the in-flight + // stream. CASSANDRA-21520. + boolean rebuildRejected = node1.callOnInstance(IndexRebuildDuringEntireSSTableStreamingTest::rebuildAndReturnRejected); + assertThat(rebuildRejected) + .describedAs("SAI rebuild must fail fast while entire-sstable streaming is in progress") + .isTrue(); + assertThat(node1.logs() + .grep("while entire-sstable \\(zero-copy\\) streaming is in progress") + .getResult()) + .describedAs("Expected the rebuild to be rejected due to active streaming") + .isNotEmpty(); + + // Let the paused stream continue; its components were never mutated, so it completes cleanly. + node1.runOnInstance(() -> BBHelper.proceed.decrement()); + + NodeToolResult result = streaming.get(3, TimeUnit.MINUTES); + result.asserts().success(); + } + finally + { + executor.shutdownNow(); + } + + for (int i = 0; i < ROWS; i++) + { + Object[][] byPrimaryKey = node2.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE pk = ?"), i); + assertThat(byPrimaryKey.length) + .describedAs("Receiver must retain streamed data (pk=%d)", i) + .isEqualTo(1); + Object[][] byIndex = node2.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE v = ?"), "v" + i); + assertThat(byIndex.length) + .describedAs("Receiver SAI index must return streamed row (v%d)", i) + .isEqualTo(1); + assertThat(byIndex[0][0]).isEqualTo(i); + } + + // A rejected full rebuild leaves the index marked for rebuild - this is consistent with existing + // full-rebuild failure semantics, since markIndexesBuilding makes the index non-queryable before the + // rejection happens. The sender's data is intact, and once streaming has completed and released the + // status a fresh rebuild succeeds - proving the rejection was clean and the SAI components on the sender + // were never deleted or corrupted underneath the stream. + for (int i = 0; i < ROWS; i++) + { + Object[][] onSender = node1.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE pk = ?"), i); + assertThat(onSender.length).describedAs("Sender data must be intact (pk=%d)", i).isEqualTo(1); + } + + node1.runOnInstance(() -> { + ColumnFamilyStore cfs = Keyspace.open(KEYSPACE).getColumnFamilyStore(TABLE); + cfs.indexManager.rebuildIndexesBlocking(Collections.singleton(INDEX)); + }); + SAIUtil.waitForIndexQueryable(cluster, KEYSPACE, INDEX); + + for (int i = 0; i < ROWS; i++) + { + Object[][] onSender = node1.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE v = ?"), "v" + i); + assertThat(onSender.length).describedAs("Sender index must be queryable after re-rebuild (v%d)", i).isEqualTo(1); + assertThat(onSender[0][0]).isEqualTo(i); + } + } + } + + private static boolean rebuildAndReturnRejected() + { + try + { + // Drive the real user-facing rebuild path (nodetool rebuild_index). It routes through + // SecondaryIndexManager.buildIndexesBlocking -> StorageAttachedIndexBuildingSupport.getIndexBuildTask, + // which reserves the per-sstable rebuild status BEFORE deleting any SAI components. With a stream in + // flight that reservation must fail fast, so no component is deleted underneath the stream. + ColumnFamilyStore cfs = Keyspace.open(KEYSPACE).getColumnFamilyStore(TABLE); + cfs.indexManager.rebuildIndexesBlocking(Collections.singleton(INDEX)); + return false; + } + catch (Throwable t) + { + org.slf4j.LoggerFactory.getLogger(IndexRebuildDuringEntireSSTableStreamingTest.class).error(t.getMessage(), t); + return true; + } + } + + public static class BBHelper + { + // Per-instance state, resolved in the sender's classloader. + static final AtomicBoolean armed = new AtomicBoolean(false); + static final AtomicBoolean paused = new AtomicBoolean(false); + static final CountDownLatch manifestSent = CountDownLatch.newCountDownLatch(1); + static final CountDownLatch proceed = CountDownLatch.newCountDownLatch(1); + + /** + * Intercepts {@link CassandraEntireSSTableStreamWriter#write}. On entry the caller has already serialized and + * flushed the component manifest, so this is the "manifest sent, channels not yet opened" point. Print the + * manifest that was just sent (component -> advertised size), then pause once so the test can mutate the SAI + * components before the component channels are created. + */ + @SuppressWarnings("unused") + public static void write(@FieldValue("manifest") ComponentManifest manifest, @SuperCall Callable zuper) throws Exception + { + if (armed.get() && paused.compareAndSet(false, true)) + { + StringJoiner sj = new StringJoiner(", ", "{", "}"); + for (Component component : manifest.components()) + sj.add(component.name + '=' + manifest.sizeOf(component)); + System.out.println("[ZCS-VERIFY] manifest sent (component=advertisedSize): " + sj); + + manifestSent.decrement(); + proceed.awaitUninterruptibly(2, TimeUnit.MINUTES); + } + zuper.call(); + } + + /** + * Intercepts {@link ComponentContext#channel}. This is invoked for every component right before it is copied + * onto the wire. Print the component name, the size advertised in the manifest, and the actual on-disk size + * at stream time so the divergence introduced by the concurrent rebuild is visible (the component whose size + * changed will not match, which is what fails the stream). + */ + @SuppressWarnings("unused") + public static FileChannel channel(@AllArguments Object[] args, @SuperCall Callable zuper) throws Exception + { + Component component = (Component) args[1]; + long advertised = (Long) args[2]; + try + { + FileChannel channel = zuper.call(); + System.out.println("[ZCS-VERIFY] streaming component " + component.name + + ": advertisedSize=" + advertised + ", actualSize=" + channel.size()); + return channel; + } + catch (Throwable t) + { + System.out.println("[ZCS-VERIFY] streaming component " + component.name + + ": advertisedSize=" + advertised + " -> MISMATCH: " + t.getMessage()); + throw t; + } + } + + public static void install(ClassLoader classLoader) + { + new ByteBuddy().rebase(CassandraEntireSSTableStreamWriter.class) + .method(named("write").and(takesArguments(1))) + .intercept(MethodDelegation.to(BBHelper.class)) + .make() + .load(classLoader, ClassLoadingStrategy.Default.INJECTION); + new ByteBuddy().rebase(ComponentContext.class) + .method(named("channel").and(takesArguments(3))) + .intercept(MethodDelegation.to(BBHelper.class)) + .make() + .load(classLoader, ClassLoadingStrategy.Default.INJECTION); + } + } +} diff --git a/test/distributed/org/apache/cassandra/distributed/test/sai/StreamingDuringIndexRebuildTest.java b/test/distributed/org/apache/cassandra/distributed/test/sai/StreamingDuringIndexRebuildTest.java new file mode 100644 index 000000000000..b966d070d36f --- /dev/null +++ b/test/distributed/org/apache/cassandra/distributed/test/sai/StreamingDuringIndexRebuildTest.java @@ -0,0 +1,167 @@ +/* + * 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.sai; + +import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.dynamic.loading.ClassLoadingStrategy; +import net.bytebuddy.implementation.MethodDelegation; +import net.bytebuddy.implementation.bind.annotation.SuperCall; + +import org.junit.Test; + +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.streaming.CassandraEntireSSTableStreamWriter; +import org.apache.cassandra.distributed.Cluster; +import org.apache.cassandra.distributed.api.IInvokableInstance; +import org.apache.cassandra.distributed.api.NodeToolResult; +import org.apache.cassandra.distributed.test.TestBaseImpl; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.apache.cassandra.distributed.api.Feature.GOSSIP; +import static org.apache.cassandra.distributed.api.Feature.NETWORK; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * CASSANDRA-21520 (reverse direction): when an SAI index rebuild is in progress on the sender, a new + * entire-sstable (zero-copy) stream of the same sstable must degrade to legacy section-based streaming, so the + * rebuild and the stream never mutate/ship the SAI components concurrently. The receiver still ends up with a + * correct, queryable index (it rebuilds it locally from the legacy stream). + * + *

The test deterministically models an in-progress rebuild by reserving the per-sstable rebuild status on the + * sender (exactly what {@code StorageAttachedIndexBuildingSupport.getIndexBuildTask} does before dropping SAI + * components), then streams from the receiver. It asserts that the sender never invokes + * {@link CassandraEntireSSTableStreamWriter#write} (i.e. it fell back to legacy) and that the receiver ends up + * with the streamed data and a queryable SAI index.

+ */ +public class StreamingDuringIndexRebuildTest extends TestBaseImpl +{ + private static final String TABLE = "tbl"; + private static final String INDEX = "sai_idx"; + private static final int ROWS = 200; + + @Test + public void testStreamingFallsBackToLegacyDuringRebuild() throws Exception + { + try (Cluster cluster = init(Cluster.build(2) + .withDataDirCount(1) + .withConfig(c -> c.with(NETWORK, GOSSIP) + .set("stream_entire_sstables", true) + .set("autocompaction_on_startup_enabled", false)) + // Only the sender (node1) needs the entire-sstable writer counter installed. + .withInstanceInitializer((cl, num) -> { + if (num == 1) + BBHelper.install(cl); + }) + .start())) + { + cluster.disableAutoCompaction(KEYSPACE); + cluster.schemaChange(withKeyspace("CREATE TABLE %s." + TABLE + " (pk int PRIMARY KEY, v text)")); + cluster.schemaChange(withKeyspace("CREATE INDEX " + INDEX + " ON %s." + TABLE + "(v) USING 'sai'")); + SAIUtil.waitForIndexQueryable(cluster, KEYSPACE, INDEX); + + IInvokableInstance node1 = cluster.get(1); // sender + IInvokableInstance node2 = cluster.get(2); // receiver + + for (int i = 0; i < ROWS; i++) + node1.executeInternal(withKeyspace("INSERT INTO %s." + TABLE + "(pk, v) VALUES (?, ?)"), i, "v" + i); + node1.flush(KEYSPACE); + + // Model an in-progress SAI rebuild on the sender by reserving the per-sstable rebuild status for every + // live sstable. While this is held, an entire-sstable stream of the same sstable must fall back to legacy. + node1.runOnInstance(() -> { + ColumnFamilyStore cfs = Keyspace.open(KEYSPACE).getColumnFamilyStore(TABLE); + cfs.getLiveSSTables().forEach(sstable -> { + if (!sstable.streamRebuildState().tryBeginRebuild()) + throw new IllegalStateException("Could not reserve rebuild status for " + sstable.descriptor); + }); + }); + + node1.runOnInstance(() -> BBHelper.armed.set(true)); + + try + { + // node2 rebuilds its data from node1. The sender must NOT use entire-sstable streaming while a + // rebuild holds the sstable, so it degrades to legacy streaming (blocking call). + NodeToolResult result = node2.nodetoolResult("rebuild", "--keyspace", KEYSPACE); + result.asserts().success(); + + int entireWrites = node1.callOnInstance(() -> BBHelper.entireSSTableWrites.get()); + assertThat(entireWrites) + .describedAs("Entire-sstable streaming must not run while a rebuild holds the sstable") + .isZero(); + } + finally + { + // Release the reserved status so the sender returns to normal. + node1.runOnInstance(() -> { + ColumnFamilyStore cfs = Keyspace.open(KEYSPACE).getColumnFamilyStore(TABLE); + cfs.getLiveSSTables().forEach(sstable -> sstable.streamRebuildState().endRebuild()); + }); + } + + // The receiver built its SAI index locally from the legacy stream and must expose the streamed rows. + SAIUtil.waitForIndexQueryable(cluster, KEYSPACE, INDEX); + for (int i = 0; i < ROWS; i++) + { + Object[][] byPrimaryKey = node2.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE pk = ?"), i); + assertThat(byPrimaryKey.length) + .describedAs("Receiver must retain streamed data (pk=%d)", i) + .isEqualTo(1); + Object[][] byIndex = node2.executeInternal(withKeyspace("SELECT pk FROM %s." + TABLE + " WHERE v = ?"), "v" + i); + assertThat(byIndex.length) + .describedAs("Receiver SAI index must return streamed row (v%d)", i) + .isEqualTo(1); + assertThat(byIndex[0][0]).isEqualTo(i); + } + } + } + + public static class BBHelper + { + static final AtomicBoolean armed = new AtomicBoolean(false); + static final AtomicInteger entireSSTableWrites = new AtomicInteger(0); + + /** + * Counts invocations of {@link CassandraEntireSSTableStreamWriter#write}. Any invocation while armed means + * entire-sstable (zero-copy) streaming was used rather than the expected legacy fallback. + */ + @SuppressWarnings("unused") + public static void write(@SuperCall Callable zuper) throws Exception + { + if (armed.get()) + entireSSTableWrites.incrementAndGet(); + zuper.call(); + } + + public static void install(ClassLoader classLoader) + { + new ByteBuddy().rebase(CassandraEntireSSTableStreamWriter.class) + .method(named("write").and(takesArguments(1))) + .intercept(MethodDelegation.to(BBHelper.class)) + .make() + .load(classLoader, ClassLoadingStrategy.Default.INJECTION); + } + } +} diff --git a/test/unit/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildStateTest.java b/test/unit/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildStateTest.java new file mode 100644 index 000000000000..e8e83bd08f68 --- /dev/null +++ b/test/unit/org/apache/cassandra/io/sstable/format/SSTableStreamRebuildStateTest.java @@ -0,0 +1,103 @@ +/* + * 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.io.sstable.format; + +import org.junit.Test; + +import org.apache.cassandra.io.sstable.format.SSTableStreamRebuildState.State; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class SSTableStreamRebuildStateTest +{ + @Test + public void startsNormal() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + assertEquals(State.NORMAL, s.state()); + assertEquals(0, s.zcsStreamCount()); + } + + @Test + public void multipleStreamsAllowedAndCounted() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + assertTrue(s.tryBeginStreaming()); + assertTrue(s.tryBeginStreaming()); + assertEquals(State.ZCS_STREAMING, s.state()); + assertEquals(2, s.zcsStreamCount()); + + s.endStreaming(); + assertEquals(State.ZCS_STREAMING, s.state()); + assertEquals(1, s.zcsStreamCount()); + + s.endStreaming(); + assertEquals(State.NORMAL, s.state()); + assertEquals(0, s.zcsStreamCount()); + } + + @Test + public void rebuildBlockedWhileStreaming() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + assertTrue(s.tryBeginStreaming()); + assertFalse(s.tryBeginRebuild()); + assertEquals(State.ZCS_STREAMING, s.state()); + } + + @Test + public void streamingBlockedWhileRebuilding() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + assertTrue(s.tryBeginRebuild()); + assertFalse(s.tryBeginStreaming()); + assertEquals(State.REBUILDING, s.state()); + assertEquals(0, s.zcsStreamCount()); + } + + @Test + public void rebuildExcludesSecondRebuild() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + assertTrue(s.tryBeginRebuild()); + assertFalse(s.tryBeginRebuild()); + } + + @Test + public void rebuildResetsToNormal() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + assertTrue(s.tryBeginRebuild()); + s.endRebuild(); + assertEquals(State.NORMAL, s.state()); + assertTrue(s.tryBeginStreaming()); + } + + @Test + public void endIsDefensiveAgainstOverRelease() + { + SSTableStreamRebuildState s = new SSTableStreamRebuildState(); + s.endStreaming(); // no-op, must not go negative + s.endRebuild(); // no-op + assertEquals(State.NORMAL, s.state()); + assertEquals(0, s.zcsStreamCount()); + } +}