KAFKA-20597: Fix prefixScan overflow for prefixes with no upper bound in in-memory state stores

clients/src/main/java/org/apache/kafka/common/utils/internals/ByteUtils.java

@@ -66,6 +66,23 @@ public static Bytes increment(Bytes input) throws IndexOutOfBoundsException {
         }
     }
 
+    /**
+     * Same as {@link #increment(Bytes)} but returns {@code null} instead of throwing
+     * {@code IndexOutOfBoundsException} when incrementing would overflow the byte array
+     * (i.e. every byte is {@code 0xFF}). A {@code null} result represents an unbounded
+     * upper range, which callers performing prefix scans should treat as "scan to the end".
+     *
+     * @param input the byte array to increment
+     * @return A new copy of the incremented byte array, or {@code null} if incrementing would overflow
+     */
+    public static Bytes incrementWithoutOverflow(final Bytes input) {
+        try {
+            return increment(input);
+        } catch (final IndexOutOfBoundsException e) {
+            return null;
+        }
+    }
+
     /**
      * A byte array comparator based on lexicographic ordering.
      */

clients/src/test/java/org/apache/kafka/common/utils/internals/ByteUtilsTest.java

@@ -41,6 +41,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -77,6 +78,21 @@ public void testIncrementUpperBoundary() {
         byte[] input = new byte[]{(byte) 0xFF, (byte) 0xFF, (byte) 0xFF};
         assertThrows(IndexOutOfBoundsException.class, () -> ByteUtils.increment(Bytes.wrap(input)));
     }
+
+    @Test
+    public void testIncrementWithoutOverflow() {
+        byte[] input = new byte[]{(byte) 0xAB, (byte) 0xCD, (byte) 0xFF};
+        byte[] expected = new byte[]{(byte) 0xAB, (byte) 0xCE, (byte) 0x00};
+        Bytes output = ByteUtils.incrementWithoutOverflow(Bytes.wrap(input));
+        assertArrayEquals(expected, output.get());
+    }
+
+    @Test
+    public void testIncrementWithoutOverflowReturnsNullOnOverflow() {
+        byte[] input = new byte[]{(byte) 0xFF, (byte) 0xFF, (byte) 0xFF};
+        assertNull(ByteUtils.incrementWithoutOverflow(Bytes.wrap(input)));
+    }
+
     @Test
     public void testIncrementWithSubmap() {
         final NavigableMap<Bytes, byte[]> map = new TreeMap<>();

streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java

@@ -439,7 +439,7 @@ private <PS extends Serializer<P>, P> KeyValueIterator<Bytes, byte[]> prefixScan
         validateStoreOpen();
         final KeyValueIterator<Bytes, byte[]> storeIterator = underlying.prefixScan(prefix, prefixKeySerializer);
         final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null, prefix));
-        final Bytes to = ByteUtils.increment(from);
+        final Bytes to = ByteUtils.incrementWithoutOverflow(from);
         final ThreadCache.MemoryLRUCacheBytesIterator cacheIterator = internalContext.cache().range(cacheName, from, to, false);
         return new MergedSortedCacheKeyValueBytesStoreIterator(cacheIterator, storeIterator, true);
     }

streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java

@@ -35,7 +35,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 
-import static org.apache.kafka.streams.state.internals.RocksDBStore.incrementWithoutOverflow;
+import static org.apache.kafka.common.utils.internals.ByteUtils.incrementWithoutOverflow;
 
 /**
  * A generic implementation of {@link RocksDBStore.ColumnFamilyAccessor} that supports dual-column-family

streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java

@@ -239,13 +239,16 @@ public <PS extends Serializer<P>, P> KeyValueIterator<Bytes, byte[]> prefixScan(
     private <PS extends Serializer<P>, P> KeyValueIterator<Bytes, byte[]> prefixScan(final P prefix, final PS prefixKeySerializer,
                                                                                      final IsolationLevel isolationLevel) {
         final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null, prefix));
-        final Bytes to = ByteUtils.increment(from);
+        final Bytes to = ByteUtils.incrementWithoutOverflow(from);
 
         if (transactionBuffer != null) {
             return transactionBuffer.range(from, to, true, false, isolationLevel);
         }
         synchronized (this) {
-            return new InMemoryKeyValueIterator(map.subMap(from, true, to, false).keySet(), true);
+            final NavigableMap<Bytes, byte[]> subMap = to == null
+                ? map.tailMap(from, true)
+                : map.subMap(from, true, to, false);
+            return new InMemoryKeyValueIterator(subMap.keySet(), true);
         }
     }
 

streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegment.java

@@ -45,7 +45,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-import static org.apache.kafka.streams.state.internals.RocksDBStore.incrementWithoutOverflow;
+import static org.apache.kafka.common.utils.internals.ByteUtils.incrementWithoutOverflow;
 
 /**
  * This "logical segment" is a segment which shares its underlying physical store with other

streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryNavigableLRUCache.java

@@ -31,6 +31,7 @@
 
 import java.util.Iterator;
 import java.util.Map;
+import java.util.NavigableMap;
 import java.util.Objects;
 import java.util.TreeMap;
 
@@ -90,13 +91,16 @@ private Iterator<Bytes> getIterator(final TreeMap<Bytes, byte[]> treeMap, final
     public <PS extends Serializer<P>, P> KeyValueIterator<Bytes, byte[]> prefixScan(final P prefix, final PS prefixKeySerializer) {
 
         final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null, prefix));
-        final Bytes to = ByteUtils.increment(from);
+        final Bytes to = ByteUtils.incrementWithoutOverflow(from);
 
         final TreeMap<Bytes, byte[]> treeMap = toTreeMap();
+        final NavigableMap<Bytes, byte[]> subMap = to == null
+            ? treeMap.tailMap(from, true)
+            : treeMap.subMap(from, true, to, false);
 
         return new DelegatingPeekingKeyValueIterator<>(
             name(),
-            new MemoryNavigableLRUCache.CacheIterator(treeMap.subMap(from, true, to, false).keySet().iterator(), treeMap)
+            new MemoryNavigableLRUCache.CacheIterator(subMap.keySet().iterator(), treeMap)
         );
     }
 

streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java

@@ -1478,7 +1478,7 @@ public ManagedKeyValueIterator<Bytes, byte[]> all(final DBAccessor accessor, fin
 
         @Override
         public ManagedKeyValueIterator<Bytes, byte[]> prefixScan(final DBAccessor accessor, final Bytes prefix) {
-            final Bytes to = incrementWithoutOverflow(prefix);
+            final Bytes to = ByteUtils.incrementWithoutOverflow(prefix);
             return accessor.prefixScan(columnFamily, name, prefix, to);
         }
 
@@ -1551,20 +1551,4 @@ public Position getPosition() {
             return position.copy().merge(dbAccessor.uncommittedPositionDeltas());
         }
     }
-
-    /**
-     * Same as {@link ByteUtils#increment(Bytes)} but {@code null} is returned instead of throwing
-     * {@code IndexOutOfBoundsException} in the event of overflow.
-     *
-     * @param input bytes to increment
-     * @return A new copy of the incremented byte array, or {@code null} if incrementing would
-     *         result in overflow.
-     */
-    static Bytes incrementWithoutOverflow(final Bytes input) {
-        try {
-            return ByteUtils.increment(input);
-        } catch (final IndexOutOfBoundsException e) {
-            return null;
-        }
-    }
-}
+}

streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java

@@ -660,5 +660,27 @@ public void prefixScanShouldNotThrowConcurrentModificationException() {
                 iter.next();
             }
         }
-    }                  
+    }
+
+    @Test
+    public void prefixScanShouldReturnMatchesWhenPrefixHasNoUpperBound() {
+        // A prefix that serializes to all 0xFF bytes has no lexicographically larger successor,
+        // so incrementing it would overflow. The scan must treat this as an unbounded upper range
+        // and run to the end of the keyspace rather than throwing IndexOutOfBoundsException.
+        // IntegerSerializer encodes -1 as 0xFFFFFFFF, which reproduces the overflow case (KAFKA-20597).
+        store.put(0, "zero");
+        store.put(1, "one");
+        store.put(2, "two");
+        store.put(-1, "minus-one");
+
+        final List<KeyValue<Integer, String>> result = new ArrayList<>();
+        try (final KeyValueIterator<Integer, String> iter = store.prefixScan(-1, new IntegerSerializer())) {
+            while (iter.hasNext()) {
+                result.add(iter.next());
+            }
+        }
+
+        // Only the key whose bytes are 0xFFFFFFFF matches the prefix; the lower-valued keys are excluded.
+        assertEquals(Collections.singletonList(KeyValue.pair(-1, "minus-one")), result);
+    }
 }