fix: Improvements to encoding ranges (int vs long)

parallel-consumer-core/src/main/java/io/confluent/csid/utils/MathUtils.java

@@ -0,0 +1,29 @@
+package io.confluent.csid.utils;
+
+/*-
+ * Copyright (C) 2020-2022 Confluent, Inc.
+ */
+
+import lombok.experimental.UtilityClass;
+
+/**
+ * @author Antony Stubbs
+ */
+@UtilityClass
+public class MathUtils {
+
+    /**
+     * Ensures exact conversion from a Long to a Short.
+     * <p>
+     * {@link Math} doesn't have an exact conversion from Long to Short.
+     *
+     * @see Math#toIntExact
+     */
+    public static short toShortExact(long value) {
+        final short shortCast = (short) value;
+        if (shortCast != value) {
+            throw new ArithmeticException("short overflow");
+        }
+        return shortCast;
+    }
+}

parallel-consumer-core/src/main/java/io/confluent/csid/utils/Range.java

@@ -1,48 +1,87 @@
 package io.confluent.csid.utils;
 
 /*-
- * Copyright (C) 2020-2021 Confluent, Inc.
+ * Copyright (C) 2020-2022 Confluent, Inc.
  */
 
-import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.NoSuchElementException;
-import java.util.function.Consumer;
 import java.util.stream.IntStream;
+import java.util.stream.LongStream;
+
+import static java.util.stream.Collectors.toList;
 
 /**
- * https://stackoverflow.com/a/16570509/105741
+ * Class for simple ranges.
+ * <p>
+ * <a href="https://stackoverflow.com/a/16570509/105741">For loop - like Python range function</a>
+ *
+ * @see #range(long)
  */
-public class Range implements Iterable<Integer> {
+public class Range implements Iterable<Long> {
+
+    private final long start;
 
     private final long limit;
 
+    /**
+     * @see this#range(long)
+     */
+    public Range(int start, long max) {
+        this.start = start;
+        this.limit = max;
+    }
+
     public Range(long limit) {
+        this.start = 0L;
         this.limit = limit;
     }
 
     /**
-     * Exclusive of max
+     * Provides an {@link Iterable} for the range of numbers from 0 to the given limit.
+     * <p>
+     * Exclusive of max.
+     * <p>
+     * Consider using {@link IntStream#range(int, int)#forEachOrdered} instead:
+     * <pre>
+     * IntStream.range(0, originalBitsetSize).forEachOrdered(offset -> {
+     * </pre>
+     * However, if you don't want o use a closure, this is a good alternative.
      */
     public static Range range(long max) {
         return new Range(max);
     }
 
+    /**
+     * @see #range(long)
+     */
+    public static Range range(int start, long max) {
+        return new Range(start, max);
+    }
+
+    /**
+     * Potentially slow, but useful for tests
+     */
+    public static List<Integer> listOfIntegers(int max) {
+        return Range.range(max).listAsIntegers();
+    }
+
+
     @Override
-    public Iterator<Integer> iterator() {
+    public Iterator<Long> iterator() {
         final long max = limit;
-        return new Iterator<Integer>() {
+        return new Iterator<>() {
 
-            private int current = 0;
+            private long current = start;
 
             @Override
             public boolean hasNext() {
                 return current < max;
             }
 
             @Override
-            public Integer next() {
+            public Long next() {
                 if (hasNext()) {
                     return current++;
                 } else {
@@ -57,23 +96,14 @@ public void remove() {
         };
     }
 
-    public List<Integer> list() {
-        ArrayList<Integer> integers = new ArrayList<Integer>();
-        forEach(integers::add);
-        return integers;
-    }
-
-    public IntStream toStream() {
-        return IntStream.range(0, (int) limit);
-    }
-
-    static IntStream rangeStream(int i) {
-        return IntStream.range(0, i);
+    public List<Integer> listAsIntegers() {
+        return IntStream.range(Math.toIntExact(start), Math.toIntExact(limit))
+                .boxed()
+                .collect(toList());
     }
 
-    static void range(int max, Consumer<Integer> consumer) {
-        IntStream.range(0, max)
-                .forEach(consumer::accept);
+    public LongStream toStream() {
+        return LongStream.range(start, limit);
     }
 
 }

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/offsets/BitSetEncoder.java

@@ -4,6 +4,7 @@
  * Copyright (C) 2020-2022 Confluent, Inc.
  */
 
+import io.confluent.csid.utils.MathUtils;
 import io.confluent.csid.utils.StringUtils;
 import io.confluent.parallelconsumer.internal.InternalRuntimeException;
 import io.confluent.parallelconsumer.state.PartitionState;
@@ -46,33 +47,37 @@ public class BitSetEncoder extends OffsetEncoder {
 
     private static final Version DEFAULT_VERSION = Version.v2;
 
+    /**
+     * {@link BitSet} only supports {@link Integer#MAX_VALUE) bits
+     */
     public static final Integer MAX_LENGTH_ENCODABLE = Integer.MAX_VALUE;
 
     @Getter
     private final BitSet bitSet;
 
-    private final int originalLength;
+    private final long originalLength;
 
     private Optional<byte[]> encodedBytes = Optional.empty();
 
-    public BitSetEncoder(int length, OffsetSimultaneousEncoder offsetSimultaneousEncoder) throws BitSetEncodingNotSupportedException {
-        this(length, offsetSimultaneousEncoder, DEFAULT_VERSION);
-    }
-
     /**
      * @param length the difference between the highest and lowest offset to be encoded
      */
-    public BitSetEncoder(int length, OffsetSimultaneousEncoder offsetSimultaneousEncoder, Version newVersion) throws BitSetEncodingNotSupportedException {
+    public BitSetEncoder(long length, OffsetSimultaneousEncoder offsetSimultaneousEncoder, Version newVersion) throws BitSetEncodingNotSupportedException {
         super(offsetSimultaneousEncoder);
 
         this.version = newVersion;
 
-        bitSet = new BitSet(length);
+        // prep bit set buffer, range check above
+        try {
+            bitSet = new BitSet(Math.toIntExact(length));
+        } catch (ArithmeticException e) {
+            throw new BitSetEncodingNotSupportedException("BitSet only supports " + MAX_LENGTH_ENCODABLE + " bits, but " + length + " were requested", e);
+        }
 
         this.originalLength = length;
     }
 
-    private ByteBuffer constructWrappedByteBuffer(final int length, final Version newVersion) throws BitSetEncodingNotSupportedException {
+    private ByteBuffer constructWrappedByteBuffer(long length, Version newVersion) throws BitSetEncodingNotSupportedException {
         return switch (newVersion) {
             case v1 -> initV1(length);
             case v2 -> initV2(length);
@@ -85,43 +90,39 @@ private ByteBuffer constructWrappedByteBuffer(final int length, final Version ne
      * Integer.MAX_VALUE should always be good enough as system restricts large from being processed at once.
      */
     // TODO refactor inivtV2 and V1 together, passing in the Short or Integer
-    private ByteBuffer initV2(int bitsetEntriesRequired) throws BitSetEncodingNotSupportedException {
+    private ByteBuffer initV2(long bitsetEntriesRequired) throws BitSetEncodingNotSupportedException {
         if (bitsetEntriesRequired > MAX_LENGTH_ENCODABLE) {
             // need to upgrade to using Integer for the bitset length, but can't change serialisation format in-place
             throw new BitSetEncodingNotSupportedException(StringUtils.msg("BitSet V2 too long to encode, as length overflows Integer.MAX_VALUE. Length: {}. (max: {})", bitsetEntriesRequired, MAX_LENGTH_ENCODABLE));
         }
 
-        // prep bit set buffer
         int bytesRequiredForEntries = (int) (Math.ceil((double) bitsetEntriesRequired / Byte.SIZE));
         int lengthEntryWidth = Integer.BYTES;
         int wrappedBufferLength = lengthEntryWidth + bytesRequiredForEntries + 1;
         final ByteBuffer wrappedBitSetBytesBuffer = ByteBuffer.allocate(wrappedBufferLength);
 
         // bitset doesn't serialise it's set capacity, so we have to as the unused capacity actually means something
-        wrappedBitSetBytesBuffer.putInt(bitsetEntriesRequired);
+        wrappedBitSetBytesBuffer.putInt(Math.toIntExact(bitsetEntriesRequired));
 
         return wrappedBitSetBytesBuffer;
     }
 
     /**
      * This was a bit "short" sighted of me....
-     *
-     * @return
      */
-    private ByteBuffer initV1(int bitsetEntriesRequired) throws BitSetEncodingNotSupportedException {
+    private ByteBuffer initV1(long bitsetEntriesRequired) throws BitSetEncodingNotSupportedException {
         if (bitsetEntriesRequired > Short.MAX_VALUE) {
             // need to upgrade to using Integer for the bitset length, but can't change serialisation format in-place
             throw new BitSetEncodingNotSupportedException("Input too long to encode for BitSet V1, length overflows Short.MAX_VALUE: " + bitsetEntriesRequired + ". (max: " + Short.MAX_VALUE + ")");
         }
 
-        // prep bit set buffer
         int bytesRequiredForEntries = (int) (Math.ceil((double) bitsetEntriesRequired / Byte.SIZE));
         int lengthEntryWidth = Short.BYTES;
         int wrappedBufferLength = lengthEntryWidth + bytesRequiredForEntries + 1;
         final ByteBuffer wrappedBitSetBytesBuffer = ByteBuffer.allocate(wrappedBufferLength);
 
         // bitset doesn't serialise it's set capacity, so we have to as the unused capacity actually means something
-        wrappedBitSetBytesBuffer.putShort((short) bitsetEntriesRequired);
+        wrappedBitSetBytesBuffer.putShort(MathUtils.toShortExact(bitsetEntriesRequired));
 
         return wrappedBitSetBytesBuffer;
     }
@@ -143,13 +144,14 @@ protected OffsetEncoding getEncodingTypeCompressed() {
     }
 
     @Override
-    public void encodeIncompleteOffset(final int index) {
+    public void encodeIncompleteOffset(final long relativeOffset) {
         // noop - bitset defaults to 0's (`unset`)
     }
 
     @Override
-    public void encodeCompletedOffset(final int index) {
-        bitSet.set(index);
+    public void encodeCompletedOffset(final long relativeOffset) {
+        // range will already have been checked at initialization
+        bitSet.set(Math.toIntExact(relativeOffset));
     }
 
     @Override

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/offsets/ByteBufferEncoder.java

@@ -19,9 +19,11 @@ public class ByteBufferEncoder extends OffsetEncoder {
 
     private final ByteBuffer bytesBuffer;
 
-    public ByteBufferEncoder(final int length, OffsetSimultaneousEncoder offsetSimultaneousEncoder) {
+    public ByteBufferEncoder(long length, OffsetSimultaneousEncoder offsetSimultaneousEncoder) {
         super(offsetSimultaneousEncoder);
-        this.bytesBuffer = ByteBuffer.allocate(1 + length);
+        // safe cast the length to an int, as we're not expecting to have more than 2^31 offsets
+        final int safeCast = Math.toIntExact(length);
+        this.bytesBuffer = ByteBuffer.allocate(1 + safeCast);
     }
 
     @Override
@@ -35,12 +37,12 @@ protected OffsetEncoding getEncodingTypeCompressed() {
     }
 
     @Override
-    public void encodeIncompleteOffset(final int rangeIndex) {
+    public void encodeIncompleteOffset(final long relativeOffset) {
         this.bytesBuffer.put((byte) 0);
     }
 
     @Override
-    public void encodeCompletedOffset(final int rangeIndex) {
+    public void encodeCompletedOffset(final long relativeOffset) {
         this.bytesBuffer.put((byte) 1);
     }
 

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/offsets/OffsetBitSet.java

@@ -15,9 +15,9 @@
 import static io.confluent.csid.utils.Range.range;
 
 /**
- * Deserialisation tools for {@link BitSetEncoder}.
+ * Deserialization tools for {@link BitSetEncoder}.
  * <p>
- * todo unify or refactor with {@link BitSetEncoder}. Why was it ever seperate?
+ * todo unify or refactor with {@link BitSetEncoder}. Why was it ever separate?
  *
  * @author Antony Stubbs
  * @see BitSetEncoder
@@ -41,8 +41,9 @@ static String deserialiseBitSet(int originalBitsetSize, ByteBuffer s) {
         BitSet bitSet = BitSet.valueOf(s);
 
         StringBuilder result = new StringBuilder(bitSet.size());
-        for (var offset : range(originalBitsetSize)) {
-            if (bitSet.get(offset)) {
+        for (Long offset : range(originalBitsetSize)) {
+            // range will already have been checked at initialization
+            if (bitSet.get(Math.toIntExact(offset))) {
                 result.append('x');
             } else {
                 result.append('o');
@@ -69,7 +70,9 @@ static Set<Long> deserialiseBitSetToIncompletes(long baseOffset, int originalBit
         BitSet bitSet = BitSet.valueOf(inputBuffer);
         int numberOfIncompletes = originalBitsetSize - bitSet.cardinality();
         var incompletes = new HashSet<Long>(numberOfIncompletes);
-        for (var relativeOffset : range(originalBitsetSize)) {
+        for (long relativeOffsetLong : range(originalBitsetSize)) {
+            // range will already have been checked at initialization
+            var relativeOffset = Math.toIntExact(relativeOffsetLong);
             long offset = baseOffset + relativeOffset;
             if (bitSet.get(relativeOffset)) {
                 log.trace("Ignoring completed offset {}", relativeOffset);

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/offsets/OffsetEncoder.java

@@ -28,16 +28,16 @@ protected OffsetEncoder(OffsetSimultaneousEncoder offsetSimultaneousEncoder) {
 
     protected abstract OffsetEncoding getEncodingTypeCompressed();
 
-    abstract void encodeIncompleteOffset(final int rangeIndex);
+    abstract void encodeIncompleteOffset(final long relativeOffset) throws EncodingNotSupportedException;
 
-    abstract void encodeCompletedOffset(final int rangeIndex);
+    abstract void encodeCompletedOffset(final long relativeOffset) throws EncodingNotSupportedException;
 
     abstract byte[] serialise() throws EncodingNotSupportedException;
 
     abstract int getEncodedSize();
 
     boolean quiteSmall() {
-        return this.getEncodedSize() < OffsetSimultaneousEncoder.LARGE_INPUT_MAP_SIZE_THRESHOLD;
+        return this.getEncodedSize() < OffsetSimultaneousEncoder.LARGE_ENCODED_SIZE_THRESHOLD_BYTES;
     }
 
     byte[] compress() throws IOException {