CASSANALYTICS-180: TokenPartitioner fails to detect range gap in reader

CHANGES.txt

@@ -1,5 +1,6 @@
 0.5.0
 -----
+ * TokenPartitioner fails to detect range gap in reader (CASSANALYTICS-180)
  * 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-common/src/main/java/org/apache/cassandra/spark/data/partitioner/TokenPartitioner.java

@@ -31,6 +31,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import com.google.common.base.Preconditions;
@@ -206,14 +207,15 @@ private void validateRangesDoNotOverlap()
     private void validateCompleteRangeCoverage()
     {
         RangeSet<BigInteger> missingRangeSet = TreeRangeSet.create();
-        missingRangeSet.add(Range.closed(ring.partitioner().minToken(),
-                                         ring.partitioner().maxToken()));
+        // The ring must be open-closed, matching the sub-ranges it is compared against; a closed lower bound would
+        // report minToken as a spurious gap, because open-closed sub-ranges never cover their own lower endpoint
+        missingRangeSet.add(Range.openClosed(ring.partitioner().minToken(),
+                                             ring.partitioner().maxToken()));
 
         partitionMap.asMapOfRanges().keySet().forEach(missingRangeSet::remove);
 
-        List<Range<BigInteger>> missingRanges = missingRangeSet.asRanges().stream()
-                                                                          .filter(Range::isEmpty)
-                                                                          .collect(Collectors.toList());
+        // Whatever is left is a real gap: TreeRangeSet never retains empty ranges
+        Set<Range<BigInteger>> missingRanges = missingRangeSet.asRanges();
         Preconditions.checkState(missingRanges.isEmpty(),
                                  "There should be no missing ranges, but found " + missingRanges.toString());
     }

cassandra-analytics-common/src/test/java/org/apache/cassandra/spark/data/partitioner/TokenPartitionerValidationTest.java

@@ -0,0 +1,111 @@
+/*
+ * 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.data.partitioner;
+
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Range;
+import org.junit.jupiter.api.Test;
+
+import org.apache.cassandra.spark.data.ReplicationFactor;
+import org.apache.cassandra.spark.utils.RangeUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class TokenPartitionerValidationTest
+{
+    private static final Partitioner PARTITIONER = Partitioner.Murmur3Partitioner;
+
+    @Test
+    public void testValidationDetectsRangeGap()
+    {
+        List<Range<BigInteger>> subRanges = RangeUtils.split(wholeRing(), 4);
+
+        assertThatThrownBy(() -> new TokenPartitioner(withGapAt(subRanges, 2), ring()))
+        .isInstanceOf(IllegalStateException.class)
+        .hasMessageContaining("There should be no missing ranges")
+        .hasMessageContaining(gapPunchedInto(subRanges.get(2)).toString());
+    }
+
+    @Test
+    public void testValidationDetectsRangeGapAtRingLowerEdge()
+    {
+        // Guards the bound type at minToken from both sides: minToken itself is owned by no sub-range and must not
+        // be reported, yet a gap starting immediately above it must still be caught. The gap is punched into the
+        // first sub-range rather than dropping it, so that the partition count stays put and validateMapSizes
+        // cannot fail first with an unrelated message.
+        List<Range<BigInteger>> subRanges = RangeUtils.split(wholeRing(), 4);
+
+        assertThatThrownBy(() -> new TokenPartitioner(withGapAt(subRanges, 0), ring()))
+        .isInstanceOf(IllegalStateException.class)
+        .hasMessageContaining("There should be no missing ranges")
+        .hasMessageContaining(gapPunchedInto(subRanges.get(0)).toString());
+    }
+
+    @Test
+    public void testValidationAcceptsCompleteRangeCoverage()
+    {
+        // minToken is deliberately left uncovered: the sub-ranges are open-closed, so it belongs to none of them.
+        // Validation must not report it as a gap, otherwise every job fails on a healthy ring.
+        TokenPartitioner partitioner = new TokenPartitioner(RangeUtils.split(wholeRing(), 4), ring());
+        assertThat(partitioner.numPartitions()).isEqualTo(4);
+    }
+
+    private static Range<BigInteger> wholeRing()
+    {
+        return Range.openClosed(PARTITIONER.minToken(), PARTITIONER.maxToken());
+    }
+
+    /**
+     * Punches a real, non-empty gap into the sub-range at {@code gapIndex} by moving its lower endpoint up, so that
+     * the returned ranges leave exactly {@link #gapPunchedInto} uncovered.
+     */
+    private static List<Range<BigInteger>> withGapAt(List<Range<BigInteger>> gapFreeRanges, int gapIndex)
+    {
+        List<Range<BigInteger>> ranges = new ArrayList<>(gapFreeRanges);
+        Range<BigInteger> covered = ranges.get(gapIndex);
+        ranges.set(gapIndex, Range.openClosed(gapPunchedInto(covered).upperEndpoint(), covered.upperEndpoint()));
+        return ranges;
+    }
+
+    /**
+     * @return the sub-range that {@link #withGapAt} leaves uncovered when it punches a gap into {@code range}
+     */
+    private static Range<BigInteger> gapPunchedInto(Range<BigInteger> range)
+    {
+        return Range.openClosed(range.lowerEndpoint(), range.lowerEndpoint().add(BigInteger.TEN));
+    }
+
+    private static CassandraRing ring()
+    {
+        List<CassandraInstance> instances = Arrays.asList(new CassandraInstance("0", "local0-i1", "DEV"),
+                                                          new CassandraInstance("100", "local0-i2", "DEV"),
+                                                          new CassandraInstance("200", "local0-i3", "DEV"));
+        return new CassandraRing(PARTITIONER,
+                                 "test",
+                                 new ReplicationFactor(ImmutableMap.of("class", "NetworkTopologyStrategy", "DEV", "3")),
+                                 instances);
+    }
+}

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

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
@@ -230,15 +231,15 @@ private void validateRangesDoNotOverlap()
     private void validateCompleteRangeCoverage()
     {
         RangeSet<BigInteger> missingRangeSet = TreeRangeSet.create();
-        missingRangeSet.add(Range.closed(tokenRangeMapping.partitioner().minToken(),
-                                         tokenRangeMapping.partitioner().maxToken()));
+        // The ring must be open-closed, matching the sub-ranges it is compared against; a closed lower bound would
+        // report minToken as a spurious gap, because open-closed sub-ranges never cover their own lower endpoint
+        missingRangeSet.add(Range.openClosed(tokenRangeMapping.partitioner().minToken(),
+                                             tokenRangeMapping.partitioner().maxToken()));
 
         partitionMap.asMapOfRanges().keySet().forEach(missingRangeSet::remove);
 
-        List<Range<BigInteger>> missingRanges = missingRangeSet.asRanges().stream()
-                                                               .filter(Range::isEmpty)
-                                                               .collect(Collectors.toList());
-        // noinspection unchecked
+        // Whatever is left is a real gap: TreeRangeSet never retains empty ranges
+        Set<Range<BigInteger>> missingRanges = missingRangeSet.asRanges();
         Preconditions.checkState(missingRanges.isEmpty(),
                                  "There should be no missing ranges, but found " + missingRanges.toString());
     }

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

@@ -21,17 +21,31 @@
 
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 
+import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeMap;
+import com.google.common.collect.TreeRangeMap;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.apache.cassandra.spark.bulkwriter.token.TokenRangeMapping;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+import org.apache.cassandra.spark.utils.RangeUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class TokenPartitionerTest
 {
+    private static final Partitioner RING_PARTITIONER = Partitioner.Murmur3Partitioner;
+
     private TokenPartitioner partitioner;
 
     @BeforeEach
@@ -172,6 +186,92 @@ public void testSplitCalculationWithMultipleDcs()
         assertThat(partitioner.numPartitions()).isGreaterThanOrEqualTo(200);
     }
 
+    // Range coverage validation must reject a partition map that leaves a token uncovered.
+    @Test
+    public void testValidationDetectsRangeGap()
+    {
+        List<Range<BigInteger>> subRanges = RangeUtils.split(wholeRing(), 4);
+
+        // numberSplits of 1 leaves the ranges untouched, so they reach the partition map as-is
+        assertThatThrownBy(() -> new TokenPartitioner(mappingCovering(withGapAt(subRanges, 2)), 1, 2, 1, false))
+        .isInstanceOf(IllegalStateException.class)
+        .hasMessageContaining("There should be no missing ranges")
+        .hasMessageContaining(gapPunchedInto(subRanges.get(2)).toString());
+    }
+
+    @Test
+    public void testValidationDetectsRangeGapAtRingLowerEdge()
+    {
+        // Guards the bound type at minToken from both sides: minToken itself is owned by no sub-range and must not
+        // be reported, yet a gap starting immediately above it must still be caught. The gap is punched into the
+        // first sub-range rather than dropping it, so that the partition count stays put and validateMapSizes
+        // cannot fail first with an unrelated message.
+        List<Range<BigInteger>> subRanges = RangeUtils.split(wholeRing(), 4);
+
+        assertThatThrownBy(() -> new TokenPartitioner(mappingCovering(withGapAt(subRanges, 0)), 1, 2, 1, false))
+        .isInstanceOf(IllegalStateException.class)
+        .hasMessageContaining("There should be no missing ranges")
+        .hasMessageContaining(gapPunchedInto(subRanges.get(0)).toString());
+    }
+
+    // Guards against over-correcting the fix for the above: the partition map is built from open-closed sub-ranges,
+    // so minToken belongs to none of them. Validation that expected [minToken, maxToken] to be covered would report
+    // a spurious [minToken, minToken] gap and fail every bulk write on a perfectly healthy ring.
+    @Test
+    public void testValidationAcceptsRingNotCoveringMinToken()
+    {
+        TokenRangeMapping<RingInstance> tokenRangeMapping = TokenRangeMappingUtils.buildTokenRangeMapping(0, ImmutableMap.of("DC1", 3), 3);
+        // Validation runs in the driver as part of construction, so not throwing here is the assertion
+        TokenPartitioner tokenPartitioner = new TokenPartitioner(tokenRangeMapping, 2, 2, 1, false);
+        // ... and the premise of the test holds: no partition owns minToken, as the sub-ranges are open-closed
+        assertThat(tokenPartitioner.getTokenRange(0).contains(RING_PARTITIONER.minToken())).isFalse();
+    }
+
+    private static Range<BigInteger> wholeRing()
+    {
+        return Range.openClosed(RING_PARTITIONER.minToken(), RING_PARTITIONER.maxToken());
+    }
+
+    /**
+     * Mocking is the only way to feed a gapped range map to the partitioner: {@link TokenRangeMapping} seeds its
+     * range map with the whole ring, so a mapping built the normal way is always gap-free and cannot exercise the
+     * coverage check.
+     *
+     * @return a mapping whose range map covers exactly {@code ranges}
+     */
+    private static TokenRangeMapping<RingInstance> mappingCovering(List<Range<BigInteger>> ranges)
+    {
+        RangeMap<BigInteger, List<RingInstance>> rangeMap = TreeRangeMap.create();
+        ranges.forEach(range -> rangeMap.put(range, Collections.emptyList()));
+
+        @SuppressWarnings("unchecked")
+        TokenRangeMapping<RingInstance> tokenRangeMapping = mock(TokenRangeMapping.class);
+        when(tokenRangeMapping.partitioner()).thenReturn(RING_PARTITIONER);
+        when(tokenRangeMapping.getRangeMap()).thenReturn(rangeMap);
+        when(tokenRangeMapping.getTokenRanges()).thenReturn(ArrayListMultimap.create());
+        return tokenRangeMapping;
+    }
+
+    /**
+     * Punches a real, non-empty gap into the sub-range at {@code gapIndex} by moving its lower endpoint up, so that
+     * the returned ranges leave exactly {@link #gapPunchedInto} uncovered.
+     */
+    private static List<Range<BigInteger>> withGapAt(List<Range<BigInteger>> gapFreeRanges, int gapIndex)
+    {
+        List<Range<BigInteger>> ranges = new ArrayList<>(gapFreeRanges);
+        Range<BigInteger> covered = ranges.get(gapIndex);
+        ranges.set(gapIndex, Range.openClosed(gapPunchedInto(covered).upperEndpoint(), covered.upperEndpoint()));
+        return ranges;
+    }
+
+    /**
+     * @return the sub-range that {@link #withGapAt} leaves uncovered when it punches a gap into {@code range}
+     */
+    private static Range<BigInteger> gapPunchedInto(Range<BigInteger> range)
+    {
+        return Range.openClosed(range.lowerEndpoint(), range.lowerEndpoint().add(BigInteger.TEN));
+    }
+
     private int partitionForToken(int token)
     {
         return partitionForToken(BigInteger.valueOf(token));