aliehsaeedii commented on code in PR #22583: URL: https://github.com/apache/kafka/pull/22583#discussion_r3420494313
########## streams/src/main/java/org/apache/kafka/streams/query/TimestampedKeyWithHeadersQuery.java: ########## @@ -0,0 +1,90 @@ +/* + * 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.InterfaceStability.Evolving; +import org.apache.kafka.streams.state.TimestampedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.ValueTimestampHeaders; + +import java.util.Objects; + +/** + * Interactive query for retrieving a single record, including its record headers, based on its key + * from a {@link TimestampedKeyValueStoreWithHeaders}. + * + * <p>This is the headers-aware parallel of {@link TimestampedKeyQuery}: it returns a + * {@link ValueTimestampHeaders} (value, timestamp, and headers) instead of a + * {@link org.apache.kafka.streams.state.ValueAndTimestamp} (value and timestamp only). + * + * <p>Headers are only returned when the queried store was built with a KIP-1271 + * {@code WithHeaders} supplier. Against a plain (non-headers) store, this query type is + * unsupported and fails with {@link FailureReason#UNKNOWN_QUERY_TYPE}. + * + * @param <K> Type of keys + * @param <V> Type of values + */ +@Evolving +public final class TimestampedKeyWithHeadersQuery<K, V> implements Query<ValueTimestampHeaders<V>> { + + private final K key; + private final boolean skipCache; + + private TimestampedKeyWithHeadersQuery(final K key, final boolean skipCache) { + this.key = key; + this.skipCache = skipCache; + } + + /** + * Creates a query that will retrieve the record (value, timestamp, and headers) identified by + * {@code key} if it exists (or {@code null} otherwise). + * @param key The key to retrieve + * @param <K> The type of the key + * @param <V> The type of the value that will be retrieved + */ + public static <K, V> TimestampedKeyWithHeadersQuery<K, V> withKey(final K key) { + Objects.requireNonNull(key, "the key should not be null"); + return new TimestampedKeyWithHeadersQuery<>(key, false); + } + + /** + * Specifies that the cache should be skipped during query evaluation. This means, that the query will always + * get forwarded to the underlying store. + * + * <p><b>PoC limitation:</b> this flag is currently a no-op. The header-aware store handler does not yet + * propagate it to the underlying cache, so queries are evaluated against the cache regardless. Wiring + * this through is out of scope for the proof-of-concept. + */ + public TimestampedKeyWithHeadersQuery<K, V> skipCache() { Review Comment: Then do we need to have this at all? Just to be consitent with other query types? or if we added caching to headers-aware stores? If keeping this, we may need to add tests for that as well. ########## streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreWithHeaders.java: ########## @@ -199,6 +200,22 @@ private void openFromTimestampedStore(final DBOptions dbOptions, public <R> QueryResult<R> query(final Query<R> query, final PositionBound positionBound, final QueryConfig config) { + // PoC (KIP-1356): serve raw KeyQuery from the persistent store so that the header-aware + // IQv2 query forwarded by MeteredTimestampedKeyValueStoreWithHeaders keeps working after the + // record cache has been flushed (a cache miss falls through to this layer). The returned bytes + // are the serialized ValueTimestampHeaders; the metered store performs the header-aware + // deserialization. Other query types remain unsupported for now. + if (query instanceof KeyQuery) { Review Comment: This new `KeyQuery` branch returns before entering the `synchronized (position)` block that the rest of `query()` uses to take a consistent `position.copy()`. `handleBasicQueries` reads `position` without holding that lock, so a concurrent position update during the query is now possible for this path where the other query types are protected. Worth confirming whether that's intentional or whether the branch should run inside the synchronized block for consistency with the existing code. ########## streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/TimestampedKeyWithHeadersQueryIntegrationTest.java: ########## @@ -0,0 +1,282 @@ +/* + * 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.integration; + +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.serialization.IntegerDeserializer; +import org.apache.kafka.common.serialization.IntegerSerializer; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.common.serialization.StringDeserializer; +import org.apache.kafka.common.serialization.StringSerializer; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.StreamsBuilder; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster; +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.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.query.FailureReason; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.StateQueryRequest; +import org.apache.kafka.streams.query.StateQueryResult; +import org.apache.kafka.streams.query.TimestampedKeyWithHeadersQuery; +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.ValueAndTimestamp; +import org.apache.kafka.streams.state.ValueTimestampHeaders; +import org.apache.kafka.test.TestUtils; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; + +import java.io.IOException; +import java.time.Duration; +import java.util.Arrays; +import java.util.Properties; + +import static org.apache.kafka.streams.query.StateQueryRequest.inStore; +import static org.apache.kafka.streams.utils.TestUtils.safeUniqueTestName; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * End-to-end PoC test for KIP-1356 {@link TimestampedKeyWithHeadersQuery}. + * + * <p>Builds a KIP-1271 {@code WithHeaders} store, writes records (with headers) into it through a + * processor, then queries the store through IQv2 and asserts that the returned + * {@link ValueTimestampHeaders} carries value, timestamp, and the exact headers written. + */ +@Tag("integration") +public class TimestampedKeyWithHeadersQueryIntegrationTest { + + private static final String STORE_NAME = "headers-store"; + + private static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(1); + + private static final Headers HEADERS = new RecordHeaders() + .add("source", "test".getBytes()) + .add("version", "1.0".getBytes()); + + private String inputStream; + private String outputStream; + private long baseTimestamp; + private KafkaStreams kafkaStreams; + private TestInfo testInfo; + + @BeforeAll + public static void before() throws IOException { + CLUSTER.start(); + } + + @AfterAll + public static void after() { + CLUSTER.stop(); + } + + @BeforeEach + public void beforeTest(final TestInfo testInfo) throws InterruptedException { + this.testInfo = testInfo; + final String uniqueTestName = safeUniqueTestName(testInfo); + inputStream = "input-stream-" + uniqueTestName; + outputStream = "output-stream-" + uniqueTestName; + CLUSTER.createTopic(inputStream); + CLUSTER.createTopic(outputStream); + baseTimestamp = CLUSTER.time.milliseconds(); + } + + @AfterEach + public void afterTest() { + if (kafkaStreams != null) { + kafkaStreams.close(Duration.ofSeconds(30L)); + kafkaStreams.cleanUp(); + } + } + + @Test + public void shouldHandleTimestampedKeyWithHeadersQuery() throws Exception { + startStreams(); + + // key 1 has headers, key 2 has empty headers, key 3 is tombstoned (null value) + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, + KeyValue.pair(1, "a0")); + produceDataToTopicWithHeaders(inputStream, baseTimestamp + 1, new RecordHeaders(), + KeyValue.pair(2, "b0")); + produceDataToTopicWithHeaders(inputStream, baseTimestamp + 2, HEADERS, + KeyValue.pair(3, "c0"), KeyValue.pair(3, null)); + + // wait until all 4 input records have been processed (one output per input record) + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + TestUtils.consumerConfig(CLUSTER.bootstrapServers(), IntegerDeserializer.class, StringDeserializer.class), + outputStream, + 4); + + // key 1: value + timestamp + headers round-trip + final ValueTimestampHeaders<String> result1 = query(1); + assertEquals("a0", result1.value()); + assertEquals(baseTimestamp, result1.timestamp()); + assertEquals(HEADERS, result1.headers()); + + // key 2: written with no headers -> empty (never null) headers + final ValueTimestampHeaders<String> result2 = query(2); + assertEquals("b0", result2.value()); + assertEquals(baseTimestamp + 1, result2.timestamp()); + assertEquals(new RecordHeaders(), result2.headers()); + + // key 3: tombstoned -> null result, never a partially-populated wrapper + assertNull(query(3)); + + // never-written key -> null result + assertNull(query(999)); + } + + @Test + public void shouldFailWithUnknownQueryTypeAgainstNonHeadersStore() throws Exception { + // store built WITHOUT a WithHeaders supplier -> the query type is unsupported + final StreamsBuilder builder = new StreamsBuilder(); + builder + .addStateStore( + Stores.timestampedKeyValueStoreBuilder( + Stores.persistentTimestampedKeyValueStore(STORE_NAME), + Serdes.Integer(), + Serdes.String())) + .stream(inputStream, Consumed.with(Serdes.Integer(), Serdes.String())) + .process(() -> new PlainStoreWriterProcessor(), STORE_NAME) + .to(outputStream, Produced.with(Serdes.Integer(), Serdes.String())); + + kafkaStreams = new KafkaStreams(builder.build(), props()); + IntegrationTestUtils.startApplicationAndWaitUntilRunning(kafkaStreams); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS, KeyValue.pair(1, "a0")); + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + TestUtils.consumerConfig(CLUSTER.bootstrapServers(), IntegerDeserializer.class, StringDeserializer.class), + outputStream, + 1); + + final StateQueryResult<ValueTimestampHeaders<String>> result = + kafkaStreams.query(inStore(STORE_NAME).withQuery(TimestampedKeyWithHeadersQuery.withKey(1))); + + assertTrue(result.getOnlyPartitionResult().isFailure()); + assertEquals(FailureReason.UNKNOWN_QUERY_TYPE, result.getOnlyPartitionResult().getFailureReason()); + } + + private void startStreams() throws Exception { + final StreamsBuilder builder = new StreamsBuilder(); + builder + .addStateStore( + Stores.timestampedKeyValueStoreWithHeadersBuilder( + Stores.persistentTimestampedKeyValueStoreWithHeaders(STORE_NAME), + Serdes.Integer(), + Serdes.String()) + // Disable caching so every IQv2 query is forced down to the persistent + // RocksDBTimestampedStoreWithHeaders layer, exercising its KeyQuery handling + // (rather than being short-circuited by a cache hit). + .withCachingDisabled()) Review Comment: Caching is disabled here so the query is forced down to the RocksDB layer — good for exercising the new bottom-layer `KeyQuery` handling. But it means the cache-*hit* path (query served from `CachingKeyValueStore` before reaching RocksDB) is never exercised end-to-end. Consider an additional case (or a parameterized caching on/off variant) that queries with caching enabled and an unflushed record, to cover the metered store's behavior on a cache hit. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
