[CASSANALYTICS-181] - Support Sidecar behind a load balancer for time…

CHANGES.txt

@@ -1,5 +1,6 @@
 0.5.0
 -----
+ * Support Sidecar behind a load balancer for time-skew validation in single-cluster and coordinated writes (CASSANALYTICS-181)
  * CDC reader stats silently dropped in SidecarCdcBuilder (CASSANALYTICS-191)
  * Add CapturePublishedSchema metric to SidecarCdcStats (CASSANALYTICS-189)
  * Expand list of architecture that supports unaligned access in FastByteOperations (CASSANALYTICS-188)

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

@@ -100,6 +100,6 @@ public String clusterId()
     @Override
     public ClusterInfo reconstruct()
     {
-        return new CassandraClusterInfo(this);
+        return CassandraClusterInfo.create(this);
     }
 }

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

@@ -146,6 +146,7 @@ public class BulkSparkConf implements Serializable
     public final double importCoordinatorTimeoutMultiplier;
     public boolean quoteIdentifiers;
     public final boolean skipSecondaryIndexCheck;
+    public final boolean sidecarBehindLoadBalancer;
     protected final String keystorePassword;
     protected final String keystorePath;
     protected final String keystoreBase64Encoded;
@@ -229,6 +230,7 @@ public BulkSparkConf(SparkConf conf, Map<String, String> options, @Nullable Logg
         this.timestamp = MapUtils.getOrDefault(options, WriterOptions.TIMESTAMP.name(), null);
         this.quoteIdentifiers = MapUtils.getBoolean(options, WriterOptions.QUOTE_IDENTIFIERS.name(), false, "quote identifiers");
         this.skipSecondaryIndexCheck = MapUtils.getBoolean(options, WriterOptions.SKIP_SECONDARY_INDEX_CHECK.name(), false, "skip secondary index check");
+        this.sidecarBehindLoadBalancer = MapUtils.getBoolean(options, WriterOptions.SIDECAR_BEHIND_LOAD_BALANCER.name(), false, "sidecar behind load balancer");
         int storageClientConcurrency = MapUtils.getInt(options, WriterOptions.STORAGE_CLIENT_CONCURRENCY.name(),
                                                        DEFAULT_STORAGE_CLIENT_CONCURRENCY, "storage client concurrency");
         long storageClientKeepAliveSeconds = MapUtils.getLong(options, WriterOptions.STORAGE_CLIENT_THREAD_KEEP_ALIVE_SECONDS.name(),

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

@@ -59,7 +59,7 @@ protected CassandraBulkWriterContext(@NotNull BulkWriterConfig config)
     @Override
     protected ClusterInfo buildClusterInfo()
     {
-        return new CassandraClusterInfo(bulkSparkConf());
+        return CassandraClusterInfo.create(bulkSparkConf());
     }
 
     @Override

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

@@ -138,6 +138,57 @@ public CassandraClusterInfo(BroadcastableClusterInfo broadcastable)
         this.allNodeSettingFutures = null;
     }
 
+    /**
+     * Creates a {@link CassandraClusterInfo} for a single cluster, selecting the concrete type based on the
+     * {@link BulkSparkConf#sidecarBehindLoadBalancer} flag.
+     *
+     * @param conf bulk write conf
+     * @return {@link LoadBalancedCassandraClusterInfo} when Sidecar is behind a load balancer, otherwise a plain
+     *         {@link CassandraClusterInfo}
+     */
+    public static CassandraClusterInfo create(BulkSparkConf conf)
+    {
+        return create(conf, null);
+    }
+
+    /**
+     * Creates a {@link CassandraClusterInfo}, selecting the concrete type based on the
+     * {@link BulkSparkConf#sidecarBehindLoadBalancer} flag. Kept centralized so the driver-side factory and the
+     * executor-side broadcast reconstruction stay in lockstep, for both single-cluster and coordinated writes.
+     *
+     * @param conf bulk write conf
+     * @param clusterId cluster identifier, or {@code null} for a single (non-coordinated) cluster
+     * @return {@link LoadBalancedCassandraClusterInfo} when Sidecar is behind a load balancer, otherwise a plain
+     *         {@link CassandraClusterInfo}
+     */
+    public static CassandraClusterInfo create(BulkSparkConf conf, String clusterId)
+    {
+        if (conf.sidecarBehindLoadBalancer)
+        {
+            LOGGER.info("Using LoadBalancedCassandraClusterInfo for load-balanced Sidecar. clusterId={}", clusterId);
+            return new LoadBalancedCassandraClusterInfo(conf, clusterId);
+        }
+        return new CassandraClusterInfo(conf, clusterId);
+    }
+
+    /**
+     * Reconstructs a {@link CassandraClusterInfo} on an executor from broadcast, selecting the concrete type based on
+     * the {@link BulkSparkConf#sidecarBehindLoadBalancer} flag so it matches the driver-side selection in
+     * {@link #create(BulkSparkConf, String)}.
+     *
+     * @param broadcastable the broadcastable cluster info from broadcast
+     * @return {@link LoadBalancedCassandraClusterInfo} when Sidecar is behind a load balancer, otherwise a plain
+     *         {@link CassandraClusterInfo}
+     */
+    public static CassandraClusterInfo create(BroadcastableClusterInfo broadcastable)
+    {
+        if (broadcastable.getConf().sidecarBehindLoadBalancer)
+        {
+            return new LoadBalancedCassandraClusterInfo(broadcastable);
+        }
+        return new CassandraClusterInfo(broadcastable);
+    }
+
     @Override
     public void checkBulkWriterIsEnabledOrThrow()
     {
@@ -246,16 +297,7 @@ void validateTimeSkewWithLocalNow(Range<BigInteger> range, Instant localNow) thr
         TimeSkewResponse timeSkew;
         try
         {
-            TokenRangeMapping<RingInstance> topology = getTokenRangeMapping(true);
-            List<SidecarInstance> instances = topology.getSubRanges(range)
-                                                      .asMapOfRanges()
-                                                      .values()
-                                                      .stream()
-                                                      .flatMap(Collection::stream)
-                                                      .distinct() // remove duplications
-                                                      .map(replica -> new SidecarInstanceImpl(replica.nodeName(), getCassandraContext().sidecarPort()))
-                                                      .collect(Collectors.toList());
-            timeSkew = getCassandraContext().getSidecarClient().timeSkew(instances).get();
+            timeSkew = fetchTimeSkew(range).get();
         }
         catch (InterruptedException | ExecutionException exception)
         {
@@ -270,6 +312,26 @@ void validateTimeSkewWithLocalNow(Range<BigInteger> range, Instant localNow) thr
         }
     }
 
+    /**
+     * Fetches time-skew information from Sidecar. The default implementation queries the replicas
+     * that own {@code range}. Subclasses may override to target a different endpoint set — for
+     * example, coordinated writes route through shared contact points when replica FQDNs are not
+     * reachable from Spark executors.
+     */
+    protected CompletableFuture<TimeSkewResponse> fetchTimeSkew(Range<BigInteger> range)
+    {
+        TokenRangeMapping<RingInstance> topology = getTokenRangeMapping(true);
+        List<SidecarInstance> instances = topology.getSubRanges(range)
+                                                  .asMapOfRanges()
+                                                  .values()
+                                                  .stream()
+                                                  .flatMap(Collection::stream)
+                                                  .distinct() // remove duplications
+                                                  .map(replica -> new SidecarInstanceImpl(replica.nodeName(), getCassandraContext().sidecarPort()))
+                                                  .collect(Collectors.toList());
+        return getCassandraContext().getSidecarClient().timeSkew(instances);
+    }
+
     @Override
     public synchronized void refreshClusterInfo()
     {

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

@@ -0,0 +1,60 @@
+/*
+ * 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.concurrent.CompletableFuture;
+
+import com.google.common.collect.Range;
+
+import o.a.c.sidecar.client.shaded.common.response.TimeSkewResponse;
+
+/**
+ * Variant of {@link CassandraClusterInfo} used when Sidecars are fronted by a load balancer.
+ * Replica FQDNs from the token map are not routable from Spark executors in that topology,
+ * so requests that would otherwise fan out to per-replica Sidecar addresses are routed through
+ * the configured contact points (load balancer endpoints) instead.
+ * <p>
+ * This applies to both single-cluster and coordinated writes. It is selected by
+ * {@link CassandraClusterInfo#create(BulkSparkConf, String)} and
+ * {@link CassandraClusterInfo#create(BroadcastableClusterInfo)} when
+ * {@link org.apache.cassandra.spark.bulkwriter.WriterOptions#SIDECAR_BEHIND_LOAD_BALANCER}
+ * is set.
+ */
+public class LoadBalancedCassandraClusterInfo extends CassandraClusterInfo
+{
+    public LoadBalancedCassandraClusterInfo(BulkSparkConf conf, String clusterId)
+    {
+        super(conf, clusterId);
+    }
+
+    public LoadBalancedCassandraClusterInfo(BroadcastableClusterInfo broadcastable)
+    {
+        super(broadcastable);
+    }
+
+    @Override
+    protected CompletableFuture<TimeSkewResponse> fetchTimeSkew(Range<BigInteger> range)
+    {
+        // range is irrelevant; the load balancer contact points are queried directly rather
+        // than the per-range replicas, whose FQDNs are not routable from Spark executors.
+        return getCassandraContext().getSidecarClient().timeSkew();
+    }
+}

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

@@ -154,4 +154,13 @@ public enum WriterOptions implements WriterOption
      * </ul>
      */
     STORAGE_CREDENTIAL_TYPE,
+    /**
+     * Option declaring that the Sidecar cluster is fronted by a load balancer, so replica FQDNs
+     * from the token map are not directly routable from Spark executors. When {@code true},
+     * requests that would otherwise fan out to per-replica Sidecar addresses are routed through
+     * the configured contact points instead. Defaults to {@code false}.
+     * <p>
+     * Today this affects time-skew validation, and is honored for both single-cluster and coordinated writes.
+     */
+    SIDECAR_BEHIND_LOAD_BALANCER,
 }

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/cloudstorage/coordinated/CassandraClusterInfoGroup.java

@@ -95,7 +95,7 @@ public class CassandraClusterInfoGroup implements ClusterInfo, MultiClusterSuppo
      */
     public static CassandraClusterInfoGroup fromBulkSparkConf(BulkSparkConf conf)
     {
-        return fromBulkSparkConf(conf, clusterId -> new CassandraClusterInfo(conf, clusterId));
+        return fromBulkSparkConf(conf, clusterId -> CassandraClusterInfo.create(conf, clusterId));
     }
 
     /**
@@ -170,7 +170,7 @@ private CassandraClusterInfoGroup(BroadcastableClusterInfoGroup broadcastable)
         // Build list of ClusterInfo from broadcastable data
         List<ClusterInfo> clusterInfosList = new ArrayList<>();
         broadcastable.forEach((clusterId, broadcastableInfo) -> {
-            clusterInfosList.add(new CassandraClusterInfo((BroadcastableClusterInfo) broadcastableInfo));
+            clusterInfosList.add(CassandraClusterInfo.create((BroadcastableClusterInfo) broadcastableInfo));
         });
         this.clusterInfos = Collections.unmodifiableList(clusterInfosList);
 

cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfoTest.java

@@ -24,29 +24,37 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.google.common.util.concurrent.Uninterruptibles;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 
+import o.a.c.sidecar.client.shaded.client.SidecarClient;
 import o.a.c.sidecar.client.shaded.common.response.NodeSettings;
 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.spark.SparkConf;
 
 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;
+import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class CassandraClusterInfoTest
@@ -63,6 +71,75 @@ void testTimeSkewAcceptable()
         .isThrownBy(() -> ci.validateTimeSkewWithLocalNow(range(10, 20), localNow));
     }
 
+    @Test
+    void testLoadBalancedClusterInfoUsesSidecarClientContactPointsForTimeSkew()
+    {
+        Instant localNow = Instant.now();
+        int allowanceMinutes = 10;
+
+        SidecarClient sidecarClient = mock(SidecarClient.class);
+        CassandraContext ctx = mock(CassandraContext.class);
+        when(ctx.getSidecarClient()).thenReturn(sidecarClient);
+        TimeSkewResponse tsr = new TimeSkewResponse(localNow.toEpochMilli(), allowanceMinutes);
+        when(sidecarClient.timeSkew()).thenReturn(CompletableFuture.completedFuture(tsr));
+
+        CassandraClusterInfo ci = new LoadBalancedCassandraClusterInfo((BulkSparkConf) null, null)
+        {
+            @Override
+            protected CassandraContext buildCassandraContext()
+            {
+                return ctx;
+            }
+
+            @Override
+            public TokenRangeMapping<RingInstance> getTokenRangeMapping(boolean cached)
+            {
+                return TokenRangeMappingUtils.buildTokenRangeMapping(0, ImmutableMap.of("dc1", 3), 5);
+            }
+        };
+
+        assertThatNoException()
+        .describedAs("Load-balanced cluster info must use timeSkew() so load balancer contact points stay reachable")
+        .isThrownBy(() -> ci.validateTimeSkewWithLocalNow(range(10, 20), localNow));
+        verify(sidecarClient).timeSkew();
+        verify(sidecarClient, never()).timeSkew(anyList());
+    }
+
+    @Test
+    void testCreateSelectsLoadBalancedClusterInfoWhenSidecarBehindLoadBalancer()
+    {
+        // Single-cluster path: the SIDECAR_BEHIND_LOAD_BALANCER flag must be honored, not just for coordinated writes
+        try (CassandraClusterInfo ci = CassandraClusterInfo.create(bulkSparkConf(true)))
+        {
+            assertThat(ci)
+            .describedAs("Sidecar behind a load balancer must select the load-balanced variant on the single-cluster path")
+            .isExactlyInstanceOf(LoadBalancedCassandraClusterInfo.class);
+        }
+    }
+
+    @Test
+    void testCreateSelectsPlainClusterInfoByDefault()
+    {
+        try (CassandraClusterInfo ci = CassandraClusterInfo.create(bulkSparkConf(false)))
+        {
+            assertThat(ci)
+            .describedAs("Without the load balancer flag, the plain per-replica variant must be selected")
+            .isExactlyInstanceOf(CassandraClusterInfo.class);
+        }
+    }
+
+    private static BulkSparkConf bulkSparkConf(boolean sidecarBehindLoadBalancer)
+    {
+        // No keystore options: this keeps SSL disabled so create() can build a real (non-TLS) Sidecar client
+        // offline. We only assert on the concrete type create() selects, never issuing a request.
+        Map<String, String> options = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER);
+        options.put(WriterOptions.SIDECAR_CONTACT_POINTS.name(), "127.0.0.1");
+        options.put(WriterOptions.KEYSPACE.name(), "ks");
+        options.put(WriterOptions.TABLE.name(), "table");
+        options.put(WriterOptions.SIDECAR_BEHIND_LOAD_BALANCER.name(), String.valueOf(sidecarBehindLoadBalancer));
+        return new BulkSparkConf(new SparkConf(), options);
+    }
+
     @Test
     void testTimeSkewTooLarge()
     {