This is an automated email from the ASF dual-hosted git repository.
mjsax pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new 34e04e35e05 KAFKA-20412: Fix prefixScan for KV-Store with headers
(#21971)
34e04e35e05 is described below
commit 34e04e35e057b22ec6b43d3ffbc73c037b7ce093
Author: Uladzislau Blok <[email protected]>
AuthorDate: Sun Jul 26 03:24:25 2026 +0200
KAFKA-20412: Fix prefixScan for KV-Store with headers (#21971)
Part of KIP-1271.
This PR closes a gap in headers passing. For `prefixScan`, we are not
yet correctly passing
Headers into the serdes.
Approach is similar to
`MeteredTimestampedKeyValueStoreWithHeaders#delete(key)`
Reviewers: Alieh Saeedi <[email protected]>, Matthias J. Sax
<[email protected]>
---
.../state/internals/CachingKeyValueStore.java | 1 +
.../state/internals/InMemoryKeyValueStore.java | 1 +
.../state/internals/MemoryNavigableLRUCache.java | 2 +-
.../state/internals/MeteredKeyValueStore.java | 6 ++-
...MeteredTimestampedKeyValueStoreWithHeaders.java | 6 ++-
.../state/internals/MeteredKeyValueStoreTest.java | 41 ++++++++++----
...redTimestampedKeyValueStoreWithHeadersTest.java | 62 ++++++++++++++++++++++
7 files changed, 105 insertions(+), 14 deletions(-)
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 b6d85590091..1a1edf71692 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
@@ -438,6 +438,7 @@ public class CachingKeyValueStore
final ReadOnlyKeyValueStore<Bytes, byte[]> underlying) {
validateStoreOpen();
final KeyValueIterator<Bytes, byte[]> storeIterator =
underlying.prefixScan(prefix, prefixKeySerializer);
+ // headers aren't needed because the prefix already arrives serialized
final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null,
prefix));
final Bytes to = ByteUtils.increment(from);
final ThreadCache.MemoryLRUCacheBytesIterator cacheIterator =
internalContext.cache().range(cacheName, from, to, false);
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 fcf570a101a..793d5b4cf2a 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
@@ -238,6 +238,7 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
private <PS extends Serializer<P>, P> KeyValueIterator<Bytes, byte[]>
prefixScan(final P prefix, final PS prefixKeySerializer,
final IsolationLevel isolationLevel) {
+ // headers aren't needed because the prefix already arrives serialized
final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null,
prefix));
final Bytes to = ByteUtils.increment(from);
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 0406f217aec..83506bc1e14 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
@@ -88,7 +88,7 @@ public class MemoryNavigableLRUCache extends MemoryLRUCache {
@Override
public <PS extends Serializer<P>, P> KeyValueIterator<Bytes, byte[]>
prefixScan(final P prefix, final PS prefixKeySerializer) {
-
+ // headers aren't needed because the prefix already arrives serialized
final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null,
prefix));
final Bytes to = ByteUtils.increment(from);
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java
index ca8f12ac743..68c3795345f 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java
@@ -19,6 +19,7 @@ package org.apache.kafka.streams.state.internals;
import org.apache.kafka.common.IsolationLevel;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.Bytes;
@@ -79,6 +80,8 @@ public class MeteredKeyValueStore<K, V>
extends WrappedStateStore<KeyValueStore<Bytes, byte[]>, K, V>
implements KeyValueStore<K, V>, MeteredStateStore {
+ private static final Serializer<byte[]> BYTE_ARRAY_SERIALIZER = new
ByteArraySerializer();
+
final Serde<K> keySerde;
final Serde<V> valueSerde;
StateSerdes<K, V> serdes;
@@ -408,7 +411,8 @@ public class MeteredKeyValueStore<K, V>
private <PS extends Serializer<P>, P> KeyValueIterator<K, V>
prefixScanInternal(
final ReadOnlyKeyValueStore<Bytes, byte[]> store, final P prefix,
final PS prefixKeySerializer
) {
- return meteredKeyValueIterator(store.prefixScan(prefix,
prefixKeySerializer), prefixScanSensor);
+ final byte[] keyBytes = prefixKeySerializer.serialize(null,
internalContext.headers(), prefix);
+ return meteredKeyValueIterator(store.prefixScan(keyBytes,
BYTE_ARRAY_SERIALIZER), prefixScanSensor);
}
@Override
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java
index 3d10e5569cc..4aa0b63982d 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeaders.java
@@ -20,6 +20,7 @@ import org.apache.kafka.common.IsolationLevel;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.Bytes;
@@ -79,6 +80,8 @@ public class MeteredTimestampedKeyValueStoreWithHeaders<K, V>
extends MeteredKeyValueStore<K, ValueTimestampHeaders<V>>
implements TimestampedKeyValueStoreWithHeaders<K, V> {
+ private static final Serializer<byte[]> BYTE_ARRAY_SERIALIZER = new
ByteArraySerializer();
+
MeteredTimestampedKeyValueStoreWithHeaders(
final KeyValueStore<Bytes, byte[]> inner,
final String metricScope,
@@ -592,8 +595,9 @@ public class MeteredTimestampedKeyValueStoreWithHeaders<K,
V>
private <PS extends Serializer<P>, P> KeyValueIterator<K,
ValueTimestampHeaders<V>> prefixScanInternal(
final ReadOnlyKeyValueStore<Bytes, byte[]> store, final P prefix,
final PS prefixKeySerializer
) {
+ final byte[] keyBytes = prefixKeySerializer.serialize(null,
internalContext.headers(), prefix);
return new MeteredTimestampedKeyValueStoreWithHeadersIterator(
- store.prefixScan(prefix, prefixKeySerializer), prefixScanSensor
+ store.prefixScan(keyBytes, BYTE_ARRAY_SERIALIZER), prefixScanSensor
);
}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStoreTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStoreTest.java
index d5d65e56be5..94911f5ece6 100644
---
a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStoreTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStoreTest.java
@@ -27,6 +27,7 @@ import org.apache.kafka.common.metrics.KafkaMetricsContext;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.metrics.MetricsContext;
import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serdes;
@@ -521,20 +522,26 @@ public class MeteredKeyValueStoreTest {
assertThrows(NullPointerException.class, () ->
metered.reverseRange("from", null));
}
+ @SuppressWarnings("unchecked")
@Test
public void shouldGetRecordsWithPrefixKey() {
setUp();
- final StringSerializer stringSerializer = new StringSerializer();
- when(inner.prefixScan(KEY, stringSerializer))
+ final Serializer<String> mockSerializer = mock(Serializer.class);
+ final Headers headers = new RecordHeaders();
+ when(context.headers()).thenReturn(headers);
+ when(mockSerializer.serialize(null, headers,
KEY)).thenReturn(KEY.getBytes(StandardCharsets.UTF_8));
+
+ when(inner.prefixScan(eq(KEY.getBytes(StandardCharsets.UTF_8)),
any(ByteArraySerializer.class)))
.thenReturn(new
KeyValueIteratorStub<>(Collections.singletonList(BYTE_KEY_VALUE_PAIR).iterator()));
init();
- final KeyValueIterator<String, String> iterator =
metered.prefixScan(KEY, stringSerializer);
+ final KeyValueIterator<String, String> iterator =
metered.prefixScan(KEY, mockSerializer);
assertThat(iterator.next().value, equalTo(VALUE));
iterator.close();
final KafkaMetric metric = metrics.metric(new
MetricName("prefix-scan-rate", STORE_LEVEL_GROUP, "", tags));
assertTrue((Double) metric.metricValue() > 0);
+ verify(mockSerializer).serialize(null, headers, KEY);
}
@Test
@@ -548,12 +555,17 @@ public class MeteredKeyValueStoreTest {
assertThat((Long) numKeysMetric.metricValue(), equalTo(-1L));
}
- @SuppressWarnings("unused")
+ @SuppressWarnings({"unused", "unchecked"})
@Test
public void shouldTrackOpenIteratorsMetric() {
setUp();
- final StringSerializer stringSerializer = new StringSerializer();
- when(inner.prefixScan(KEY,
stringSerializer)).thenReturn(KeyValueIterators.emptyIterator());
+ final Serializer<String> mockSerializer = mock(Serializer.class);
+ final Headers headers = new RecordHeaders();
+ when(context.headers()).thenReturn(headers);
+ when(mockSerializer.serialize(null, headers,
KEY)).thenReturn(KEY.getBytes(StandardCharsets.UTF_8));
+
+ when(inner.prefixScan(eq(KEY.getBytes(StandardCharsets.UTF_8)),
any(ByteArraySerializer.class)))
+ .thenReturn(KeyValueIterators.emptyIterator());
init();
final KafkaMetric openIteratorsMetric = metric("num-open-iterators");
@@ -561,11 +573,12 @@ public class MeteredKeyValueStoreTest {
assertThat((Long) openIteratorsMetric.metricValue(), equalTo(0L));
- try (final KeyValueIterator<String, String> unused =
metered.prefixScan(KEY, stringSerializer)) {
+ try (final KeyValueIterator<String, String> unused =
metered.prefixScan(KEY, mockSerializer)) {
assertThat((Long) openIteratorsMetric.metricValue(), equalTo(1L));
}
assertThat((Long) openIteratorsMetric.metricValue(), equalTo(0L));
+ verify(mockSerializer).serialize(null, headers, KEY);
}
@SuppressWarnings("unused")
@@ -734,19 +747,25 @@ public class MeteredKeyValueStoreTest {
public void
shouldReadOnlyViewPrefixScanApplySerdesAndRecordPrefixScanMetric() {
setUp();
final ReadOnlyKeyValueStore<Bytes, byte[]> innerView =
mock(ReadOnlyKeyValueStore.class);
- final StringSerializer stringSerializer = new StringSerializer();
+ final Serializer<String> mockSerializer = mock(Serializer.class);
+ final Headers headers = new RecordHeaders();
+ when(context.headers()).thenReturn(headers);
+ when(mockSerializer.serialize(null, headers,
KEY)).thenReturn(KEY.getBytes(StandardCharsets.UTF_8));
+
when(inner.readOnly(IsolationLevel.READ_UNCOMMITTED)).thenReturn(innerView);
- when(innerView.prefixScan(KEY, stringSerializer))
+ when(innerView.prefixScan(eq(KEY.getBytes(StandardCharsets.UTF_8)),
any(ByteArraySerializer.class)))
.thenReturn(new
KeyValueIteratorStub<>(Collections.singletonList(BYTE_KEY_VALUE_PAIR).iterator()));
init();
final ReadOnlyKeyValueStore<String, String> view =
metered.readOnly(IsolationLevel.READ_UNCOMMITTED);
- try (final KeyValueIterator<String, String> it = view.prefixScan(KEY,
stringSerializer)) {
+ try (final KeyValueIterator<String, String> it = view.prefixScan(KEY,
mockSerializer)) {
assertThat(it.next().value, equalTo(VALUE));
assertFalse(it.hasNext());
}
- assertTrue((Double) metrics.metric(new MetricName("prefix-scan-rate",
STORE_LEVEL_GROUP, "", tags)).metricValue() > 0);
+ final KafkaMetric metric = metric("prefix-scan-rate");
+ assertTrue((Double) metric.metricValue() > 0);
+ verify(mockSerializer).serialize(null, headers, KEY);
}
@SuppressWarnings("unchecked")
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java
index 7a45afae1e8..14087a855ea 100644
---
a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreWithHeadersTest.java
@@ -15,7 +15,10 @@
* limitations under the License.
*/
package org.apache.kafka.streams.state.internals;
+
+import org.apache.kafka.common.IsolationLevel;
import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.metrics.JmxReporter;
import org.apache.kafka.common.metrics.KafkaMetric;
@@ -29,6 +32,7 @@ import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.internals.LogContext;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.errors.StreamsException;
@@ -63,6 +67,7 @@ import org.mockito.quality.Strictness;
import java.io.Closeable;
import java.io.IOException;
+import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -81,6 +86,7 @@ import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -868,4 +874,60 @@ public class
MeteredTimestampedKeyValueStoreWithHeadersTest {
// The critical verification: key deserializer must have been called
with HEADERS (not empty headers)
verify(keyDeserializer).deserialize(any(), eq(HEADERS),
eq(KEY.getBytes()));
}
+
+ @Test
+ public void shouldReturnRightEntriesForHeaderDependentSerdeInPrefixScan() {
+ setUp();
+ reset(inner);
+ final ThreadCache cache = new ThreadCache(new LogContext("testCache"),
1024L, context.metrics());
+ when(context.cache()).thenReturn(cache);
+ when(context.headers()).thenReturn(HEADERS);
+
+ final InMemoryKeyValueStore inMemoryStore = new
InMemoryKeyValueStore(STORE_NAME);
+
+ final Serializer<String> headerDependentSerializer = new
Serializer<>() {
+ @Override
+ public byte[] serialize(final String topic, final String data) {
+ return serialize(topic, null, data);
+ }
+
+ @Override
+ public byte[] serialize(final String topic, final Headers headers,
final String data) {
+ final byte[] bytes = data.getBytes(StandardCharsets.UTF_8);
+ final byte[] result = new byte[bytes.length + 1];
+ result[0] = headers == null ? (byte) 1 : (byte) 2;
+ System.arraycopy(bytes, 0, result, 1, bytes.length);
+ return result;
+ }
+ };
+
+ final MeteredTimestampedKeyValueStoreWithHeaders<String, String> store
= new MeteredTimestampedKeyValueStoreWithHeaders<>(
+ new CachingKeyValueStoreWithHeaders(inMemoryStore),
+ "scope",
+ new MockTime(),
+ Serdes.serdeFrom(headerDependentSerializer,
Serdes.String().deserializer()),
+ new ValueTimestampHeadersSerde<>(Serdes.String())
+ );
+
+ store.init(context, store);
+
+ store.put("key1", ValueTimestampHeaders.make("value1", 100L, HEADERS));
+ store.put("key2", ValueTimestampHeaders.make("value2", 100L, HEADERS));
+
+ // 1. Test default prefixScan
+ try (final KeyValueIterator<String, ValueTimestampHeaders<String>>
iterator = store.prefixScan("key", headerDependentSerializer)) {
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.next().key.endsWith("key1"));
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.next().key.endsWith("key2"));
+ }
+
+ // 2. Test readOnly prefixScan
+ try (final KeyValueIterator<String, ValueTimestampHeaders<String>>
iterator = store.readOnly(IsolationLevel.READ_UNCOMMITTED).prefixScan("key",
headerDependentSerializer)) {
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.next().key.endsWith("key1"));
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.next().key.endsWith("key2"));
+ }
+ }
}
