diff --git a/clients/src/main/java/org/apache/kafka/common/utils/internals/ByteUtils.java b/clients/src/main/java/org/apache/kafka/common/utils/internals/ByteUtils.java index 9a20b8cb91fa2..ad58d661a442f 100644 --- a/clients/src/main/java/org/apache/kafka/common/utils/internals/ByteUtils.java +++ b/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. */ diff --git a/clients/src/test/java/org/apache/kafka/common/utils/internals/ByteUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/internals/ByteUtilsTest.java index 139e3efdd4b53..051ededbb040c 100644 --- a/clients/src/test/java/org/apache/kafka/common/utils/internals/ByteUtilsTest.java +++ b/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 map = new TreeMap<>(); diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java index b6d855900911a..e342849f14667 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java @@ -439,7 +439,7 @@ private , P> KeyValueIterator prefixScan validateStoreOpen(); final KeyValueIterator 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); } diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java index c41aaf4a357ff..55b1de554d7d1 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java +++ b/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 diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java index fcf570a101a49..c676c1ab39d7b 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java @@ -239,13 +239,16 @@ public , P> KeyValueIterator prefixScan( private , P> KeyValueIterator 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 subMap = to == null + ? map.tailMap(from, true) + : map.subMap(from, true, to, false); + return new InMemoryKeyValueIterator(subMap.keySet(), true); } } diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegment.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegment.java index 2aa5182fc2dec..0aaea383295c2 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegment.java +++ b/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 diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryNavigableLRUCache.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryNavigableLRUCache.java index 0406f217aecae..b57846fc06811 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryNavigableLRUCache.java +++ b/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 getIterator(final TreeMap treeMap, final public , P> KeyValueIterator 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 treeMap = toTreeMap(); + final NavigableMap 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) ); } diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java index e7b2eef3a4cdb..dfb4775af1f2e 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java @@ -1478,7 +1478,7 @@ public ManagedKeyValueIterator all(final DBAccessor accessor, fin @Override public ManagedKeyValueIterator 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; - } - } -} \ No newline at end of file +} diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java index 3727add492d13..ab0e99235d8ab 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java +++ b/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> result = new ArrayList<>(); + try (final KeyValueIterator 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); + } }