diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeaders.java index cbddac333cc79..65758c4959150 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeaders.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.state.internals; +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.serialization.Serde; @@ -35,6 +36,7 @@ import org.apache.kafka.streams.query.internals.InternalQueryResultUtil; import org.apache.kafka.streams.state.AggregationWithHeaders; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.ReadOnlySessionStore; import org.apache.kafka.streams.state.SessionStore; import org.apache.kafka.streams.state.SessionStoreWithHeaders; @@ -360,6 +362,101 @@ private QueryResult runRangeQuery( return queryResult; } + @Override + public ReadOnlySessionStore> readOnly(final IsolationLevel isolationLevel) { + Objects.requireNonNull(isolationLevel, "isolationLevel cannot be null"); + return new ReadOnlyHeadersView(wrapped().readOnly(isolationLevel)); + } + + private final class ReadOnlyHeadersView implements ReadOnlySessionStore> { + + private final ReadOnlySessionStore underlying; + + ReadOnlyHeadersView(final ReadOnlySessionStore underlying) { + this.underlying = underlying; + } + + @Override + public AggregationWithHeaders fetchSession( + final K key, final long earliestSessionEndTime, final long latestSessionStartTime) { + Objects.requireNonNull(key, "key cannot be null"); + return maybeMeasureLatency( + () -> deserializeValue(underlying.fetchSession( + serializeKey(key, internalContext.headers()), earliestSessionEndTime, latestSessionStartTime)), + time, + fetchSensor + ); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> fetch(final K key) { + Objects.requireNonNull(key, "key cannot be null"); + return new MeteredSessionStoreWithHeadersIterator(underlying.fetch(serializeKey(key, internalContext.headers()))); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> backwardFetch(final K key) { + Objects.requireNonNull(key, "key cannot be null"); + return new MeteredSessionStoreWithHeadersIterator(underlying.backwardFetch(serializeKey(key, internalContext.headers()))); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> fetch(final K keyFrom, final K keyTo) { + return new MeteredSessionStoreWithHeadersIterator( + underlying.fetch(serializeKey(keyFrom, internalContext.headers()), serializeKey(keyTo, internalContext.headers())) + ); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> backwardFetch(final K keyFrom, final K keyTo) { + return new MeteredSessionStoreWithHeadersIterator( + underlying.backwardFetch(serializeKey(keyFrom, internalContext.headers()), serializeKey(keyTo, internalContext.headers())) + ); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> findSessions( + final K key, final long earliestSessionEndTime, final long latestSessionStartTime) { + Objects.requireNonNull(key, "key cannot be null"); + return new MeteredSessionStoreWithHeadersIterator( + underlying.findSessions(serializeKey(key, internalContext.headers()), earliestSessionEndTime, latestSessionStartTime) + ); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> backwardFindSessions( + final K key, final long earliestSessionEndTime, final long latestSessionStartTime) { + Objects.requireNonNull(key, "key cannot be null"); + return new MeteredSessionStoreWithHeadersIterator( + underlying.backwardFindSessions(serializeKey(key, internalContext.headers()), earliestSessionEndTime, latestSessionStartTime) + ); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> findSessions( + final K keyFrom, final K keyTo, final long earliestSessionEndTime, final long latestSessionStartTime) { + return new MeteredSessionStoreWithHeadersIterator( + underlying.findSessions( + serializeKey(keyFrom, internalContext.headers()), + serializeKey(keyTo, internalContext.headers()), + earliestSessionEndTime, + latestSessionStartTime) + ); + } + + @Override + public KeyValueIterator, AggregationWithHeaders> backwardFindSessions( + final K keyFrom, final K keyTo, final long earliestSessionEndTime, final long latestSessionStartTime) { + return new MeteredSessionStoreWithHeadersIterator( + underlying.backwardFindSessions( + serializeKey(keyFrom, internalContext.headers()), + serializeKey(keyTo, internalContext.headers()), + earliestSessionEndTime, + latestSessionStartTime) + ); + } + } + private class MeteredSessionStoreWithHeadersIterator implements KeyValueIterator, AggregationWithHeaders>, MeteredIterator { diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java index c1ab307001baa..eb02894fbf608 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeaders.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.state.internals; +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.serialization.Serde; @@ -36,6 +37,7 @@ import org.apache.kafka.streams.query.WindowRangeQuery; import org.apache.kafka.streams.query.internals.InternalQueryResultUtil; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.ReadOnlyWindowStore; import org.apache.kafka.streams.state.TimestampedBytesStore; import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; import org.apache.kafka.streams.state.ValueAndTimestamp; @@ -43,6 +45,7 @@ import org.apache.kafka.streams.state.WindowStore; import org.apache.kafka.streams.state.WindowStoreIterator; +import java.time.Instant; import java.util.Objects; import java.util.function.Function; @@ -347,6 +350,110 @@ public KeyValueIterator, ValueTimestampHeaders> backwardAll() { ); } + @Override + public ReadOnlyWindowStore> readOnly(final IsolationLevel isolationLevel) { + Objects.requireNonNull(isolationLevel, "isolationLevel cannot be null"); + return new ReadOnlyHeadersView(wrapped().readOnly(isolationLevel)); + } + + private final class ReadOnlyHeadersView implements ReadOnlyWindowStore> { + + private final ReadOnlyWindowStore underlying; + + ReadOnlyHeadersView(final ReadOnlyWindowStore underlying) { + this.underlying = underlying; + } + + @Override + public ValueTimestampHeaders fetch(final K key, final long windowStartTimestamp) { + Objects.requireNonNull(key, "key cannot be null"); + return maybeMeasureLatency( + () -> { + final byte[] result = underlying.fetch(serializeKey(key, internalContext.headers()), windowStartTimestamp); + return result == null ? null : deserializeValue(result); + }, + time, + fetchSensor + ); + } + + @Override + public WindowStoreIterator> fetch( + final K key, final Instant timeFrom, final Instant timeTo) { + Objects.requireNonNull(key, "key cannot be null"); + return new MeteredWindowStoreIterator<>( + underlying.fetch(serializeKey(key, internalContext.headers()), timeFrom, timeTo), + fetchSensor, + iteratorDurationSensor, + MeteredTimestampedWindowStoreWithHeaders.this::deserializeValue, + time, + numOpenIterators, + openIterators + ); + } + + @Override + public WindowStoreIterator> backwardFetch( + final K key, final Instant timeFrom, final Instant timeTo) { + Objects.requireNonNull(key, "key cannot be null"); + return new MeteredWindowStoreIterator<>( + underlying.backwardFetch(serializeKey(key, internalContext.headers()), timeFrom, timeTo), + fetchSensor, + iteratorDurationSensor, + MeteredTimestampedWindowStoreWithHeaders.this::deserializeValue, + time, + numOpenIterators, + openIterators + ); + } + + @Override + public KeyValueIterator, ValueTimestampHeaders> fetch( + final K keyFrom, final K keyTo, final Instant timeFrom, final Instant timeTo) { + return new MeteredTimestampedWindowStoreWithHeadersKeyValueIterator( + underlying.fetch( + serializeKey(keyFrom, internalContext.headers()), + serializeKey(keyTo, internalContext.headers()), + timeFrom, + timeTo) + ); + } + + @Override + public KeyValueIterator, ValueTimestampHeaders> backwardFetch( + final K keyFrom, final K keyTo, final Instant timeFrom, final Instant timeTo) { + return new MeteredTimestampedWindowStoreWithHeadersKeyValueIterator( + underlying.backwardFetch( + serializeKey(keyFrom, internalContext.headers()), + serializeKey(keyTo, internalContext.headers()), + timeFrom, + timeTo) + ); + } + + @Override + public KeyValueIterator, ValueTimestampHeaders> all() { + return new MeteredTimestampedWindowStoreWithHeadersKeyValueIterator(underlying.all()); + } + + @Override + public KeyValueIterator, ValueTimestampHeaders> backwardAll() { + return new MeteredTimestampedWindowStoreWithHeadersKeyValueIterator(underlying.backwardAll()); + } + + @Override + public KeyValueIterator, ValueTimestampHeaders> fetchAll( + final Instant timeFrom, final Instant timeTo) { + return new MeteredTimestampedWindowStoreWithHeadersKeyValueIterator(underlying.fetchAll(timeFrom, timeTo)); + } + + @Override + public KeyValueIterator, ValueTimestampHeaders> backwardFetchAll( + final Instant timeFrom, final Instant timeTo) { + return new MeteredTimestampedWindowStoreWithHeadersKeyValueIterator(underlying.backwardFetchAll(timeFrom, timeTo)); + } + } + private class MeteredTimestampedWindowStoreWithHeadersKeyValueIterator implements KeyValueIterator, ValueTimestampHeaders>, MeteredIterator { diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeadersTest.java index aafa28d00b27b..6b9d602cf8e41 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredSessionStoreWithHeadersTest.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.state.internals; +import org.apache.kafka.common.IsolationLevel; import org.apache.kafka.common.Metric; import org.apache.kafka.common.MetricName; import org.apache.kafka.common.header.Headers; @@ -47,6 +48,7 @@ import org.apache.kafka.streams.query.WindowRangeQuery; import org.apache.kafka.streams.state.AggregationWithHeaders; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.ReadOnlySessionStore; import org.apache.kafka.streams.state.SessionStore; import org.apache.kafka.test.KeyValueIteratorStub; @@ -1050,6 +1052,33 @@ public void shouldUseHeadersFromValueToDeserializeKeyInBackwardFindSessionsRange verify(keySerde.deserializer()).deserialize(any(), eq(HEADERS), eq(KEY.getBytes())); } + @SuppressWarnings("unchecked") + @Test + public void shouldUseHeadersFromValueToDeserializeKeyInReadOnlyFindSessions() { + setUp(); + final Serde keySerde = mock(Serde.class); + final MeteredSessionStoreWithHeaders store = createStoreWithMockSerdes(keySerde); + + final ReadOnlySessionStore readOnlyInner = mock(ReadOnlySessionStore.class); + when(innerStore.readOnly(IsolationLevel.READ_COMMITTED)).thenReturn(readOnlyInner); + when(readOnlyInner.findSessions(any(Bytes.class), eq(0L), eq(100L))) + .thenReturn(new KeyValueIteratorStub<>( + List.of(KeyValue.pair(WINDOWED_KEY_BYTES, SERIALIZED_VALUE)).iterator())); + + final KeyValueIterator, AggregationWithHeaders> iterator = + store.readOnly(IsolationLevel.READ_COMMITTED).findSessions(KEY, 0L, 100L); + + assertTrue(iterator.hasNext()); + assertEquals(KEY, iterator.peekNextKey().key()); + final KeyValue, AggregationWithHeaders> result = iterator.next(); + assertEquals(KEY, result.key.key()); + assertEquals(AGG_WITH_HEADERS, result.value); + assertFalse(iterator.hasNext()); + iterator.close(); + + verify(keySerde.deserializer()).deserialize(any(), eq(HEADERS), eq(KEY.getBytes())); + } + @SuppressWarnings("unchecked") @Test public void shouldUseHeadersFromValueToDeserializeKeyInFindSessionsByTime() { diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java index eb2c6a65f8ebb..1f0dcbb1fdcd3 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedWindowStoreWithHeadersTest.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.state.internals; +import org.apache.kafka.common.IsolationLevel; import org.apache.kafka.common.header.internals.RecordHeaders; import org.apache.kafka.common.metrics.MetricConfig; import org.apache.kafka.common.metrics.Metrics; @@ -39,6 +40,7 @@ import org.apache.kafka.streams.processor.internals.ProcessorStateManager; import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.ReadOnlyWindowStore; import org.apache.kafka.streams.state.ValueTimestampHeaders; import org.apache.kafka.streams.state.WindowStore; import org.apache.kafka.test.InternalMockProcessorContext; @@ -54,6 +56,7 @@ import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import java.time.Instant; import java.util.List; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -484,4 +487,33 @@ public void shouldUseHeadersFromValueToDeserializeKeyInBackwardFetchRange() { verify(keyDeserializer).deserialize(any(), eq(HEADERS), eq(KEY.getBytes())); } + + @Test + public void shouldUseHeadersFromValueToDeserializeKeyInReadOnlyFetchAll() { + setUp(); + + final Windowed windowedKey = new Windowed<>(KEY_BYTES, new TimeWindow(0, WINDOW_SIZE_MS)); + final KeyValue, byte[]> testData = KeyValue.pair(windowedKey, VALUE_TIMESTAMP_HEADERS_BYTES); + + final ReadOnlyWindowStore readOnlyInner = mock(ReadOnlyWindowStore.class); + when(innerStoreMock.readOnly(IsolationLevel.READ_COMMITTED)).thenReturn(readOnlyInner); + when(readOnlyInner.fetchAll(Instant.ofEpochMilli(0), Instant.ofEpochMilli(100))) + .thenReturn(new KeyValueIteratorStub<>(List.of(testData).iterator())); + + store = createStoreWithMockSerdes(); + + final KeyValueIterator, ValueTimestampHeaders> iterator = + store.readOnly(IsolationLevel.READ_COMMITTED).fetchAll(Instant.ofEpochMilli(0), Instant.ofEpochMilli(100)); + + assertTrue(iterator.hasNext()); + assertEquals(KEY, iterator.peekNextKey().key()); + final KeyValue, ValueTimestampHeaders> result = iterator.next(); + + assertEquals(KEY, result.key.key()); + assertEquals(VALUE_TIMESTAMP_HEADERS, result.value); + assertFalse(iterator.hasNext()); + iterator.close(); + + verify(keyDeserializer).deserialize(any(), eq(HEADERS), eq(KEY.getBytes())); + } }