diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java index 9b1442883ca77..a67e5784c8826 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/IQv2HeadersStoreIntegrationTest.java @@ -35,6 +35,7 @@ import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; import org.apache.kafka.streams.kstream.Consumed; import org.apache.kafka.streams.kstream.Produced; +import org.apache.kafka.streams.kstream.Windowed; import org.apache.kafka.streams.processor.api.Processor; import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.ProcessorSupplier; @@ -49,11 +50,14 @@ import org.apache.kafka.streams.query.StateQueryResult; import org.apache.kafka.streams.query.TimestampedKeyWithHeadersQuery; import org.apache.kafka.streams.query.TimestampedRangeWithHeadersQuery; +import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery; import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.StoreBuilder; import org.apache.kafka.streams.state.Stores; import org.apache.kafka.streams.state.TimestampedKeyValueStore; import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.TimestampedWindowStore; +import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; import org.apache.kafka.streams.state.ValueAndTimestamp; import org.apache.kafka.streams.state.ValueTimestampHeaders; import org.apache.kafka.test.TestUtils; @@ -68,6 +72,7 @@ import java.io.IOException; import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.LinkedList; import java.util.List; @@ -87,9 +92,9 @@ * IQv2 integration tests for KIP-1271/KIP-1356 headers-aware state stores. * *

It builds a KIP-1271 {@code WithHeaders} store, writes records (with headers) into it through a processor, - * and queries it through IQv2. Covers {@link TimestampedKeyWithHeadersQuery} and - * {@link TimestampedRangeWithHeadersQuery}; the remaining KIP-1356 headers query types (window/session) are - * expected to extend this class as they land. + * and queries it through IQv2. Covers {@link TimestampedKeyWithHeadersQuery}, + * {@link TimestampedRangeWithHeadersQuery}, and {@link TimestampedWindowKeyWithHeadersQuery}; the remaining + * KIP-1356 headers query type (window range / session) is expected to extend this class as it lands. */ @Tag("integration") public class IQv2HeadersStoreIntegrationTest { @@ -102,6 +107,14 @@ public class IQv2HeadersStoreIntegrationTest { .add("source", "test".getBytes()) .add("version", "1.0".getBytes()); + // Window-store test parameters (used by the window headers query tests). + private static final long WINDOW_SIZE_MS = 10L; + // Stored event-time is offset a few ms into the window so tests can tell ReadOnlyRecord.timestamp() + // (the event-time) apart from the window's start and end. Must satisfy 0 < offset < WINDOW_SIZE_MS. + private static final long EVENT_TIME_OFFSET_MS = 3L; + private static final Duration WINDOW_SIZE = Duration.ofMillis(WINDOW_SIZE_MS); + private static final Duration RETENTION = Duration.ofMinutes(1); + private String inputStream; private String outputStream; private long baseTimestamp; @@ -392,6 +405,140 @@ public void shouldReturnEmptyHeadersForTimestampedRangeWithHeadersQueryOnAdapter assertThrows(IllegalStateException.class, () -> range.get(0).headers().remove("source")); } + @Test + public void shouldHandleTimestampedWindowKeyWithHeadersQuery() throws Exception { + startStreamsWithWindowHeadersStore(); + + final long window0 = baseTimestamp; + final long window1 = baseTimestamp + 100; + final long window2 = baseTimestamp + 200; + final long window3 = baseTimestamp + 300; + final long window4 = baseTimestamp + 400; + + // Queried key 1 across four window starts: headers, empty headers, tombstoned, headers. + produceDataToTopicWithHeaders(inputStream, window0, HEADERS, KeyValue.pair(1, "v0")); + produceDataToTopicWithHeaders(inputStream, window1, new RecordHeaders(), KeyValue.pair(1, "v1")); + produceDataToTopicWithHeaders(inputStream, window2, HEADERS, KeyValue.pair(1, "v2"), KeyValue.pair(1, null)); + produceDataToTopicWithHeaders(inputStream, window3, HEADERS, KeyValue.pair(1, "v3")); + + // Noise key 2 must never leak into a query for key 1: a tombstone at window0 (which must not remove + // key 1's window0), a live entry at the shared window1, and a live entry at window4 (a window key 1 + // never uses). + produceDataToTopicWithHeaders(inputStream, window0, HEADERS, KeyValue.pair(2, "o0"), KeyValue.pair(2, null)); + produceDataToTopicWithHeaders(inputStream, window1, HEADERS, KeyValue.pair(2, "o1")); + produceDataToTopicWithHeaders(inputStream, window4, HEADERS, KeyValue.pair(2, "o4")); + + // Full range over key 1: window0, window1, window3 in window-start order. window2 is tombstoned; + // key 2's entries are excluded; and key 2's window0 tombstone did not drop key 1's window0. + final List, String>> records = + windowKeyQuery(1, Instant.ofEpochMilli(window0), Instant.ofEpochMilli(window4)); + assertEquals( + List.of(window0, window1, window3), + records.stream().map(r -> r.key().window().start()).collect(Collectors.toList()), + "expected window0, window1, window3 in order (window2 tombstoned, key 2 excluded)"); + assertWindowedRecord(records.get(0), 1, window0, "v0", HEADERS); + assertWindowedRecord(records.get(1), 1, window1, "v1", new RecordHeaders()); + assertWindowedRecord(records.get(2), 1, window3, "v3", HEADERS); + + // Sub-range [window1, window2]: only window1 (window0 is below the lower bound, window2 is + // tombstoned, window3 is above the upper bound). + final List, String>> subRange = + windowKeyQuery(1, Instant.ofEpochMilli(window1), Instant.ofEpochMilli(window2)); + assertEquals(1, subRange.size()); + assertWindowedRecord(subRange.get(0), 1, window1, "v1", new RecordHeaders()); + + // A never-written key -> empty result. + assertTrue(windowKeyQuery(999, Instant.ofEpochMilli(window0), Instant.ofEpochMilli(window4)).isEmpty(), + "expected no records for a never-written key"); + // key 1, a window-start range entirely after its last window -> empty result. + assertTrue(windowKeyQuery(1, Instant.ofEpochMilli(window3 + WINDOW_SIZE_MS), Instant.ofEpochMilli(window4)).isEmpty(), + "expected no records for a window-start range that excludes all of key 1's windows"); + + // Querying key 2 returns only its live entries (window1, window4; window0 tombstoned) -- isolation is + // bidirectional and key 2's own tombstone is key-scoped. + final List, String>> key2Records = + windowKeyQuery(2, Instant.ofEpochMilli(window0), Instant.ofEpochMilli(window4)); + assertEquals( + List.of(window1, window4), + key2Records.stream().map(r -> r.key().window().start()).collect(Collectors.toList()), + "expected key 2's window1 and window4 (window0 tombstoned)"); + assertWindowedRecord(key2Records.get(0), 2, window1, "o1", HEADERS); + assertWindowedRecord(key2Records.get(1), 2, window4, "o4", HEADERS); + } + + @Test + public void shouldNotSeeUnflushedWriteInWindowKeyQueryWhenCachingEnabled() throws Exception { + // A window key query bypasses the record cache: CachingWindowStore does not intercept IQv2 queries + // (unlike CachingKeyValueStore), so WrappedStateStore.query() forwards straight to the persistent + // store. That store's Position is only advanced by writes that actually reach it, so a query bound on + // the input position never catches up while the write lives only in the cache -- it fails + // NOT_UP_TO_BOUND, rather than returning a result. Use a very large commit interval so the record is + // never flushed during the test. + commitIntervalMs = Duration.ofMinutes(10).toMillis(); + startStreamsWithWindowHeadersStore(true); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "v0")); + + final StateQueryRequest, String>> request = + inStore(STORE_NAME) + .withQuery(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + 1, Instant.ofEpochMilli(baseTimestamp), Instant.ofEpochMilli(baseTimestamp))) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult, String>> result = kafkaStreams.query(request); + + final QueryResult, String>> onlyResult = result.getOnlyPartitionResult(); + assertTrue(onlyResult.isFailure(), "A window key query bound on the input position must not catch up while the write is cache-only"); + assertEquals(FailureReason.NOT_UP_TO_BOUND, onlyResult.getFailureReason()); + } + + @Test + public void shouldFailWithUnknownQueryTypeForWindowKeyQueryAgainstNonHeadersStore() throws Exception { + startStreamsWithWindowNonHeadersStore(); + assertUnknownQueryType(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + 1, Instant.ofEpochMilli(baseTimestamp), Instant.ofEpochMilli(baseTimestamp + 1))); + } + + @Test + public void shouldThrowForTimestampedWindowKeyWithHeadersQueryOnPlainSupplier() throws Exception { + // A WithHeaders window builder over a plain (non-timestamped) window supplier: entries come back + // with timestamp = -1. The query succeeds, but iterating throws because -1 cannot be a ReadOnlyRecord. + startStreamsWithWindowPlainSupplierStore(); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "one")); + + final StateQueryRequest, String>> request = + inStore(STORE_NAME) + .withQuery(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + 1, Instant.ofEpochMilli(baseTimestamp), Instant.ofEpochMilli(baseTimestamp + 1))) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult, String>> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + + final QueryResult, String>> onlyResult = result.getOnlyPartitionResult(); + assertTrue(onlyResult.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = onlyResult.getResult()) { + final StreamsException e = assertThrows(StreamsException.class, iterator::next); + assertTrue(e.getMessage().contains("as a ReadOnlyRecord") && e.getMessage().contains("is negative"), + "unexpected message: " + e.getMessage()); + } + } + + @Test + public void shouldReturnEmptyHeadersForTimestampedWindowKeyWithHeadersQueryOnAdapterStore() throws Exception { + // A WithHeaders window builder over a plain *timestamped* window supplier keeps timestamps + // (unlike the plain-supplier build above) but drops headers via + // TimestampedToHeadersWindowStoreAdapter. A window key query reads the underlying store + // directly, so headers come back empty (never null), even though they were written. + startStreamsWithWindowAdapterStore(); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "one")); + + final List, String>> records = + windowKeyQuery(1, Instant.ofEpochMilli(baseTimestamp), Instant.ofEpochMilli(baseTimestamp)); + assertEquals(1, records.size()); + assertWindowedRecord(records.get(0), 1, baseTimestamp, "one", new RecordHeaders()); + } + private void startStreams(final StoreBuilder storeBuilder, final ProcessorSupplier processorSupplier) throws Exception { final StreamsBuilder builder = new StreamsBuilder(); @@ -455,8 +602,62 @@ private void startStreamsWithKeyValueAdapterStore() throws Exception { KeyValueHeadersStoreWriterProcessor::new); } + private void startStreamsWithWindowHeadersStore() throws Exception { + // Caching disabled: every IQv2 query is forced down to the persistent + // RocksDBTimestampedWindowStoreWithHeaders layer. + startStreamsWithWindowHeadersStore(false); + } + + private void startStreamsWithWindowHeadersStore(final boolean cachingEnabled) throws Exception { + final StoreBuilder> storeBuilder = + Stores.timestampedWindowStoreWithHeadersBuilder( + Stores.persistentTimestampedWindowStoreWithHeaders(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()); + startStreams( + cachingEnabled ? storeBuilder.withCachingEnabled() : storeBuilder.withCachingDisabled(), + WindowHeadersStoreWriterProcessor::new); + } + + private void startStreamsWithWindowNonHeadersStore() throws Exception { + // A plain (non-WithHeaders) timestamped window store: the headers-aware query types are unsupported here. + startStreams( + Stores.timestampedWindowStoreBuilder( + Stores.persistentTimestampedWindowStore(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()), + WindowPlainStoreWriterProcessor::new); + } + + private void startStreamsWithWindowPlainSupplierStore() throws Exception { + // A WithHeaders window builder over a plain (non-timestamped) window supplier: entries come back + // with timestamp = -1, which cannot be represented as a ReadOnlyRecord. + startStreams( + Stores.timestampedWindowStoreWithHeadersBuilder( + Stores.persistentWindowStore(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()).withCachingDisabled(), + WindowHeadersStoreWriterProcessor::new); + } + + private void startStreamsWithWindowAdapterStore() throws Exception { + // A WithHeaders window builder over a plain *timestamped* window supplier: the adapter keeps + // timestamps but cannot persist headers, so window key results carry value + timestamp with empty + // headers. + startStreams( + Stores.timestampedWindowStoreWithHeadersBuilder( + Stores.persistentTimestampedWindowStore(STORE_NAME, RETENTION, WINDOW_SIZE, false), + Serdes.Integer(), + Serdes.String()).withCachingDisabled(), + WindowHeadersStoreWriterProcessor::new); + } + private void assertUnknownQueryTypeAgainstNonHeadersStore(final Query query) throws Exception { startStreamsWithKeyValueNonHeadersStore(); + assertUnknownQueryType(query); + } + + private void assertUnknownQueryType(final Query query) { produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); final StateQueryRequest request = @@ -508,6 +709,42 @@ private List> rangeQuery(final TimestampedRangeW return records; } + private List, String>> windowKeyQuery(final int key, + final Instant timeFrom, + final Instant timeTo) { + final StateQueryRequest, String>> request = + inStore(STORE_NAME) + .withQuery(TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + key, timeFrom, timeTo)) + .withPositionBound(PositionBound.at(inputPosition)); + final StateQueryResult, String>> result = + IntegrationTestUtils.iqv2WaitForResult(kafkaStreams, request); + final List, String>> records = new ArrayList<>(); + try (ReadOnlyRecordIterator, String> iterator = result.getOnlyPartitionResult().getResult()) { + while (iterator.hasNext()) { + records.add(iterator.next()); + } + } + return records; + } + + private void assertWindowedRecord(final ReadOnlyRecord, String> record, + final int key, + final long windowStart, + final String value, + final Headers expectedHeaders) { + assertEquals(Integer.valueOf(key), record.key().key()); + assertEquals(windowStart, record.key().window().start()); + assertEquals(windowStart + WINDOW_SIZE_MS, record.key().window().end()); + // timestamp() is the stored event-time, distinct from the window's start and end. + assertEquals(windowStart + EVENT_TIME_OFFSET_MS, record.timestamp()); + assertEquals(value, record.value()); + assertEquals(expectedHeaders, record.headers()); + // The IQ result is a read-only snapshot: neither add nor remove is allowed. + assertThrows(IllegalStateException.class, () -> record.headers().add("x", new byte[0])); + assertThrows(IllegalStateException.class, () -> record.headers().remove("source")); + } + private static List keys(final List> records) { return records.stream().map(ReadOnlyRecord::key).collect(Collectors.toList()); } @@ -596,4 +833,51 @@ public void process(final Record record) { context.forward(record); } } + + private static class WindowHeadersStoreWriterProcessor implements Processor { + private ProcessorContext context; + private TimestampedWindowStoreWithHeaders store; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + store = context.getStateStore(STORE_NAME); + } + + @Override + public void process(final Record record) { + // The record timestamp is the window start; a null value tombstones that window. The value's + // event-time is stored a few ms into the window (EVENT_TIME_OFFSET_MS) so tests can tell + // ReadOnlyRecord.timestamp() apart from the window start/end. + if (record.value() == null) { + store.put(record.key(), null, record.timestamp()); + } else { + store.put( + record.key(), + ValueTimestampHeaders.make(record.value(), record.timestamp() + EVENT_TIME_OFFSET_MS, record.headers()), + record.timestamp()); + } + context.forward(record); + } + } + + private static class WindowPlainStoreWriterProcessor implements Processor { + private ProcessorContext context; + private TimestampedWindowStore store; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + store = context.getStateStore(STORE_NAME); + } + + @Override + public void process(final Record record) { + store.put( + record.key(), + ValueAndTimestamp.make(record.value(), record.timestamp()), + record.timestamp()); + context.forward(record); + } + } } diff --git a/streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java new file mode 100644 index 0000000000000..c437f2ef6e301 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.streams.query; + +import org.apache.kafka.common.annotation.InterfaceAudience; +import org.apache.kafka.common.annotation.InterfaceStability.Evolving; +import org.apache.kafka.streams.kstream.Windowed; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; +import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; + +import java.time.Instant; +import java.util.Objects; +import java.util.Optional; + +/** + * Interactive query for retrieving records of a single key across a window-start range, including + * their record headers, from a {@link TimestampedWindowStoreWithHeaders}. + * + *

This is the headers-aware parallel of {@link WindowKeyQuery}: it returns a + * {@link ReadOnlyRecordIterator} of {@link ReadOnlyRecord} elements, each carrying the windowed key, + * value, timestamp, and headers, whereas {@link WindowKeyQuery} returns a + * {@link org.apache.kafka.streams.state.WindowStoreIterator} of values keyed by window-start + * timestamp (value only, no headers). Each element's key is a {@link Windowed} whose window + * describes the record's window; the record's stored event-time is exposed via + * {@link ReadOnlyRecord#timestamp()}. + * + *

As with {@link WindowKeyQuery}, a closed window-start range must be supplied (both + * {@code timeFrom} and {@code timeTo}); open-ended ranges are not supported. + * + *

Headers are persisted and returned only when the store is backed by a native headers store, + * i.e. built with a KIP-1271 {@code WithHeaders} byte-store supplier (e.g. + * {@code Stores.persistentTimestampedWindowStoreWithHeaders}). A {@code WithHeaders} store built + * over a legacy (non-headers) supplier cannot persist headers, so the store-served reads come back + * with empty {@code headers()}. + * + *

Against a plain store not built with the {@code WithHeaders} builder at all, this query type is + * unsupported and fails with {@link FailureReason#UNKNOWN_QUERY_TYPE}. + * + *

Each element is a {@link ReadOnlyRecord}, whose timestamp is non-negative. If the backing store + * does not persist timestamps -- for example a {@code WithHeaders} store built over a plain + * window-store supplier that surfaces entries with {@code NO_TIMESTAMP} ({@code -1}) -- that entry + * cannot be represented, so advancing the returned {@link ReadOnlyRecordIterator} throws + * {@link org.apache.kafka.streams.errors.StreamsException} at that entry. Back the store with + * {@code Stores.persistentTimestampedWindowStoreWithHeaders(...)} to persist timestamps and headers, + * or use {@link WindowKeyQuery} if headers are not needed. + * + *

Because iteration can throw mid-stream, always close the returned {@link ReadOnlyRecordIterator} + * (for example with try-with-resources), even when a call to {@code next()} throws; otherwise the + * underlying store iterator leaks and the {@code num-open-iterators} metric stays incremented. + * + * @param Type of keys + * @param Type of values + */ +@Evolving +@InterfaceAudience.Public +public final class TimestampedWindowKeyWithHeadersQuery implements Query, V>> { + + private final K key; + private final Optional timeFrom; + private final Optional timeTo; + + private TimestampedWindowKeyWithHeadersQuery(final K key, + final Optional timeFrom, + final Optional timeTo) { + this.key = key; + this.timeFrom = timeFrom; + this.timeTo = timeTo; + } + + /** + * Creates a query that will retrieve the records (value, timestamp, and headers) identified by + * {@code key} whose window start falls within the closed range {@code [timeFrom, timeTo]}. + * @param key The key to retrieve + * @param timeFrom The inclusive lower bound of the window-start range + * @param timeTo The inclusive upper bound of the window-start range + * @param The type of the key + * @param The type of the value that will be retrieved + */ + public static TimestampedWindowKeyWithHeadersQuery withKeyAndWindowStartRange(final K key, + final Instant timeFrom, + final Instant timeTo) { + Objects.requireNonNull(key, "the key should not be null"); + return new TimestampedWindowKeyWithHeadersQuery<>(key, Optional.of(timeFrom), Optional.of(timeTo)); + } + + /** + * The key that was specified for this query. + */ + public K key() { + return key; + } + + /** + * The inclusive lower bound of the window-start range, if specified. + */ + public Optional timeFrom() { + return timeFrom; + } + + /** + * The inclusive upper bound of the window-start range, if specified. + */ + public Optional timeTo() { + return timeTo; + } + + @Override + public String toString() { + return "TimestampedWindowKeyWithHeadersQuery{" + + "key=" + key + + ", timeFrom=" + timeFrom + + ", timeTo=" + timeTo + + '}'; + } +} diff --git a/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java b/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java index ab14cfc728529..9d37ebf38d0ae 100644 --- a/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java +++ b/streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java @@ -33,8 +33,8 @@ public class WindowKeyQuery implements Query> { private final Optional timeTo; private WindowKeyQuery(final K key, - final Optional timeTo, - final Optional timeFrom) { + final Optional timeFrom, + final Optional timeTo) { this.key = key; this.timeFrom = timeFrom; this.timeTo = timeTo; 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 eb02894fbf608..bb84cb542a281 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 @@ -24,8 +24,12 @@ import org.apache.kafka.common.utils.Time; import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.errors.ProcessorStateException; +import org.apache.kafka.streams.errors.StreamsException; import org.apache.kafka.streams.kstream.Windowed; +import org.apache.kafka.streams.kstream.internals.TimeWindow; import org.apache.kafka.streams.processor.StateStore; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.processor.api.Record; import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; import org.apache.kafka.streams.processor.internals.SerdeGetter; import org.apache.kafka.streams.query.FailureReason; @@ -33,10 +37,12 @@ import org.apache.kafka.streams.query.Query; import org.apache.kafka.streams.query.QueryConfig; import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery; import org.apache.kafka.streams.query.WindowKeyQuery; 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.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.ReadOnlyWindowStore; import org.apache.kafka.streams.state.TimestampedBytesStore; import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; @@ -142,6 +148,8 @@ public QueryResult query( if (query instanceof WindowKeyQuery) { result = runWindowKeyQuery((WindowKeyQuery>) query, positionBound, config); + } else if (query instanceof TimestampedWindowKeyWithHeadersQuery) { + result = runTimestampedWindowKeyWithHeadersQuery((TimestampedWindowKeyWithHeadersQuery) query, positionBound, config); } else if (query instanceof WindowRangeQuery) { result = runWindowRangeQuery((WindowRangeQuery>) query, positionBound, config); } else { @@ -216,6 +224,50 @@ private QueryResult runWindowKeyQuery( return queryResult; } + /** + * Handles {@link TimestampedWindowKeyWithHeadersQuery} by forwarding a raw byte-level + * {@link WindowKeyQuery} to the wrapped store and surfacing each entry as a {@link ReadOnlyRecord} + * (via {@link Record}) whose key is a {@link Windowed} of the queried key and the entry's window, + * carrying the stored value, event-time timestamp, and headers. + */ + @SuppressWarnings("unchecked") + private QueryResult runTimestampedWindowKeyWithHeadersQuery( + final TimestampedWindowKeyWithHeadersQuery query, + final PositionBound positionBound, + final QueryConfig config + ) { + final QueryResult queryResult; + + // Mirrors WindowKeyQuery handling: the closed-range contract (both bounds present) is guaranteed + // by the sole withKeyAndWindowStartRange factory today, but the Optional accessors are @Evolving, + // so we reject an open-ended range with UNKNOWN_QUERY_TYPE rather than assuming .get() is safe. + if (query.timeFrom().isPresent() && query.timeTo().isPresent()) { + final WindowKeyQuery rawKeyQuery = + WindowKeyQuery.withKeyAndWindowStartRange( + serializeKey(query.key(), internalContext.headers()), + query.timeFrom().get(), + query.timeTo().get() + ); + final QueryResult> rawResult = wrapped().query(rawKeyQuery, positionBound, config); + if (rawResult.isSuccess()) { + final ReadOnlyRecordIterator, V> resultIterator = + new MeteredWindowStoreWithHeadersReadOnlyRecordIterator(rawResult.getResult(), query.key()); + queryResult = (QueryResult) InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, resultIterator); + } else { + queryResult = (QueryResult) rawResult; + } + } else { + queryResult = QueryResult.forFailure( + FailureReason.UNKNOWN_QUERY_TYPE, + "This store (" + getClass() + ") doesn't know how to execute" + + " the given query (" + query + ") because it only supports closed-range" + + " queries." + + " Contact the store maintainer if you need support for a new query type." + ); + } + return queryResult; + } + private MeteredWindowStoreIterator meteredIterator( final QueryResult> rawResult, final Function valueDeserializer @@ -519,6 +571,92 @@ public Windowed peekNextKey() { } } + /** + * Iterator backing {@link TimestampedWindowKeyWithHeadersQuery}: yields each entry as a + * {@link ReadOnlyRecord} (implemented by {@link Record}) whose key is a {@link Windowed} of the + * queried key and the entry's window, carrying the value, stored event-time timestamp, and the + * stored headers, with the headers frozen so a caller cannot mutate the read-only result. + * + *

A {@link ReadOnlyRecord} timestamp is contractually non-negative, so an entry with a + * negative stored timestamp cannot be represented -- for example a {@code WithHeaders} store + * built over a plain, non-timestamped window supplier surfaces entries with + * {@code NO_TIMESTAMP} (-1). This mirrors the rule the point/range headers queries apply. But + * because a lazily-evaluated iterator has already returned a successful {@link QueryResult} + * before any entry is read, such an entry cannot be surfaced as a query-level failure; it is + * instead reported by throwing a {@link StreamsException} while advancing the iterator. + * + *

Because {@link #next()} can throw mid-iteration, the caller must always close this iterator + * (via try-with-resources or a {@code finally} block) even when iteration throws: a caller that + * catches the exception and abandons the iterator without closing it leaks the underlying store + * iterator and leaves the {@code num-open-iterators} metric permanently incremented. + */ + private class MeteredWindowStoreWithHeadersReadOnlyRecordIterator + implements ReadOnlyRecordIterator, V>, MeteredIterator { + + private final WindowStoreIterator iter; + private final K key; + private final long startNs; + private final long startTimestampMs; + + private MeteredWindowStoreWithHeadersReadOnlyRecordIterator( + final WindowStoreIterator iter, + final K key + ) { + this.iter = iter; + this.key = key; + this.startNs = time.nanoseconds(); + this.startTimestampMs = time.milliseconds(); + numOpenIterators.increment(); + openIterators.add(this); + } + + @Override + public long startTimestamp() { + return startTimestampMs; + } + + @Override + public boolean hasNext() { + return iter.hasNext(); + } + + @Override + public ReadOnlyRecord, V> next() { + final KeyValue next = iter.next(); + final long windowStart = next.key; + final ValueTimestampHeaders valueTimestampHeaders = deserializeValue(next.value); + final Headers headers = valueTimestampHeaders.headers(); + if (valueTimestampHeaders.timestamp() < 0) { + throw new StreamsException( + "Cannot represent the stored record for key [" + key + "] and window start [" + + windowStart + "] as a ReadOnlyRecord: its timestamp (" + + valueTimestampHeaders.timestamp() + ") is negative."); + } + final Windowed windowedKey = + new Windowed<>(key, new TimeWindow(windowStart, windowStart + windowSizeMs)); + final Record, V> record = new Record<>( + windowedKey, + valueTimestampHeaders.value(), + valueTimestampHeaders.timestamp(), + headers); + ((RecordHeaders) record.headers()).setReadOnly(); + return record; + } + + @Override + public void close() { + try { + iter.close(); + } finally { + final long duration = time.nanoseconds() - startNs; + fetchSensor.record(duration); + iteratorDurationSensor.record(duration); + numOpenIterators.decrement(); + openIterators.remove(this); + } + } + } + private MeteredWindowedKeyValueIterator meteredWindowedIterator( final QueryResult, byte[]>> rawResult, final Function, ValueType> valueConverter diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredWindowStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredWindowStore.java index 31c5dbfc0cfa4..fdf4b6103553d 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredWindowStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredWindowStore.java @@ -67,7 +67,7 @@ public class MeteredWindowStore extends WrappedStateStore, Windowed, V> implements WindowStore, MeteredStateStore { - private final long windowSizeMs; + protected final long windowSizeMs; private final String metricsScope; protected final Time time; private final Serde keySerde; diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java index 8d4d5a23b4419..08a8f1771b774 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeaders.java @@ -16,12 +16,12 @@ */ package org.apache.kafka.streams.state.internals; -import org.apache.kafka.streams.processor.StateStore; import org.apache.kafka.streams.query.Position; import org.apache.kafka.streams.query.PositionBound; import org.apache.kafka.streams.query.Query; import org.apache.kafka.streams.query.QueryConfig; import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.WindowKeyQuery; import org.apache.kafka.streams.state.HeadersBytesStore; import org.apache.kafka.streams.state.TimestampedBytesStore; @@ -32,9 +32,6 @@ * {@link TimestampedBytesStore} (for timestamp support) and {@link HeadersBytesStore} * (for header support) marker interfaces. *

- * This store returns {@link QueryResult#forUnknownQueryType(Query, StateStore)} for all queries, - * as IQv2 query handling is done at the metered layer. - *

* The storage format for values is: [headersSize(varint)][headersBytes][timestamp(8)][value] * * @see RocksDBTimeOrderedWindowStore @@ -53,6 +50,15 @@ class RocksDBTimeOrderedWindowStoreWithHeaders extends RocksDBTimeOrderedWindowS public QueryResult query(final Query query, final PositionBound positionBound, final QueryConfig config) { + // KIP-1356 (window key query only): the headers-aware TimestampedWindowKeyWithHeadersQuery + // forwards a raw WindowKeyQuery to this native store, so enable WindowKeyQuery via the inherited + // RocksDBTimeOrderedWindowStore handling (StoreQueryUtils). Every other query type (e.g. + // WindowRangeQuery) stays UNKNOWN_QUERY_TYPE here, unchanged from before; those are enabled by + // their own KIP-1356 follow-ups. + if (query instanceof WindowKeyQuery) { + return super.query(query, positionBound, config); + } + final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L; final QueryResult result; final Position position = getPosition(); diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java index fa22aa5090780..304975279f307 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeaders.java @@ -21,6 +21,7 @@ import org.apache.kafka.streams.query.Query; import org.apache.kafka.streams.query.QueryConfig; import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.WindowKeyQuery; import org.apache.kafka.streams.state.HeadersBytesStore; import org.apache.kafka.streams.state.TimestampedBytesStore; @@ -57,6 +58,15 @@ class RocksDBTimestampedWindowStoreWithHeaders extends RocksDBWindowStore implem public QueryResult query(final Query query, final PositionBound positionBound, final QueryConfig config) { + // KIP-1356 (window key query only): the headers-aware TimestampedWindowKeyWithHeadersQuery + // forwards a raw WindowKeyQuery to this native store, so enable WindowKeyQuery via the inherited + // RocksDBWindowStore handling (StoreQueryUtils). Every other query type (e.g. WindowRangeQuery) + // stays UNKNOWN_QUERY_TYPE here, unchanged from before; those are enabled by their own KIP-1356 + // follow-ups. + if (query instanceof WindowKeyQuery) { + return super.query(query, positionBound, config); + } + final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L; final QueryResult result; final Position position = getPosition(); 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 1f0dcbb1fdcd3..3aa0a1b87ff44 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 @@ -18,6 +18,7 @@ import org.apache.kafka.common.IsolationLevel; import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.metrics.KafkaMetric; import org.apache.kafka.common.metrics.MetricConfig; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.metrics.Sensor; @@ -39,10 +40,19 @@ import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; import org.apache.kafka.streams.processor.internals.ProcessorStateManager; import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl; +import org.apache.kafka.streams.query.FailureReason; +import org.apache.kafka.streams.query.PositionBound; +import org.apache.kafka.streams.query.Query; +import org.apache.kafka.streams.query.QueryConfig; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery; +import org.apache.kafka.streams.query.WindowKeyQuery; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; 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.streams.state.WindowStoreIterator; import org.apache.kafka.test.InternalMockProcessorContext; import org.apache.kafka.test.KeyValueIteratorStub; import org.apache.kafka.test.MockRecordCollector; @@ -51,17 +61,21 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; import java.time.Instant; +import java.util.Iterator; import java.util.List; +import java.util.Optional; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; @@ -516,4 +530,158 @@ public void shouldUseHeadersFromValueToDeserializeKeyInReadOnlyFetchAll() { verify(keyDeserializer).deserialize(any(), eq(HEADERS), eq(KEY.getBytes())); } + + @Test + public void shouldPropagateWindowKeyQueryBoundsForTimestampedWindowKeyWithHeadersQuery() { + setUp(); + store.init(context, store); + + final Instant timeFrom = Instant.ofEpochMilli(5); + final Instant timeTo = Instant.ofEpochMilli(100); + final WindowKeyQuery rawQuery = forwardedRawWindowKeyQuery( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange(KEY, timeFrom, timeTo)); + + // The typed query is translated into a raw byte-level WindowKeyQuery with the serialized key + // and the same window-start range, forwarded to the wrapped store. + assertEquals(KEY_BYTES, rawQuery.getKey()); + assertEquals(Optional.of(timeFrom), rawQuery.getTimeFrom()); + assertEquals(Optional.of(timeTo), rawQuery.getTimeTo()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + @Test + public void shouldPropagateWrappedStoreFailureForTimestampedWindowKeyWithHeadersQuery() { + setUp(); + store.init(context, store); + when(innerStoreMock.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forFailure(FailureReason.STORE_EXCEPTION, "boom")); + + final QueryResult, String>> result = store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + KEY, Instant.ofEpochMilli(5), Instant.ofEpochMilli(100)), + PositionBound.unbounded(), + new QueryConfig(false)); + + assertFalse(result.isSuccess()); + assertEquals(FailureReason.STORE_EXCEPTION, result.getFailureReason()); + assertEquals("boom", result.getFailureMessage()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + @Test + public void shouldTrackNumOpenIteratorsForTimestampedWindowKeyWithHeadersQuery() { + setUp(); + store.init(context, store); + when(innerStoreMock.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forResult(windowStoreIterator(List.of()))); + + final KafkaMetric openIterators = numOpenIteratorsMetric(); + assertEquals(0L, (Long) openIterators.metricValue()); + + final QueryResult, String>> result = store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + KEY, Instant.ofEpochMilli(5), Instant.ofEpochMilli(100)), + PositionBound.unbounded(), + new QueryConfig(false)); + assertTrue(result.isSuccess()); + + // The query's ReadOnlyRecordIterator registers itself on open and deregisters on close. + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertEquals(1L, (Long) openIterators.metricValue()); + } + assertEquals(0L, (Long) openIterators.metricValue()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + @Test + public void shouldDecrementOpenIteratorsTwiceWhenClosedTwiceForTimestampedWindowKeyWithHeadersQuery() { + setUp(); + store.init(context, store); + when(innerStoreMock.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forResult(windowStoreIterator(List.of()))); + + final KafkaMetric openIterators = numOpenIteratorsMetric(); + final ReadOnlyRecordIterator, String> iterator = store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + KEY, Instant.ofEpochMilli(5), Instant.ofEpochMilli(100)), + PositionBound.unbounded(), + new QueryConfig(false)).getResult(); + + assertEquals(1L, (Long) openIterators.metricValue()); + iterator.close(); + assertEquals(0L, (Long) openIterators.metricValue()); + // close() is intentionally not idempotent (matching the sibling metered iterators): each call + // decrements, so a repeated close drives the gauge below zero. Callers must close exactly once. + iterator.close(); + assertEquals(-1L, (Long) openIterators.metricValue()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + @Test + public void shouldLeaveIteratorOpenWhenNextThrowsAndNotClosed() { + setUp(); + store.init(context, store); + // A stored entry with a negative timestamp cannot be represented as a ReadOnlyRecord, so next() throws. + final byte[] negativeTimestampBytes = new ValueTimestampHeadersSerializer<>(new StringSerializer()) + .serialize("topic", ValueTimestampHeaders.make("value", -1L, HEADERS)); + when(innerStoreMock.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forResult( + windowStoreIterator(List.of(KeyValue.pair(5L, negativeTimestampBytes))))); + + final KafkaMetric openIterators = numOpenIteratorsMetric(); + final ReadOnlyRecordIterator, String> iterator = store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + KEY, Instant.ofEpochMilli(5), Instant.ofEpochMilli(100)), + PositionBound.unbounded(), + new QueryConfig(false)).getResult(); + + assertEquals(1L, (Long) openIterators.metricValue()); + // next() throws on the negative-timestamp entry but does not close the iterator, so it stays + // registered (num-open-iterators stays incremented) until the caller closes it. + assertThrows(StreamsException.class, iterator::next); + assertEquals(1L, (Long) openIterators.metricValue()); + iterator.close(); + assertEquals(0L, (Long) openIterators.metricValue()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private WindowKeyQuery forwardedRawWindowKeyQuery(final Query query) { + when(innerStoreMock.query(any(), any(PositionBound.class), any(QueryConfig.class))) + .thenReturn((QueryResult) QueryResult.forResult(windowStoreIterator(List.of()))); + store.query(query, PositionBound.unbounded(), new QueryConfig(false)); + final ArgumentCaptor captor = ArgumentCaptor.forClass(WindowKeyQuery.class); + verify(innerStoreMock).query(captor.capture(), any(PositionBound.class), any(QueryConfig.class)); + return captor.getValue(); + } + + private static WindowStoreIterator windowStoreIterator(final List> data) { + final Iterator> iterator = data.iterator(); + return new WindowStoreIterator<>() { + @Override + public void close() { } + + @Override + public Long peekNextKey() { + throw new UnsupportedOperationException(); + } + + @Override + public boolean hasNext() { + return iterator.hasNext(); + } + + @Override + public KeyValue next() { + return iterator.next(); + } + }; + } + + private KafkaMetric numOpenIteratorsMetric() { + return metrics.metrics().entrySet().stream() + .filter(entry -> entry.getKey().name().equals("num-open-iterators")) + .findFirst() + .orElseThrow(() -> new AssertionError("num-open-iterators metric not registered")) + .getValue(); + } } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java index 13994dd36af45..6b997c668b035 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStoreWithHeadersTest.java @@ -18,6 +18,7 @@ import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.query.FailureReason; import org.apache.kafka.streams.query.PositionBound; @@ -39,6 +40,7 @@ import java.time.Instant; import java.util.Properties; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -87,17 +89,32 @@ public void tearDown() { } @Test - public void shouldReturnUnknownQueryTypeForWindowKeyQuery() { + public void shouldHandleWindowKeyQuery() { + // KIP-1356 (window key query only): the native time-ordered window header store special-cases + // WindowKeyQuery to the inherited RocksDBTimeOrderedWindowStore handling (StoreQueryUtils), + // returning the raw stored header-format bytes (previously UNKNOWN_QUERY_TYPE); the metered store + // does the header-aware deserialization. This matches the adapter build path. + final Bytes key = new Bytes("test-key".getBytes()); + final byte[] storedBytes = "headers+timestamp+value".getBytes(); + final long windowStart = 1_000L; + windowStore.put(key, storedBytes, windowStart); + final WindowKeyQuery query = WindowKeyQuery.withKeyAndWindowStartRange( - new Bytes("test-key".getBytes()), + key, Instant.ofEpochMilli(0), - Instant.ofEpochMilli(Long.MAX_VALUE) + Instant.ofEpochMilli(RETENTION_PERIOD) ); final QueryResult> result = windowStore.query(query, PositionBound.unbounded(), new QueryConfig(false)); - assertFalse(result.isSuccess()); - assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getFailureReason()); + assertTrue(result.isSuccess(), "Expected WindowKeyQuery to succeed"); + try (WindowStoreIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext(), "Expected the stored entry in the window key result"); + final KeyValue keyValue = iterator.next(); + assertEquals(windowStart, keyValue.key); + assertArrayEquals(storedBytes, keyValue.value, "Expected the raw stored bytes to be returned"); + assertFalse(iterator.hasNext(), "Expected exactly one entry in the window key result"); + } assertNotNull(result.getPosition()); } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java index 04ee23555a70b..6d77fb3dd0e37 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedWindowStoreWithHeadersTest.java @@ -18,6 +18,7 @@ import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.query.FailureReason; import org.apache.kafka.streams.query.PositionBound; @@ -39,6 +40,7 @@ import java.time.Instant; import java.util.Properties; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -91,24 +93,32 @@ public void tearDown() { } @Test - public void shouldReturnUnknownQueryTypeForWindowKeyQuery() { + public void shouldHandleWindowKeyQuery() { + // KIP-1356 (window key query only): the native window header store special-cases WindowKeyQuery + // to the inherited RocksDBWindowStore handling (StoreQueryUtils), returning the raw stored + // header-format bytes (previously UNKNOWN_QUERY_TYPE); the metered store does the header-aware + // deserialization. This matches the adapter build path. + final Bytes key = new Bytes("test-key".getBytes()); + final byte[] storedBytes = "headers+timestamp+value".getBytes(); + final long windowStart = 1_000L; + windowStore.put(key, storedBytes, windowStart); + final WindowKeyQuery query = WindowKeyQuery.withKeyAndWindowStartRange( - new Bytes("test-key".getBytes()), + key, Instant.ofEpochMilli(0), - Instant.ofEpochMilli(Long.MAX_VALUE) - ); - final PositionBound positionBound = PositionBound.unbounded(); - final QueryConfig config = new QueryConfig(false); - - final QueryResult> result = windowStore.query(query, positionBound, config); - - // Verify: Window store with headers currently returns UNKNOWN_QUERY_TYPE - assertFalse(result.isSuccess(), "Expected query to fail with unknown query type"); - assertEquals( - FailureReason.UNKNOWN_QUERY_TYPE, - result.getFailureReason(), - "Expected UNKNOWN_QUERY_TYPE failure reason" + Instant.ofEpochMilli(RETENTION_PERIOD) ); + final QueryResult> result = + windowStore.query(query, PositionBound.unbounded(), new QueryConfig(false)); + + assertTrue(result.isSuccess(), "Expected WindowKeyQuery to succeed"); + try (WindowStoreIterator iterator = result.getResult()) { + assertTrue(iterator.hasNext(), "Expected the stored entry in the window key result"); + final KeyValue keyValue = iterator.next(); + assertEquals(windowStart, keyValue.key); + assertArrayEquals(storedBytes, keyValue.value, "Expected the raw stored bytes to be returned"); + assertFalse(iterator.hasNext(), "Expected exactly one entry in the window key result"); + } assertNotNull(result.getPosition(), "Expected position to be set"); } diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java index 67216bca41157..093d5c8022f75 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilderTest.java @@ -17,31 +17,67 @@ package org.apache.kafka.streams.state.internals; +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.serialization.Serdes; +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.errors.StreamsException; +import org.apache.kafka.streams.kstream.Windowed; import org.apache.kafka.streams.processor.StateStore; +import org.apache.kafka.streams.processor.api.ReadOnlyRecord; +import org.apache.kafka.streams.processor.internals.MockStreamsMetrics; +import org.apache.kafka.streams.processor.internals.ProcessorRecordContext; +import org.apache.kafka.streams.query.PositionBound; +import org.apache.kafka.streams.query.QueryConfig; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery; +import org.apache.kafka.streams.query.WindowKeyQuery; +import org.apache.kafka.streams.state.ReadOnlyRecordIterator; import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders; +import org.apache.kafka.streams.state.ValueAndTimestamp; +import org.apache.kafka.streams.state.ValueTimestampHeaders; import org.apache.kafka.streams.state.WindowBytesStoreSupplier; +import org.apache.kafka.streams.state.WindowStore; +import org.apache.kafka.streams.state.WindowStoreIterator; +import org.apache.kafka.test.InternalMockProcessorContext; +import org.apache.kafka.test.TestUtils; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.STRICT_STUBS) public class TimestampedWindowStoreWithHeadersBuilderTest { + private enum StoreType { NATIVE, ADAPTER, PLAIN_ADAPTER, IN_MEMORY } + @Nested class BuilderTests { private static final String STORE_NAME = "name"; @@ -207,4 +243,290 @@ public void shouldThrowNullPointerIfInnerIsNull() { assertThrows(NullPointerException.class, () -> new TimestampedWindowStoreWithHeadersBuilder<>(null, Serdes.String(), Serdes.String(), new MockTime())); } } + + /** + * End-to-end query behavior of {@link TimestampedWindowKeyWithHeadersQuery} against a real store + * built by the builder over each supported build path. Caching is disabled: a window key query reads + * the underlying store directly (it never consults the cache), so writes must be store-served. + */ + @Nested + class QueryTests { + private static final String STORE_NAME = "test-window-store"; + private static final String METRICS_SCOPE = "metrics-scope"; + private static final long RETENTION = 60_000L; + private static final long SEGMENT_INTERVAL = 30_000L; + private static final long WINDOW_SIZE = 10_000L; + private static final long WINDOW_START = 1_000L; + + @Mock + private WindowBytesStoreSupplier supplier; + + private WindowStore innerStore(final StoreType storeType) { + switch (storeType) { + case NATIVE: + return new RocksDBTimestampedWindowStoreWithHeaders( + new RocksDBTimestampedSegmentedBytesStoreWithHeaders( + STORE_NAME, METRICS_SCOPE, RETENTION, SEGMENT_INTERVAL, new WindowKeySchema()), + false, WINDOW_SIZE); + case ADAPTER: + return new RocksDBTimestampedWindowStore( + new RocksDBTimestampedSegmentedBytesStore( + STORE_NAME, METRICS_SCOPE, RETENTION, SEGMENT_INTERVAL, new WindowKeySchema()), + false, WINDOW_SIZE); + case PLAIN_ADAPTER: + return new RocksDBWindowStore( + new RocksDBSegmentedBytesStore( + STORE_NAME, METRICS_SCOPE, RETENTION, SEGMENT_INTERVAL, new WindowKeySchema()), + false, WINDOW_SIZE); + case IN_MEMORY: + // Non-persistent supplier: the builder wraps it with the in-memory headers marker, + // which (like the native store) persists the value-with-headers byte format. + return new InMemoryWindowStore(STORE_NAME, RETENTION, WINDOW_SIZE, false, METRICS_SCOPE); + default: + throw new IllegalArgumentException("unknown store type: " + storeType); + } + } + + private TimestampedWindowStoreWithHeaders buildAndInitStore(final StoreType storeType) { + lenient().when(supplier.name()).thenReturn(STORE_NAME); + lenient().when(supplier.metricsScope()).thenReturn(METRICS_SCOPE); + lenient().when(supplier.windowSize()).thenReturn(WINDOW_SIZE); + lenient().when(supplier.get()).thenReturn(innerStore(storeType)); + + final TimestampedWindowStoreWithHeaders store = + new TimestampedWindowStoreWithHeadersBuilder<>(supplier, Serdes.String(), Serdes.String(), new MockTime()) + .withLoggingDisabled() + .withCachingDisabled() + .build(); + + final ThreadCache cache = new ThreadCache(new LogContext("test "), 0, new MockStreamsMetrics(new Metrics())); + final InternalMockProcessorContext context = new InternalMockProcessorContext<>( + TestUtils.tempDirectory(), Serdes.String(), Serdes.String(), null, cache); + context.setRecordContext(new ProcessorRecordContext(0L, 0L, 0, "topic", new RecordHeaders())); + store.init(context, store); + return store; + } + + private Headers headersWith(final String key, final String value) { + return new RecordHeaders().add(key, value.getBytes(StandardCharsets.UTF_8)); + } + + private QueryResult, String>> windowKeyQuery( + final TimestampedWindowStoreWithHeaders store) { + return store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + "k", Instant.ofEpochMilli(0), Instant.ofEpochMilli(RETENTION)), + PositionBound.unbounded(), + new QueryConfig(false)); + } + + @ParameterizedTest + @CsvSource({"NATIVE", "IN_MEMORY"}) + public void shouldReturnHeadersForTimestampedWindowKeyWithHeadersQueryOnHeaderPersistingStore(final StoreType storeType) { + // The native and in-memory builds both persist headers: the native store keeps them in its + // headers column family, and the in-memory marker stores the value bytes verbatim (the metered + // layer serializes the headers into those bytes). The adapter build is the one that drops them. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(storeType); + try { + final Headers headers = headersWith("h", "x"); + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headers), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess(), "Expected TimestampedWindowKeyWithHeadersQuery to succeed"); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertTrue(iterator.hasNext()); + final ReadOnlyRecord, String> record = iterator.next(); + assertEquals("k", record.key().key()); + assertEquals(WINDOW_START, record.key().window().start()); + assertEquals(WINDOW_START + WINDOW_SIZE, record.key().window().end()); + assertEquals("v", record.value()); + assertEquals(1_005L, record.timestamp()); + assertEquals(headers, record.headers()); + // The IQ result is a read-only snapshot: its headers are immutable (neither add nor remove). + assertThrows(IllegalStateException.class, () -> record.headers().add("new", new byte[0])); + assertThrows(IllegalStateException.class, () -> record.headers().remove("h")); + assertFalse(iterator.hasNext()); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @Test + public void shouldReturnEmptyHeadersForTimestampedWindowKeyWithHeadersQueryOnAdapterStore() { + // The timestamped adapter keeps the timestamp but drops headers on write, so the record comes + // back with empty (never null) headers while value and timestamp still round-trip. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.ADAPTER); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertTrue(iterator.hasNext()); + final ReadOnlyRecord, String> record = iterator.next(); + assertEquals("k", record.key().key()); + assertEquals(WINDOW_START, record.key().window().start()); + assertEquals("v", record.value()); + assertEquals(1_005L, record.timestamp()); + assertEquals(new RecordHeaders(), record.headers()); + // Even the empty headers are a read-only snapshot: neither add nor remove is allowed. + assertThrows(IllegalStateException.class, () -> record.headers().add("new", new byte[0])); + assertThrows(IllegalStateException.class, () -> record.headers().remove("h")); + assertFalse(iterator.hasNext()); + } + assertNotNull(result.getPosition(), "Expected position to be set"); + } finally { + store.close(); + } + } + + @Test + public void shouldThrowForTimestampedWindowKeyWithHeadersQueryOnPlainSupplier() { + // A plain (non-timestamped) window supplier surfaces every entry with timestamp = -1, which + // cannot be represented as a ReadOnlyRecord. The query succeeds, but iterating throws. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.PLAIN_ADAPTER); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess(), "The query itself succeeds; the failure surfaces while iterating"); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + final StreamsException e = assertThrows(StreamsException.class, iterator::next, + "An entry with ts=-1 cannot be represented as a ReadOnlyRecord"); + assertTrue(e.getMessage().contains("as a ReadOnlyRecord") && e.getMessage().contains("is negative"), + "unexpected message: " + e.getMessage()); + } + } finally { + store.close(); + } + } + + @Test + public void shouldThrowForNegativeStoredTimestampForTimestampedWindowKeyWithHeadersQuery() { + // A caller can store a negative timestamp directly on a native (header-persisting) store. + // The query succeeds, but the lazily-evaluated iterator throws while advancing (a ReadOnlyRecord + // timestamp must be non-negative). + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + store.put("k", ValueTimestampHeaders.make("v", -1L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + final StreamsException e = assertThrows(StreamsException.class, iterator::next); + assertTrue(e.getMessage().contains("as a ReadOnlyRecord") && e.getMessage().contains("is negative"), + "unexpected message: " + e.getMessage()); + } + } finally { + store.close(); + } + } + + @Test + public void shouldCollectExecutionInfoForTimestampedWindowKeyWithHeadersQueryWhenRequested() { + // With execution info enabled, the result must carry both the wrapped (native) store's entry + // and the metered handler's entry. + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = store.query( + TimestampedWindowKeyWithHeadersQuery.withKeyAndWindowStartRange( + "k", Instant.ofEpochMilli(0), Instant.ofEpochMilli(RETENTION)), + PositionBound.unbounded(), + new QueryConfig(true)); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + final String info = String.join("\n", result.getExecutionInfo()); + assertTrue( + info.contains(RocksDBTimestampedWindowStoreWithHeaders.class.getName()) + && info.contains(MeteredTimestampedWindowStoreWithHeaders.class.getName()), + "execution info missing an entry: " + info); + } + } finally { + store.close(); + } + } + + @Test + public void shouldNotCollectExecutionInfoForTimestampedWindowKeyWithHeadersQueryWhenNotRequested() { + final TimestampedWindowStoreWithHeaders store = buildAndInitStore(StoreType.NATIVE); + try { + store.put("k", ValueTimestampHeaders.make("v", 1_005L, headersWith("h", "x")), WINDOW_START); + + final QueryResult, String>> result = windowKeyQuery(store); + + assertTrue(result.isSuccess()); + try (ReadOnlyRecordIterator, String> iterator = result.getResult()) { + assertTrue(result.getExecutionInfo().isEmpty(), "Expected no execution info: " + result.getExecutionInfo()); + } + } finally { + store.close(); + } + } + + @Test + public void shouldReturnIdenticalResultsForNativeAndAdapterBuiltStores() { + // Build-path parity for the existing (header-stripped) WindowKeyQuery: KIP-1356 makes the native + // store serve it exactly as the adapter build already did. + final ValueTimestampHeaders v1 = ValueTimestampHeaders.make("v1", 1_005L, headersWith("h", "1")); + final ValueTimestampHeaders v2 = ValueTimestampHeaders.make("v2", 3_005L, headersWith("h", "2")); + final ValueTimestampHeaders v3 = ValueTimestampHeaders.make("v3", 5_005L, headersWith("h", "3")); + final ValueTimestampHeaders o1 = ValueTimestampHeaders.make("o1", 2_005L, headersWith("h", "o")); + final ValueTimestampHeaders o2 = ValueTimestampHeaders.make("o2", 4_005L, headersWith("h", "o")); + + final TimestampedWindowStoreWithHeaders nativeStore = buildAndInitStore(StoreType.NATIVE); + final TimestampedWindowStoreWithHeaders adapterStore = buildAndInitStore(StoreType.ADAPTER); + try { + // ascending window-start insertion order, interleaved with a different key at shared windows + nativeStore.put("other", o1, WINDOW_START); + nativeStore.put("k", v1, WINDOW_START); + nativeStore.put("k", v2, WINDOW_START + 2_000L); + nativeStore.put("other", o2, WINDOW_START + 2_000L); + nativeStore.put("k", v3, WINDOW_START + 4_000L); + // descending window-start insertion order (same windows, reversed), plus the other key + adapterStore.put("k", v3, WINDOW_START + 4_000L); + adapterStore.put("other", o2, WINDOW_START + 2_000L); + adapterStore.put("k", v2, WINDOW_START + 2_000L); + adapterStore.put("k", v1, WINDOW_START); + adapterStore.put("other", o1, WINDOW_START); + + final List>> expected = List.of( + KeyValue.pair(WINDOW_START, ValueAndTimestamp.make("v1", 1_005L)), + KeyValue.pair(WINDOW_START + 2_000L, ValueAndTimestamp.make("v2", 3_005L)), + KeyValue.pair(WINDOW_START + 4_000L, ValueAndTimestamp.make("v3", 5_005L))); + + assertEquals(expected, plainWindowKeyResults(nativeStore), + "native WindowKeyQuery should return all windows in window-start order"); + assertEquals(expected, plainWindowKeyResults(adapterStore), + "adapter WindowKeyQuery should return all windows in window-start order"); + } finally { + nativeStore.close(); + adapterStore.close(); + } + } + + // Drains the (header-stripped) plain WindowKeyQuery, which yields the window-start timestamp keyed + // to a ValueAndTimestamp -- used to compare native and adapter build paths. + private List>> plainWindowKeyResults( + final TimestampedWindowStoreWithHeaders store) { + final WindowKeyQuery> query = + WindowKeyQuery.withKeyAndWindowStartRange("k", Instant.ofEpochMilli(0), Instant.ofEpochMilli(RETENTION)); + final List>> out = new ArrayList<>(); + try (WindowStoreIterator> iterator = + store.query(query, PositionBound.unbounded(), new QueryConfig(false)).getResult()) { + while (iterator.hasNext()) { + out.add(iterator.next()); + } + } + return out; + } + } }