CASSANALYTICS-160: Determine whether mutation tracking is enabled for keyspace for bulk writes

CHANGES.txt

@@ -1,5 +1,6 @@
 0.5.0
 -----
+ * 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)

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);
+    }
 }

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java

@@ -145,6 +145,37 @@ public void validateClOrFail(TokenRangeMapping<RingInstance> tokenRangeMapping)
         validateClOrFail(tokenRangeMapping, true);
     }
 
+    /**
+     * Validates bulk write success for a tracked keyspace.
+     *
+     * <p>
+     * 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)}.
+     * </p>
+     *
+     * @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<CommitResult> commitResults,
+                                                       ReplicaAwareFailureHandler<RingInstance> failureHandler,
+                                                       Logger logger,
+                                                       String phase,
+                                                       JobInfo job,
+                                                       TokenRangeMapping<RingInstance> tokenRangeMapping,
+                                                       ClusterInfo cluster)
+    {
+        commitResults.forEach(cr -> updateFailureHandler(cr, phase, failureHandler));
+        validateClOrFail(tokenRangeMapping, failureHandler, logger, phase, job, cluster);
+    }
+
     public void validateClOrFail(TokenRangeMapping<RingInstance> tokenRangeMapping, boolean refreshInstanceAvailability)
     {
         if (refreshInstanceAvailability)

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<RingInstance> getTokenRangeMapping(boolean cached)
     {

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<BigInteger> range,
-                                                   ReplicaAwareFailureHandler<RingInstance> failureHandler,
-                                                   ExecutorService executorService)
+    public StreamSession<TransportContext.DirectDataBulkWriterContext> createStreamSession(
+            BulkWriterContext writerContext,
+            String sessionId,
+            SortedSSTableWriter sstableWriter,
+            Range<BigInteger> range,
+            ReplicaAwareFailureHandler<RingInstance> 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());
     }
 }

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

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamSession.java

@@ -49,7 +49,7 @@
 public class DirectStreamSession extends StreamSession<TransportContext.DirectDataBulkWriterContext>
 {
     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<CommitResult> commit(DirectStreamResult streamResult) throws ExecutionException, InterruptedException
+    protected List<CommitResult> commit(DirectStreamResult streamResult) throws ExecutionException, InterruptedException
     {
         try (CommitCoordinator cc = CommitCoordinator.commit(writerContext, transportContext, streamResult))
         {

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.
+     *
+     * <p>
+     * <b>Note:</b> 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.
+     * </p>
+     */
+    public boolean isCoordinator()
+    {
+        // TODO(CASSANALYTICS-160): read from ReplicaMetadata.isCoordinator() once Sidecar exposes it
+        return false;
+    }
+
     // Used only in tests
     @Override
     public String token()

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TrackedDirectStreamSession.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.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.spark.bulkwriter.token.ReplicaAwareFailureHandler;
+
+/**
+ * Stream session for bulk writes to keyspaces with mutation tracking enabled.
+ *
+ * <p>
+ * Tracked stream session uploads and triggers import on the <em>coordinator node only</em>. 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.
+ * </p>
+ *
+ * <p>
+ * The coordinator selection policy is:
+ * <ol>
+ *   <li>Prefer the replica designated as coordinator by Sidecar (via {@link RingInstance#isCoordinator()}).</li>
+ *   <li>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.</li>
+ * </ol>
+ * Superseded coordinators are marked failed via the {@link ReplicaAwareFailureHandler} so that retries pick the
+ * next available candidate.
+ * </p>
+ */
+public class TrackedDirectStreamSession extends DirectStreamSession
+{
+    private static final Logger LOGGER = LoggerFactory.getLogger(TrackedDirectStreamSession.class);
+
+    public TrackedDirectStreamSession(BulkWriterContext writerContext,
+                                      SortedSSTableWriter sstableWriter,
+                                      TransportContext.DirectDataBulkWriterContext transportContext,
+                                      String sessionID,
+                                      Range<BigInteger> tokenRange,
+                                      ReplicaAwareFailureHandler<RingInstance> failureHandler,
+                                      ExecutorService executorService)
+    {
+        super(writerContext, sstableWriter, transportContext, sessionID, tokenRange, failureHandler, executorService);
+    }
+
+    /**
+     * Returns a single-element list containing the coordinator node for the token range.
+     *
+     * <p>
+     * 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.
+     * </p>
+     *
+     * <p>
+     * Selection order:
+     * <ol>
+     *   <li>The replica marked as coordinator by Sidecar ({@link RingInstance#isCoordinator()}).</li>
+     *   <li>First non-failed replica (fallback when Sidecar does not yet surface the coordinator flag).</li>
+     * </ol>
+     * </p>
+     *
+     * @return list with exactly one coordinator instance, or an empty list when all replicas have failed
+     */
+    @Override
+    @VisibleForTesting
+    List<RingInstance> getReplicas()
+    {
+        Set<RingInstance> failedInstances = failureHandler.getFailedInstances();
+        List<RingInstance> 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.
+     *
+     * <p>
+     * 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.
+     * </p>
+     */
+    @Override
+    protected StreamResult doFinalizeStream()
+    {
+        sendRemainingSSTables();
+
+        List<RingInstance> coordinator = getReplicas();
+        DirectStreamResult streamResult = new DirectStreamResult(sessionID,
+                                                                 tokenRange,
+                                                                 errors,
+                                                                 new ArrayList<>(coordinator),
+                                                                 sstableWriter.rowCount(),
+                                                                 sstableWriter.bytesWritten());
+        List<CommitResult> 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;
+    }
+}