From 9f6d7b4c021ee7042bd54ea3ebb366be6a9b31ab Mon Sep 17 00:00:00 2001 From: Saranya Krishnakumar Date: Mon, 1 Jun 2026 15:09:05 -0700 Subject: [PATCH 1/2] CASSANALYTICS-160: Analytics should identify if a keyspace is tracked to determine appropriate stream session to use for bulk writes --- CHANGES.txt | 1 + .../cassandra/spark/utils/CqlUtils.java | 32 ++++ .../bulkwriter/CassandraClusterInfo.java | 11 ++ .../CassandraDirectDataTransportContext.java | 31 +++- .../spark/bulkwriter/ClusterInfo.java | 10 ++ .../TrackedDirectStreamSession.java | 70 ++++++++ .../bulkwriter/CassandraClusterInfoTest.java | 81 ++++++++++ ...ssandraDirectDataTransportContextTest.java | 151 ++++++++++++++++++ .../bulkwriter/MockBulkWriterContext.java | 12 ++ .../cassandra/spark/utils/CqlUtilsTest.java | 37 +++++ .../cassandra/bridge/CassandraBridge.java | 13 ++ 11 files changed, 441 insertions(+), 8 deletions(-) create mode 100644 cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java create mode 100644 cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContextTest.java diff --git a/CHANGES.txt b/CHANGES.txt index 974b9ceaa..21f58d70f 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,5 +1,6 @@ 0.5.0 ----- + * Analytics should identify if a keyspace is tracked to determine appropriate stream session to use for bulk writes (CASSANALYTICS-160) * Upgrade sidecar version to 0.4.0 * Exclude IP address from RingInstance equality so node replacement does not fail bulk write jobs (CASSANALYTICS-175) * Regenerate bloom filters for CQLSSTableWriter (CASSANALYTICS-167) diff --git a/cassandra-analytics-common/src/main/java/org/apache/cassandra/spark/utils/CqlUtils.java b/cassandra-analytics-common/src/main/java/org/apache/cassandra/spark/utils/CqlUtils.java index 9d989be51..8511dbaa1 100644 --- a/cassandra-analytics-common/src/main/java/org/apache/cassandra/spark/utils/CqlUtils.java +++ b/cassandra-analytics-common/src/main/java/org/apache/cassandra/spark/utils/CqlUtils.java @@ -52,6 +52,9 @@ public final class CqlUtils "max_index_interval" ); private static final Pattern REPLICATION_FACTOR_PATTERN = Pattern.compile("WITH REPLICATION = (\\{[^\\}]*\\})"); + private static final String TRACKED_REPLICATION_TYPE = "tracked"; + private static final Pattern REPLICATION_TYPE_PATTERN = Pattern.compile("replication_type\\s*=\\s*'(\\w+)'", + Pattern.CASE_INSENSITIVE); // Initialize a mapper allowing single quotes to process the RF string from the CREATE KEYSPACE statement private static final ObjectMapper MAPPER = new ObjectMapper().configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, true); private static final Pattern ESCAPED_WHITESPACE_PATTERN = Pattern.compile("(\\\\r|\\\\n|\\\\r\\n)+"); @@ -296,4 +299,33 @@ public static boolean isTimeRangeFilterSupported(String compactionStrategy) { return compactionStrategy == null || compactionStrategy.endsWith("TimeWindowCompactionStrategy"); } + + /** + * Extracts replication type from create schema statement + * + * @param schemaStr full cluster schema string as returned by Sidecar + * @param keyspace name of the keyspace to check + * @return {@code true} if keyspace is tracked {@code false} otherwise + */ + public static String extractReplicationType(@NotNull String schemaStr, @NotNull String keyspace) + { + String createKeyspaceSchema = extractKeyspaceSchema(schemaStr, keyspace); + Matcher matcher = REPLICATION_TYPE_PATTERN.matcher(createKeyspaceSchema); + if (matcher.find()) + { + return matcher.group(1); + } + return null; + } + + /** + * Returns {@code true} if {@code replication_type = 'tracked'} in create statement otherwise {@code false} + * + * @param replicationType replication type extracted from create statement + * @return {@code true} if replication type is tracked {@code false} otherwise + */ + public static boolean isTracked(String replicationType) + { + return TRACKED_REPLICATION_TYPE.equalsIgnoreCase(replicationType); + } } diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java index a9c120871..53ac97974 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java @@ -356,6 +356,17 @@ public ReplicationFactor replicationFactor() } } + @Override + public String getReplicationType() + { + String keyspaceSchema = getKeyspaceSchema(true); + if (keyspaceSchema == null) + { + throw new RuntimeException("Could not retrieve keyspace schema information for keyspace " + conf.keyspace); + } + return CqlUtils.extractReplicationType(keyspaceSchema, conf.keyspace); + } + @Override public TokenRangeMapping getTokenRangeMapping(boolean cached) { diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContext.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContext.java index bed36ac4b..fb82483a7 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContext.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContext.java @@ -46,13 +46,24 @@ public CassandraDirectDataTransportContext(@NotNull BulkWriterContext bulkWriter } @Override - public DirectStreamSession createStreamSession(BulkWriterContext writerContext, - String sessionId, - SortedSSTableWriter sstableWriter, - Range range, - ReplicaAwareFailureHandler failureHandler, - ExecutorService executorService) + public StreamSession createStreamSession( + BulkWriterContext writerContext, + String sessionId, + SortedSSTableWriter sstableWriter, + Range range, + ReplicaAwareFailureHandler failureHandler, + ExecutorService executorService) { + if (bridge().isTracked(clusterInfo.getReplicationType())) + { + return new TrackedDirectStreamSession(writerContext, + sstableWriter, + this, + sessionId, + range, + failureHandler, + executorService); + } return new DirectStreamSession(writerContext, sstableWriter, this, @@ -71,7 +82,11 @@ public DirectDataTransferApi dataTransferApi() // only invoke in constructor protected DirectDataTransferApi createDirectDataTransferApi() { - CassandraBridge bridge = CassandraBridgeFactory.get(clusterInfo.getLowestCassandraVersion()); - return new SidecarDataTransferApi(clusterInfo.getCassandraContext(), bridge, jobInfo); + return new SidecarDataTransferApi(clusterInfo.getCassandraContext(), bridge(), jobInfo); + } + + private CassandraBridge bridge() + { + return CassandraBridgeFactory.get(clusterInfo.getLowestCassandraVersion()); } } diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java index 64b654449..e4d8a0a0b 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java @@ -82,6 +82,16 @@ public interface ClusterInfo extends StartupValidatable */ ReplicationFactor replicationFactor(); + /** + * @return {@code replication_type} of the enclosing keyspace (e.g. {@code "tracked"}, {@code "untracked"}), + * or {@code null} if replication_type is absent + */ + @Nullable + default String getReplicationType() + { + return null; + } + CassandraContext getCassandraContext(); /** diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java new file mode 100644 index 000000000..166ef2387 --- /dev/null +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java @@ -0,0 +1,70 @@ +/* + * 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.spark.bulkwriter; + +import java.math.BigInteger; +import java.util.Set; +import java.util.concurrent.ExecutorService; + +import com.google.common.collect.Range; + +import org.apache.cassandra.bridge.SSTableDescriptor; +import org.apache.cassandra.spark.bulkwriter.token.ReplicaAwareFailureHandler; + +/** + * Stream session for bulk writes to keyspaces with mutation tracking enabled. + * + *

+ * Tracked stream session uploads and triggers import on the coordinator node only. Cassandra's coordinated + * transfer then propagates the data to all other replicas, avoiding the duplicate-row updates that would occur if each + * replica independently streamed the data to its peers. + *

+ */ +public class TrackedDirectStreamSession extends StreamSession +{ + public TrackedDirectStreamSession(BulkWriterContext writerContext, + SortedSSTableWriter sstableWriter, + TransportContext.DirectDataBulkWriterContext transportContext, + String sessionID, + Range tokenRange, + ReplicaAwareFailureHandler failureHandler, + ExecutorService executorService) + { + super(writerContext, sstableWriter, transportContext, sessionID, tokenRange, failureHandler, executorService); + } + + @Override + protected void onSSTablesProduced(Set sstables) + { + throw new UnsupportedOperationException("TrackedDirectStreamSession is not yet implemented"); + } + + @Override + protected StreamResult doFinalizeStream() + { + throw new UnsupportedOperationException("TrackedDirectStreamSession is not yet implemented"); + } + + @Override + protected void sendRemainingSSTables() + { + throw new UnsupportedOperationException("TrackedDirectStreamSession is not yet implemented"); + } +} diff --git a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfoTest.java b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfoTest.java index 5ed4c5221..763082a36 100644 --- a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfoTest.java +++ b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfoTest.java @@ -30,8 +30,10 @@ import o.a.c.sidecar.client.shaded.common.response.TimeSkewResponse; import org.apache.cassandra.spark.bulkwriter.token.TokenRangeMapping; import org.apache.cassandra.spark.exception.TimeSkewTooLargeException; +import org.apache.cassandra.spark.utils.CqlUtils; import static org.apache.cassandra.spark.TestUtils.range; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatNoException; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; @@ -75,6 +77,50 @@ public static CassandraClusterInfo mockClusterInfoForTimeSkewTest(int allowanceM return new MockClusterInfoForTimeSkew(allowanceMinutes, remoteNow); } + @Test + void testGetTrackedReplicationType() + { + String schema = "CREATE KEYSPACE mykeyspace " + + "WITH replication = {'class': 'NetworkTopologyStrategy', 'dc1': '3'}" + + " AND durable_writes = true" + + " AND replication_type = 'tracked';"; + CassandraClusterInfo ci = mockClusterInfoWithSchema("mykeyspace", schema); + assertThat(ci.getReplicationType()) + .describedAs("Keyspace with replication_type = 'tracked' should return 'tracked'") + .isEqualTo("tracked"); + } + + @Test + void testGetUntrackedReplicationType() + { + String schema = "CREATE KEYSPACE k " + + "WITH replication = {'class': 'NetworkTopologyStrategy', 'dc1': '3'}" + + " AND durable_writes = true" + + " AND replication_type = 'untracked';"; + CassandraClusterInfo ci = mockClusterInfoWithSchema("k", schema); + assertThat(ci.getReplicationType()) + .describedAs("Keyspace with replication_type = 'untracked' should return 'untracked'") + .isEqualTo("untracked"); + } + + @Test + void testGetReplicationTypeReturnsNullWhenPropertyAbsent() + { + // Schema without replication_type — pre-mutation-tracking clusters + String schema = "CREATE KEYSPACE k " + + "WITH replication = {'class': 'NetworkTopologyStrategy', 'dc1': '3'}" + + " AND durable_writes = true;"; + CassandraClusterInfo ci = mockClusterInfoWithSchema("k", schema); + assertThat(ci.getReplicationType()) + .describedAs("Keyspace without replication_type property should return null") + .isNull(); + } + + private static CassandraClusterInfo mockClusterInfoWithSchema(String keyspace, String schemaStr) + { + return new MockClusterInfoForSchema(keyspace, schemaStr); + } + private static class MockClusterInfoForTimeSkew extends CassandraClusterInfo { private CassandraContext cassandraContext; @@ -107,4 +153,39 @@ private void mockCassandraContext(int allowanceMinutes, Instant remoteNow) when(cassandraContext.sidecarPort()).thenReturn(9043); } } + + /** + * Minimal ClusterInfo stub that returns a fixed keyspace schema string, used to test + * {@link CassandraClusterInfo#getReplicationType()}. + */ + private static class MockClusterInfoForSchema extends CassandraClusterInfo + { + private final String keyspaceName; + private final String schemaStr; + + MockClusterInfoForSchema(String keyspace, String schemaStr) + { + super((BulkSparkConf) null); + this.keyspaceName = keyspace; + this.schemaStr = schemaStr; + } + + @Override + protected CassandraContext buildCassandraContext() + { + return mock(CassandraContext.class, RETURNS_DEEP_STUBS); + } + + @Override + public String getKeyspaceSchema(boolean cached) + { + return schemaStr; + } + + @Override + public String getReplicationType() + { + return CqlUtils.extractReplicationType(schemaStr, keyspaceName); + } + } } diff --git a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContextTest.java b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContextTest.java new file mode 100644 index 000000000..0bd2170bb --- /dev/null +++ b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraDirectDataTransportContextTest.java @@ -0,0 +1,151 @@ +/* + * 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.spark.bulkwriter; + +import java.math.BigInteger; +import java.nio.file.Path; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +import com.google.common.collect.BoundType; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Range; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import org.apache.cassandra.spark.bulkwriter.token.MultiClusterReplicaAwareFailureHandler; +import org.apache.cassandra.spark.bulkwriter.token.ReplicaAwareFailureHandler; +import org.apache.cassandra.spark.bulkwriter.token.TokenRangeMapping; +import org.apache.cassandra.spark.data.ReplicationFactor; +import org.apache.cassandra.spark.utils.XXHash32DigestAlgorithm; + +import static org.apache.cassandra.spark.data.ReplicationFactor.ReplicationStrategy.NetworkTopologyStrategy; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests that {@link CassandraDirectDataTransportContext#createStreamSession}. + */ +public class CassandraDirectDataTransportContextTest +{ + @TempDir + private Path folder; + + private MockBulkWriterContext writerContext; + private MockTableWriter tableWriter; + private ExecutorService executor; + private Range range; + + @BeforeEach + public void setup() + { + ImmutableMap rfOptions = ImmutableMap.of("DC1", 3); + ReplicationFactor rf = new ReplicationFactor(NetworkTopologyStrategy, rfOptions); + TokenRangeMapping tokenRangeMapping = TokenRangeMappingUtils.buildTokenRangeMapping(0, rfOptions, 12); + writerContext = new MockBulkWriterContext(tokenRangeMapping); + writerContext.setReplicationFactor(rf); + tableWriter = new MockTableWriter(folder); + range = Range.range(BigInteger.valueOf(101L), BoundType.CLOSED, BigInteger.valueOf(199L), BoundType.CLOSED); + executor = Executors.newSingleThreadExecutor(); + } + + @AfterEach + public void tearDown() + { + executor.shutdownNow(); + } + + @Test + void createDirectStreamSessionForUntrackedKeyspace() + { + writerContext.setTrackedKeyspace(false); + CassandraDirectDataTransportContext transportContext = stubTransportContext(); + + StreamSession session = transportContext.createStreamSession( + writerContext, + "session-untracked", + new SortedSSTableWriter(tableWriter, folder, new XXHash32DigestAlgorithm(), 1), + range, + buildFailureHandler(), + executor); + + assertThat(session) + .describedAs("Untracked keyspace should produce a DirectStreamSession") + .isInstanceOf(DirectStreamSession.class); + } + + @Test + void createTrackedDirectStreamSessionForTrackedKeyspace() + { + writerContext.setTrackedKeyspace(true); + CassandraDirectDataTransportContext transportContext = stubTransportContext(); + + StreamSession session = transportContext.createStreamSession( + writerContext, + "session-tracked", + new SortedSSTableWriter(tableWriter, folder, new XXHash32DigestAlgorithm(), 1), + range, + buildFailureHandler(), + executor); + + assertThat(session) + .describedAs("Tracked keyspace should produce a TrackedDirectStreamSession") + .isInstanceOf(TrackedDirectStreamSession.class); + } + + @Test + void createDirectStreamSessionForAbsentReplicationType() + { + CassandraDirectDataTransportContext transportContext = stubTransportContext(); + + StreamSession session = transportContext.createStreamSession( + writerContext, + "session-null-type", + new SortedSSTableWriter(tableWriter, folder, new XXHash32DigestAlgorithm(), 1), + range, + buildFailureHandler(), + executor); + + assertThat(session) + .describedAs("Null replication type (pre-mutation-tracking) should produce a DirectStreamSession") + .isInstanceOf(DirectStreamSession.class); + } + + private CassandraDirectDataTransportContext stubTransportContext() + { + DirectDataTransferApi mockApi = + ((TransportContext.DirectDataBulkWriterContext) writerContext.transportContext()).dataTransferApi(); + + return new CassandraDirectDataTransportContext(writerContext) + { + @Override + protected DirectDataTransferApi createDirectDataTransferApi() + { + return mockApi; + } + }; + } + + private ReplicaAwareFailureHandler buildFailureHandler() + { + return new MultiClusterReplicaAwareFailureHandler<>(writerContext.cluster().getPartitioner()); + } +} diff --git a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext.java b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext.java index cb5564f0c..be05715c4 100644 --- a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext.java +++ b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext.java @@ -99,6 +99,7 @@ public interface CommitResultSupplier extends BiFunction, String, D private final UUID jobId; private boolean skipClean = false; + private String replicationType = null; public int refreshClusterInfoCallCount = 0; // CHECKSTYLE IGNORE: Public mutable field private final Map> uploads = new ConcurrentHashMap<>(); private final Map> commits = new ConcurrentHashMap<>(); @@ -204,6 +205,17 @@ public void setReplicationFactor(ReplicationFactor replicationFactor) this.replicationFactor = replicationFactor; } + @Override + public String getReplicationType() + { + return replicationType; + } + + public void setTrackedKeyspace(boolean tracked) + { + this.replicationType = tracked ? "tracked" : "untracked"; + } + @Override public CassandraContext getCassandraContext() { diff --git a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/utils/CqlUtilsTest.java b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/utils/CqlUtilsTest.java index 5bb363d97..58fb5c960 100644 --- a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/utils/CqlUtilsTest.java +++ b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/utils/CqlUtilsTest.java @@ -833,6 +833,43 @@ public void testTimeRangeFilterSupported() assertThat(isTimeRangeFilterSupported("")).isFalse(); } + @Test + public void testExtractReplicationType() + { + String schemaStrTracked + = "CREATE KEYSPACE k WITH REPLICATION " + + "= { 'class' : 'org.apache.cassandra.locator.NetworkTopologyStrategy', 'dc1': '3' }" + + " AND DURABLE_WRITES = true AND replication_type = 'tracked';"; + assertThat(CqlUtils.extractReplicationType(schemaStrTracked, "k")).isEqualTo("tracked"); + String schemaStrTrackedCaseInsensitive + = "CREATE KEYSPACE k WITH REPLICATION " + + "= { 'class' : 'org.apache.cassandra.locator.NetworkTopologyStrategy', 'dc1': '3' }" + + " AND DURABLE_WRITES = true AND REPLICATION_TYPE = 'Tracked';"; + assertThat(CqlUtils.extractReplicationType(schemaStrTrackedCaseInsensitive, "k")).isEqualTo("Tracked"); + String schemaStrUntracked + = "CREATE KEYSPACE k WITH REPLICATION " + + "= { 'class' : 'org.apache.cassandra.locator.NetworkTopologyStrategy', 'dc1': '3' }" + + " AND DURABLE_WRITES = true AND replication_type = 'untracked';"; + assertThat(CqlUtils.extractReplicationType(schemaStrUntracked, "k")).isEqualTo("untracked"); + String schemaStrNoReplicationType + = "CREATE KEYSPACE k WITH REPLICATION " + + "= { 'class' : 'org.apache.cassandra.locator.NetworkTopologyStrategy', 'dc1': '3' }" + + " AND DURABLE_WRITES = true;"; + assertThat(CqlUtils.extractReplicationType(schemaStrNoReplicationType, "k")).isNull(); + } + + @Test + public void testIsTracked() + { + assertThat(CqlUtils.isTracked("tracked")).isTrue(); + assertThat(CqlUtils.isTracked("TRACKED")).isTrue(); + assertThat(CqlUtils.isTracked("Tracked")).isTrue(); + assertThat(CqlUtils.isTracked("untracked")).isFalse(); + assertThat(CqlUtils.isTracked("random")).isFalse(); + assertThat(CqlUtils.isTracked(null)).isFalse(); + assertThat(CqlUtils.isTracked("")).isFalse(); + } + private static String loadFullSchemaSample() throws IOException { Path fullSchemaSampleFile = ResourceUtils.writeResourceToPath(CqlUtilsTest.class.getClassLoader(), tempPath, "cql/fullSchema.cql"); diff --git a/cassandra-bridge/src/main/java/org/apache/cassandra/bridge/CassandraBridge.java b/cassandra-bridge/src/main/java/org/apache/cassandra/bridge/CassandraBridge.java index 67c5bc7a0..d2d67ab92 100644 --- a/cassandra-bridge/src/main/java/org/apache/cassandra/bridge/CassandraBridge.java +++ b/cassandra-bridge/src/main/java/org/apache/cassandra/bridge/CassandraBridge.java @@ -62,6 +62,7 @@ import org.apache.cassandra.spark.sparksql.filters.SparkRangeFilter; import org.apache.cassandra.analytics.stats.Stats; import org.apache.cassandra.spark.sparksql.filters.SSTableTimeRangeFilter; +import org.apache.cassandra.spark.utils.CqlUtils; import org.apache.cassandra.spark.utils.TimeProvider; import org.apache.cassandra.util.CompressionUtil; import org.jetbrains.annotations.NotNull; @@ -167,6 +168,18 @@ public String maybeQuoteIdentifier(String identifier) return cassandraTypes().maybeQuoteIdentifier(identifier); } + /** + * @return {@code true} if mutation tracking is enabled given the replication type, {@code false} otherwise + */ + public boolean isTracked(String replicationType) + { + if (replicationType == null) + { + return false; + } + return CqlUtils.isTracked(replicationType); + } + // CQL Type Parsing public CqlField.CqlType readType(CqlField.CqlType.InternalType type, Input input) From ea9d43d943f88ce128e8f4027ca0a6d9b1c8bb45 Mon Sep 17 00:00:00 2001 From: Saranya Krishnakumar Date: Mon, 6 Jul 2026 16:10:16 -0700 Subject: [PATCH 2/2] Update JIRA title --- CHANGES.txt | 2 +- .../spark/bulkwriter/BulkWriteValidator.java | 31 ++ .../spark/bulkwriter/DirectStreamSession.java | 4 +- .../spark/bulkwriter/RingInstance.java | 16 + .../TrackedDirectStreamSession.java | 117 ++++++- .../TrackedDirectStreamSessionTest.java | 330 ++++++++++++++++++ 6 files changed, 487 insertions(+), 13 deletions(-) create mode 100644 cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSessionTest.java diff --git a/CHANGES.txt b/CHANGES.txt index 21f58d70f..19a658dc1 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,6 +1,6 @@ 0.5.0 ----- - * Analytics should identify if a keyspace is tracked to determine appropriate stream session to use for bulk writes (CASSANALYTICS-160) + * Determine whether mutation tracking is enabled for keyspace for bulk writes (CASSANALYTICS-160) * Upgrade sidecar version to 0.4.0 * Exclude IP address from RingInstance equality so node replacement does not fail bulk write jobs (CASSANALYTICS-175) * Regenerate bloom filters for CQLSSTableWriter (CASSANALYTICS-167) diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java index 02fa7e660..a5cc3f528 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java @@ -145,6 +145,37 @@ public void validateClOrFail(TokenRangeMapping tokenRangeMapping) validateClOrFail(tokenRangeMapping, true); } + /** + * Validates bulk write success for a tracked keyspace. + * + *

+ * For tracked keyspaces, only the coordinator node acknowledges the import; Cassandra's + * coordinated transfer takes care of distributing the write to all other replicas at the + * requested consistency level internally. This method therefore only checks that the + * coordinator commit succeeded, rather than the per-replica replica count used by + * {@link #validateClOrFail(TokenRangeMapping, ReplicaAwareFailureHandler, Logger, String, JobInfo, ClusterInfo)}. + *

+ * + * @param commitResults results from committing SSTables on the coordinator + * @param failureHandler failure handler used to record per-range failures + * @param logger logger + * @param phase current write phase label (used for error messages) + * @param job current job info + * @param tokenRangeMapping ring topology used to resolve failure ranges + * @param cluster cluster info + */ + public static void validateTrackedKeyspaceCLOrFail(List commitResults, + ReplicaAwareFailureHandler failureHandler, + Logger logger, + String phase, + JobInfo job, + TokenRangeMapping tokenRangeMapping, + ClusterInfo cluster) + { + commitResults.forEach(cr -> updateFailureHandler(cr, phase, failureHandler)); + validateClOrFail(tokenRangeMapping, failureHandler, logger, phase, job, cluster); + } + public void validateClOrFail(TokenRangeMapping tokenRangeMapping, boolean refreshInstanceAvailability) { if (refreshInstanceAvailability) diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamSession.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamSession.java index 7af617090..6cbb2e207 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamSession.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamSession.java @@ -49,7 +49,7 @@ public class DirectStreamSession extends StreamSession { private static final Logger LOGGER = LoggerFactory.getLogger(DirectStreamSession.class); - private static final String WRITE_PHASE = "UploadAndCommit"; + protected static final String WRITE_PHASE = "UploadAndCommit"; private final AtomicInteger nextSSTableIdx = new AtomicInteger(1); private final DirectDataTransferApi directDataTransferApi; @@ -247,7 +247,7 @@ private void sendSSTableComponent(Path componentFile, recordStreamedFile(componentFile); } - private List commit(DirectStreamResult streamResult) throws ExecutionException, InterruptedException + protected List commit(DirectStreamResult streamResult) throws ExecutionException, InterruptedException { try (CommitCoordinator cc = CommitCoordinator.commit(writerContext, transportContext, streamResult)) { diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java index 5e5ed60e9..3116dcaae 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java @@ -72,6 +72,22 @@ public RingInstance(ReplicaMetadata replica) this(replica, null); } + /** + * Returns {@code true} if Sidecar has designated this replica as the coordinated-transfer coordinator for its + * token range. + * + *

+ * Note: This always returns {@code false} until the Sidecar {@code ReplicaMetadata} response includes an + * {@code isCoordinator} field (see CASSANALYTICS-160 Layer 1b). At that point this method should be updated to + * read the flag from {@code ReplicaMetadata} and store it in the field below. + *

+ */ + public boolean isCoordinator() + { + // TODO(CASSANALYTICS-160): read from ReplicaMetadata.isCoordinator() once Sidecar exposes it + return false; + } + // Used only in tests @Override public String token() diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java index 166ef2387..f17a6f075 100644 --- a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java +++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.java @@ -20,12 +20,19 @@ package org.apache.cassandra.spark.bulkwriter; import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; import java.util.Set; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; +import java.util.stream.Collectors; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.Range; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; -import org.apache.cassandra.bridge.SSTableDescriptor; import org.apache.cassandra.spark.bulkwriter.token.ReplicaAwareFailureHandler; /** @@ -35,10 +42,23 @@ * Tracked stream session uploads and triggers import on the coordinator node only. Cassandra's coordinated * transfer then propagates the data to all other replicas, avoiding the duplicate-row updates that would occur if each * replica independently streamed the data to its peers. + *

+ * *

+ * The coordinator selection policy is: + *

    + *
  1. Prefer the replica designated as coordinator by Sidecar (via {@link RingInstance#isCoordinator()}).
    2. + *
  2. If no such replica is found (e.g. older Sidecar without the flag), fall back to the first non-failed replica + * in the token range.
    3. + *
+ * Superseded coordinators are marked failed via the {@link ReplicaAwareFailureHandler} so that retries pick the + * next available candidate. + *

*/ -public class TrackedDirectStreamSession extends StreamSession +public class TrackedDirectStreamSession extends DirectStreamSession { + private static final Logger LOGGER = LoggerFactory.getLogger(TrackedDirectStreamSession.class); + public TrackedDirectStreamSession(BulkWriterContext writerContext, SortedSSTableWriter sstableWriter, TransportContext.DirectDataBulkWriterContext transportContext, @@ -50,21 +70,98 @@ public TrackedDirectStreamSession(BulkWriterContext writerContext, super(writerContext, sstableWriter, transportContext, sessionID, tokenRange, failureHandler, executorService); } + /** + * Returns a single-element list containing the coordinator node for the token range. + * + *

+ * The coordinator is the node Cassandra designates to receive the bulk import and coordinate its propagation + * to all replicas. Analytics only needs to upload to this one node; Cassandra handles the rest. + *

+ * + *

+ * Selection order: + *

    + *
  1. The replica marked as coordinator by Sidecar ({@link RingInstance#isCoordinator()}).
    2. + *
  2. First non-failed replica (fallback when Sidecar does not yet surface the coordinator flag).
    3. + *
+ *

+ * + * @return list with exactly one coordinator instance, or an empty list when all replicas have failed + */ @Override - protected void onSSTablesProduced(Set sstables) + @VisibleForTesting + List getReplicas() { - throw new UnsupportedOperationException("TrackedDirectStreamSession is not yet implemented"); + Set failedInstances = failureHandler.getFailedInstances(); + List candidates = tokenRangeMapping.getSubRanges(tokenRange) + .asMapOfRanges() + .values() + .stream() + .flatMap(Collection::stream) + .distinct() + .filter(instance -> !failedInstances.contains(instance)) + .collect(Collectors.toList()); + + LOGGER.debug("[{}]: Selecting coordinator from {} candidate(s) for range {}", + sessionID, candidates.size(), tokenRange); + + // Prefer the Sidecar-designated coordinator; fall back to any available replica + return candidates.stream() + .filter(RingInstance::isCoordinator) + .findFirst() + .map(List::of) + .orElseGet(() -> candidates.isEmpty() ? List.of() + : List.of(candidates.get(0))); } + /** + * Finalizes the stream by uploading any remaining SSTables to the coordinator node and then triggering + * a coordinated import on that node. + * + *

+ * Unlike {@link DirectStreamSession#doFinalizeStream()}, this does not validate CL against individual replica + * counts. Instead it delegates validation to + * {@link BulkWriteValidator#validateTrackedKeyspaceCLOrFail}, which checks that the coordinator accepted + * the import. Cassandra's coordinated transfer guarantees that CL is satisfied internally. + *

+ */ @Override protected StreamResult doFinalizeStream() { - throw new UnsupportedOperationException("TrackedDirectStreamSession is not yet implemented"); - } + sendRemainingSSTables(); - @Override - protected void sendRemainingSSTables() - { - throw new UnsupportedOperationException("TrackedDirectStreamSession is not yet implemented"); + List coordinator = getReplicas(); + DirectStreamResult streamResult = new DirectStreamResult(sessionID, + tokenRange, + errors, + new ArrayList<>(coordinator), + sstableWriter.rowCount(), + sstableWriter.bytesWritten()); + List commitResults; + try + { + commitResults = commit(streamResult); + } + catch (Exception e) + { + if (e instanceof InterruptedException) + { + Thread.currentThread().interrupt(); + } + throw new RuntimeException(e); + } + streamResult.setCommitResults(commitResults); + LOGGER.debug("[{}]: TrackedStreamResult: {}", sessionID, streamResult); + + // For tracked keyspaces, the coordinator accepts the import and Cassandra distributes it + // internally; we only need to verify that the coordinator commit succeeded. + BulkWriteValidator.validateTrackedKeyspaceCLOrFail(commitResults, + failureHandler, + LOGGER, + WRITE_PHASE, + writerContext.job(), + tokenRangeMapping, + writerContext.cluster()); + return streamResult; } } diff --git a/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSessionTest.java b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSessionTest.java new file mode 100644 index 000000000..124f60110 --- /dev/null +++ b/cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSessionTest.java @@ -0,0 +1,330 @@ +/* + * 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.spark.bulkwriter; + +import java.io.IOException; +import java.math.BigInteger; +import java.nio.file.Path; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.stream.Collectors; + +import com.google.common.collect.BoundType; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Range; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import org.apache.cassandra.spark.bulkwriter.token.MultiClusterReplicaAwareFailureHandler; +import org.apache.cassandra.spark.bulkwriter.token.ReplicaAwareFailureHandler; +import org.apache.cassandra.spark.bulkwriter.token.TokenRangeMapping; +import org.apache.cassandra.spark.common.model.CassandraInstance; +import org.apache.cassandra.spark.data.ReplicationFactor; +import org.apache.cassandra.spark.exception.ConsistencyNotSatisfiedException; +import org.apache.cassandra.spark.utils.DigestAlgorithm; +import org.apache.cassandra.spark.utils.XXHash32DigestAlgorithm; +import org.jetbrains.annotations.NotNull; + +import static org.apache.cassandra.spark.data.ReplicationFactor.ReplicationStrategy.NetworkTopologyStrategy; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Tests for {@link TrackedDirectStreamSession}. + * + *

+ * These tests mirror {@link DirectStreamSessionTest} but exercise the coordinator-only + * upload path used when a keyspace has mutation tracking enabled. + *

+ * + *

+ * Integration test note: A full end-to-end integration test that actually creates + * a {@code replication_type = 'tracked'} keyspace requires a Cassandra build with CEP-45 + * mutation tracking support (not present in the standard 5.0.5 dtest jar currently in + * {@code dependencies/}). Once a compatible Cassandra build is available the integration + * test should be added under {@code cassandra-analytics-integration-tests} following the + * same pattern as {@code CassandraAnalyticsSimpleTest}, gated with + * {@code assumeThat(TestUtils.getDTestClusterVersion()...)} for the minimum version. + *

+ */ +public class TrackedDirectStreamSessionTest +{ + private static final Map COLUMN_BOUND_VALUES = + ImmutableMap.of("id", 0, "date", 1, "course", "course", "marks", 2); + + @TempDir + private Path folder; + + private static final int FILES_PER_SSTABLE = 8; + // For a tracked keyspace only 1 replica (the coordinator) is written to + private static final int COORDINATOR_COUNT = 1; + private static final int RF = 3; + + private MockBulkWriterContext writerContext; + private TransportContext.DirectDataBulkWriterContext transportContext; + private TokenRangeMapping tokenRangeMapping; + private MockScheduledExecutorService executor; + private MockTableWriter tableWriter; + private Range range; + private DigestAlgorithm digestAlgorithm; + + @BeforeEach + public void setup() + { + digestAlgorithm = new XXHash32DigestAlgorithm(); + range = Range.range(BigInteger.valueOf(101L), BoundType.CLOSED, BigInteger.valueOf(199L), BoundType.CLOSED); + ImmutableMap rfOptions = ImmutableMap.of("DC1", 3); + ReplicationFactor rf = new ReplicationFactor(NetworkTopologyStrategy, rfOptions); + tokenRangeMapping = TokenRangeMappingUtils.buildTokenRangeMapping(0, rfOptions, 12); + writerContext = new MockBulkWriterContext(tokenRangeMapping); + writerContext.setReplicationFactor(rf); + writerContext.setTrackedKeyspace(true); + tableWriter = new MockTableWriter(folder); + transportContext = (TransportContext.DirectDataBulkWriterContext) writerContext.transportContext(); + executor = new MockScheduledExecutorService(); + } + + // ----------------------------------------------------------------------- + // getReplicas — coordinator selection + // ----------------------------------------------------------------------- + + /** + * For a tracked keyspace, {@code getReplicas()} must return exactly one node even when + * multiple replicas cover the token range. + * + *

+ * Until Sidecar surfaces the {@code isCoordinator} flag, {@link RingInstance#isCoordinator()} + * always returns {@code false}, so the fallback path (first non-failed replica) is exercised. + *

+ */ + @Test + void testGetReplicasReturnsSingleCoordinator() + { + TrackedDirectStreamSession session = createStreamSession(SortedSSTableWriter::new); + List replicas = session.getReplicas(); + assertThat(replicas).as("tracked session must target exactly one coordinator node") + .hasSize(COORDINATOR_COUNT); + } + + /** + * When one replica in the token range has already failed, {@code getReplicas()} must still + * return exactly one coordinator — the first of the remaining healthy replicas. + */ + @Test + void testGetReplicasSkipsFailedInstances() + { + ReplicaAwareFailureHandler handler = replicaAwareFailureHandler(); + TrackedDirectStreamSession session = createStreamSession(SortedSSTableWriter::new, handler); + + // Mark the current coordinator as failed + List first = session.getReplicas(); + assertThat(first).hasSize(1); + handler.addFailure(range, first.get(0), "simulated coordinator failure"); + + // Re-query: the new coordinator must differ from the failed one + List second = session.getReplicas(); + assertThat(second).hasSize(1); + assertThat(second.get(0)).isNotEqualTo(first.get(0)); + } + + /** + * When all replicas for the token range have failed, {@code getReplicas()} must return an + * empty list (no coordinator available). + */ + @Test + void testGetReplicasReturnsEmptyWhenAllFailed() + { + ReplicaAwareFailureHandler handler = replicaAwareFailureHandler(); + TrackedDirectStreamSession session = createStreamSession(SortedSSTableWriter::new, handler); + + // Exhaust all candidates by marking them failed one by one + List all = tokenRangeMapping.getSubRanges(range) + .asMapOfRanges().values().stream() + .flatMap(Collection::stream) + .distinct() + .collect(Collectors.toList()); + all.forEach(instance -> handler.addFailure(range, instance, "forced failure")); + + List result = session.getReplicas(); + assertThat(result).isEmpty(); + } + + // ----------------------------------------------------------------------- + // Upload path — only coordinator receives SSTables + // ----------------------------------------------------------------------- + + /** + * A successful write must upload SSTables to exactly one node (the coordinator), + * not all RF replicas. + */ + @Test + void testWriteUploadsToCoordinatorOnly() throws IOException, ExecutionException, InterruptedException + { + TrackedDirectStreamSession session = createStreamSession(NonValidatingTestSortedSSTableWriter::new); + session.addRow(BigInteger.valueOf(102L), COLUMN_BOUND_VALUES); + StreamResult result = session.finalizeStreamAsync().get(); + + assertThat(result.rowCount).isEqualTo(1L); + executor.assertFuturesCalled(); + + int totalUploads = writerContext.getUploads().values().stream().mapToInt(Collection::size).sum(); + assertThat(totalUploads).as("only coordinator receives SSTable components") + .isEqualTo(COORDINATOR_COUNT * FILES_PER_SSTABLE); + + List uploadedInstances = writerContext.getUploads().keySet().stream() + .map(CassandraInstance::nodeName) + .collect(Collectors.toList()); + assertThat(uploadedInstances).as("exactly one coordinator is targeted").hasSize(COORDINATOR_COUNT); + } + + // ----------------------------------------------------------------------- + // Commit path — coordinator failure throws ConsistencyNotSatisfiedException + // ----------------------------------------------------------------------- + + /** + * When the coordinator import fails (commit returns failure), a + * {@link ConsistencyNotSatisfiedException} must be thrown. + * + *

+ * For tracked keyspaces the coordinator is the single node that matters; its failure + * means CL cannot be satisfied regardless of how many other replicas are up. + *

+ */ + @Test + void testCoordinatorCommitFailureThrowsConsistencyException() throws IOException + { + writerContext.setCommitResultSupplier((uuids, dc) -> + new DirectDataTransferApi.RemoteCommitResult(false, uuids, null, "simulated coordinator import failure")); + + TrackedDirectStreamSession session = createStreamSession(NonValidatingTestSortedSSTableWriter::new); + session.addRow(BigInteger.valueOf(102L), COLUMN_BOUND_VALUES); + assertThatThrownBy(() -> session.finalizeStreamAsync().get()) + .isInstanceOf(ExecutionException.class) + .hasCauseExactlyInstanceOf(ConsistencyNotSatisfiedException.class); + } + + /** + * When the coordinator import succeeds, no exception is thrown even though only + * 1 out of RF nodes acknowledged the write. + */ + @Test + void testCoordinatorCommitSuccessDoesNotThrow() throws IOException, ExecutionException, InterruptedException + { + writerContext.setCommitResultSupplier((uuids, dc) -> + new DirectDataTransferApi.RemoteCommitResult(true, null, uuids, "")); + + TrackedDirectStreamSession session = createStreamSession(NonValidatingTestSortedSSTableWriter::new); + session.addRow(BigInteger.valueOf(102L), COLUMN_BOUND_VALUES); + StreamResult result = session.finalizeStreamAsync().get(); + assertThat(result.rowCount).isEqualTo(1L); + } + + // ----------------------------------------------------------------------- + // Upload failure — clean is called on the coordinator + // ----------------------------------------------------------------------- + + /** + * When the SSTable upload to the coordinator fails, the coordinator's upload session + * must be cleaned up. + */ + @Test + void testUploadFailureCleansCooordinator() throws IOException, ExecutionException, InterruptedException + { + writerContext.setUploadSupplier(instance -> false); + TrackedDirectStreamSession session = createStreamSession(NonValidatingTestSortedSSTableWriter::new); + session.addRow(BigInteger.valueOf(102L), COLUMN_BOUND_VALUES); + Future fut = session.finalizeStreamAsync(); + // Upload fails → ConsistencyNotSatisfiedException is wrapped in ExecutionException + assertThatThrownBy(fut::get) + .isInstanceOf(ExecutionException.class) + .hasCauseInstanceOf(Exception.class); + + List cleaned = writerContext.getCleanedInstances().stream() + .map(CassandraInstance::nodeName) + .collect(Collectors.toList()); + assertThat(cleaned).as("coordinator session must be cleaned after upload failure") + .hasSize(COORDINATOR_COUNT); + } + + /** + * When {@code skipClean} is set, no cleanup call is made even after upload failure. + */ + @Test + void testUploadFailureSkipsCleanWhenConfigured() throws IOException + { + writerContext.setSkipCleanOnFailures(true); + writerContext.setUploadSupplier(instance -> false); + TrackedDirectStreamSession session = createStreamSession(NonValidatingTestSortedSSTableWriter::new); + session.addRow(BigInteger.valueOf(102L), COLUMN_BOUND_VALUES); + assertThatThrownBy(() -> session.finalizeStreamAsync().get()) + .isInstanceOf(ExecutionException.class); + + assertThat(writerContext.getCleanedInstances()).isEmpty(); + } + + /** + * An upload failure must trigger a {@code refreshClusterInfo()} call so the writer gets + * the latest topology to pick a fallback coordinator. + */ + @Test + void testUploadFailureRefreshesClusterInfo() throws IOException + { + writerContext.setUploadSupplier(instance -> false); + TrackedDirectStreamSession session = createStreamSession(NonValidatingTestSortedSSTableWriter::new); + session.addRow(BigInteger.valueOf(102L), COLUMN_BOUND_VALUES); + assertThatThrownBy(() -> session.finalizeStreamAsync().get()) + .isInstanceOf(ExecutionException.class); + + assertThat(writerContext.refreshClusterInfoCallCount) + .as("cluster info must be refreshed after coordinator upload failure") + .isGreaterThanOrEqualTo(1); + } + + // ----------------------------------------------------------------------- + // Helpers + // ----------------------------------------------------------------------- + + @NotNull + private ReplicaAwareFailureHandler replicaAwareFailureHandler() + { + return new MultiClusterReplicaAwareFailureHandler<>(writerContext.cluster().getPartitioner()); + } + + private TrackedDirectStreamSession createStreamSession(MockTableWriter.Creator writerCreator) + { + return createStreamSession(writerCreator, replicaAwareFailureHandler()); + } + + private TrackedDirectStreamSession createStreamSession(MockTableWriter.Creator writerCreator, + ReplicaAwareFailureHandler failureHandler) + { + return new TrackedDirectStreamSession(writerContext, + writerCreator.create(tableWriter, folder, digestAlgorithm, 1), + transportContext, + "sessionId", + range, + failureHandler, + executor); + } +}