aliehsaeedii commented on code in PR #21995: URL: https://github.com/apache/kafka/pull/21995#discussion_r3574155064
########## streams/src/main/java/org/apache/kafka/streams/state/internals/VersionedKeyValueStoreBuilderWithHeaders.java: ########## @@ -0,0 +1,85 @@ +/* + * 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.state.internals; + +import org.apache.kafka.common.serialization.Serde; +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.StoreBuilder; +import org.apache.kafka.streams.state.VersionedBytesStore; +import org.apache.kafka.streams.state.VersionedBytesStoreSupplier; +import org.apache.kafka.streams.state.VersionedKeyValueStore; + +import java.util.Objects; + +/** + * Builder for versioned key-value stores with headers support. + * Follows the same pattern as {@link VersionedKeyValueStoreBuilder} but wraps + * with a headers-aware changelog store. + */ +public class VersionedKeyValueStoreBuilderWithHeaders<K, V> + extends AbstractStoreBuilder<K, V, VersionedKeyValueStore<K, V>> { + + private final VersionedBytesStoreSupplier storeSupplier; + + public VersionedKeyValueStoreBuilderWithHeaders(final VersionedBytesStoreSupplier storeSupplier, + final Serde<K> keySerde, + final Serde<V> valueSerde, + final Time time) { + super( + storeSupplier.name(), + keySerde, + valueSerde, + time); + Objects.requireNonNull(storeSupplier, "storeSupplier can't be null"); + Objects.requireNonNull(storeSupplier.metricsScope(), "storeSupplier's metricsScope can't be null"); + this.storeSupplier = storeSupplier; + } + + @Override + public VersionedKeyValueStore<K, V> build() { + final KeyValueStore<Bytes, byte[]> store = storeSupplier.get(); + if (!(store instanceof VersionedBytesStore)) { + throw new IllegalStateException( + "VersionedBytesStoreSupplier.get() must return an instance of VersionedBytesStore"); + } + + return new MeteredVersionedKeyValueStore<>( Review Comment: build() returns a plain MeteredVersionedKeyValueStore, which doesn't implement VersionedKeyValueStoreWithHeaders and has no put(k,v,ts,headers). So getStateStore(name) cast to the headers interface throws ClassCastException, and headers never reach the store — the 4-arg put is unreachable from a topology. This needs a metered wrapper that implements the headers interface (cf. MeteredSessionStoreWithHeaders). ########## streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStoreWithHeaders.java: ########## @@ -0,0 +1,128 @@ +/* + * 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.state.internals; + +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.common.utils.internals.ByteUtils; +import org.apache.kafka.streams.state.HeadersBytesStore; +import org.apache.kafka.streams.state.VersionedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.VersionedRecord; + +import java.nio.ByteBuffer; +import java.util.Objects; + +/** + * A persistent, versioned key-value store based on RocksDB that additionally + * preserves record headers. + * <p> + * Headers are embedded into the value bytes using the format: + * {@code [headersSize(varint)][headersBytes][rawValue]} + * before delegating to the parent {@link RocksDBVersionedStore}. On reads, + * headers are extracted from the stored value bytes. + */ +public class RocksDBVersionedStoreWithHeaders + extends RocksDBVersionedStore + implements VersionedKeyValueStoreWithHeaders<Bytes, byte[]> { + + RocksDBVersionedStoreWithHeaders(final String name, + final String metricsScope, + final long historyRetention, + final long segmentInterval) { + super(name, metricsScope, historyRetention, segmentInterval); + } + + @Override + public long put(final Bytes key, final byte[] value, final long timestamp, final Headers headers) { Review Comment: put embeds headers into the value bytes, but the changelog logs the plain value (headers go out as native record headers) and restoreBatch writes that plain value back with no prefix. After restore, get()'s extractRawValue reads a bogus varint length and corrupts the value — this hits even the empty-headers path, so any use of this store breaks on restore. Restore needs to re-embed headers, or the store should use the changelog's native-header format instead of embedding. ########## streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStoreWithHeaders.java: ########## @@ -0,0 +1,128 @@ +/* + * 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.state.internals; + +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeaders; +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.common.utils.internals.ByteUtils; +import org.apache.kafka.streams.state.HeadersBytesStore; +import org.apache.kafka.streams.state.VersionedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.VersionedRecord; + +import java.nio.ByteBuffer; +import java.util.Objects; + +/** + * A persistent, versioned key-value store based on RocksDB that additionally + * preserves record headers. + * <p> + * Headers are embedded into the value bytes using the format: + * {@code [headersSize(varint)][headersBytes][rawValue]} + * before delegating to the parent {@link RocksDBVersionedStore}. On reads, + * headers are extracted from the stored value bytes. + */ +public class RocksDBVersionedStoreWithHeaders + extends RocksDBVersionedStore + implements VersionedKeyValueStoreWithHeaders<Bytes, byte[]> { + + RocksDBVersionedStoreWithHeaders(final String name, + final String metricsScope, + final long historyRetention, + final long segmentInterval) { + super(name, metricsScope, historyRetention, segmentInterval); + } + + @Override + public long put(final Bytes key, final byte[] value, final long timestamp, final Headers headers) { + Objects.requireNonNull(headers, "headers cannot be null"); + if (value == null) { + // tombstone: delegate directly, no headers to embed + return super.put(key, null, timestamp); + } + + // Check if headers are empty and use fast path + if (!headers.iterator().hasNext()) { + final byte[] encodedValue = HeadersBytesStore.convertToHeaderFormat(value); + return super.put(key, encodedValue, timestamp); + } + + // Use shared HeadersSerializer infrastructure + final HeadersSerializer.PreSerializedHeaders prep = HeadersSerializer.prepareSerialization(headers); + final int payloadSize = prep.requiredBufferSizeForHeaders + value.length; + final ByteBuffer buffer = ByteBuffer.allocate(ByteUtils.sizeOfVarint(prep.requiredBufferSizeForHeaders) + payloadSize); + ByteUtils.writeVarint(prep.requiredBufferSizeForHeaders, buffer); + HeadersSerializer.serialize(prep, buffer); + buffer.put(value); + return super.put(key, buffer.array(), timestamp); + } + + @Override + public long put(final Bytes key, final byte[] value, final long timestamp) { + // non-headers put: embed empty headers + return put(key, value, timestamp, new RecordHeaders()); + } + + @Override + public VersionedRecord<byte[]> get(final Bytes key) { + final VersionedRecord<byte[]> record = super.get(key); + return decodeRecord(record); + } + + @Override + public VersionedRecord<byte[]> get(final Bytes key, final long asOfTimestamp) { Review Comment: get(key), get(key,ts) and delete strip the header prefix, but the multi-version query path (get(key,from,to,order), used by MultiVersionedKeyQuery/IQv2) is inherited from the parent unchanged, so it returns the raw [headersSize][headers][value] bytes to callers. That read path needs decoding too. ########## streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/VersionedKeyValueStoreWithHeadersIntegrationTest.java: ########## @@ -0,0 +1,389 @@ +/* + * 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.StringSerializer; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster; +import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.ProcessorSupplier; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.Stores; +import org.apache.kafka.streams.state.VersionedKeyValueStoreWithHeaders; +import org.apache.kafka.streams.state.VersionedRecord; +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.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Properties; + +import static org.apache.kafka.streams.utils.TestUtils.safeUniqueTestName; +import static org.junit.jupiter.api.Assertions.assertEquals; + +@Tag("integration") +public class VersionedKeyValueStoreWithHeadersIntegrationTest { + + private static final String STORE_NAME = "versioned-headers-store"; + private static final Duration HISTORY_RETENTION = Duration.ofMinutes(10); + + private String inputStream; + private String outputStream; + private long baseTimestamp; + + private KafkaStreams kafkaStreams; + + private static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(1); + + private static final Headers HEADERS1 = new RecordHeaders() + .add("source", "test".getBytes()) + .add("version", "1.0".getBytes()); + + private static final Headers HEADERS2 = new RecordHeaders() + .add("source", "test".getBytes()) + .add("version", "2.0".getBytes()); + + private static final Headers EMPTY_HEADERS = new RecordHeaders(); + + public TestInfo testInfo; + + @BeforeAll + public static void before() throws IOException { + CLUSTER.start(); + } + + @AfterAll + public static void after() { + CLUSTER.stop(); + } + + @BeforeEach + public void setUp(final TestInfo testInfo) throws Exception { + this.testInfo = testInfo; + final String testName = safeUniqueTestName(testInfo); + inputStream = "input-" + testName; + outputStream = "output-" + testName; + baseTimestamp = System.currentTimeMillis() - Duration.ofMinutes(5).toMillis(); + CLUSTER.createTopic(inputStream, 1, 1); + CLUSTER.createTopic(outputStream, 1, 1); + } + + @AfterEach + public void tearDown() { + if (kafkaStreams != null) { + kafkaStreams.close(Duration.ofSeconds(30)); + } + } + + @Test + public void shouldPutAndGetWithHeaders() throws Exception { + // Processor that puts records with headers and verifies store contents + final VersionedStoreWithHeadersCheckerProcessor processor = + new VersionedStoreWithHeadersCheckerProcessor(true); + + final Topology topology = buildTopology(() -> processor); + kafkaStreams = new KafkaStreams(topology, streamsConfiguration()); + kafkaStreams.start(); + + // Produce records with HEADERS1 + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS1, + new KeyValue<>(1, "value1")); + + // Wait for output (processor forwards failedChecks count) + final List<KeyValue<Integer, Integer>> results = + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + consumerConfig(), outputStream, 1, Duration.ofSeconds(30).toMillis()); + + assertEquals(0, results.get(0).value, "Store content check failed"); + } + + @Test + public void shouldPutAndGetWithEmptyHeaders() throws Exception { + final VersionedStoreWithHeadersCheckerProcessor processor = + new VersionedStoreWithHeadersCheckerProcessor(true); + + final Topology topology = buildTopology(() -> processor); + kafkaStreams = new KafkaStreams(topology, streamsConfiguration()); + kafkaStreams.start(); + + // Produce records with empty headers + produceDataToTopicWithHeaders(inputStream, baseTimestamp, EMPTY_HEADERS, + new KeyValue<>(1, "value1")); + + final List<KeyValue<Integer, Integer>> results = + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + consumerConfig(), outputStream, 1, Duration.ofSeconds(30).toMillis()); + + assertEquals(0, results.get(0).value, "Store content check failed for empty headers"); + } + + @Test + public void shouldPreserveHeadersAcrossMultipleVersions() throws Exception { + final VersionedStoreWithHeadersCheckerProcessor processor = + new VersionedStoreWithHeadersCheckerProcessor(true); + + final Topology topology = buildTopology(() -> processor); + kafkaStreams = new KafkaStreams(topology, streamsConfiguration()); + kafkaStreams.start(); + + // Produce first version with HEADERS1 + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS1, + new KeyValue<>(1, "version1")); + + final List<KeyValue<Integer, Integer>> results1 = + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + consumerConfig(), outputStream, 1, Duration.ofSeconds(30).toMillis()); + assertEquals(0, results1.get(0).value, "First version check failed"); + + // Produce second version with HEADERS2 at a later timestamp + produceDataToTopicWithHeaders(inputStream, baseTimestamp + 1000, HEADERS2, + new KeyValue<>(1, "version2")); + + final List<KeyValue<Integer, Integer>> results2 = + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + consumerConfig(), outputStream, 2, Duration.ofSeconds(30).toMillis()); + assertEquals(0, results2.get(1).value, "Second version check failed"); + } + + @Test + public void shouldRestoreHeadersFromChangelog() throws Exception { + final Properties config = streamsConfiguration(); + // Use a fixed state dir so we can restart and restore + config.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getAbsolutePath()); + + // First run: write data with headers + final VersionedStoreWithHeadersCheckerProcessor processor1 = + new VersionedStoreWithHeadersCheckerProcessor(true); + + Topology topology = buildTopology(() -> processor1); + kafkaStreams = new KafkaStreams(topology, config); + kafkaStreams.start(); + + produceDataToTopicWithHeaders(inputStream, baseTimestamp, HEADERS1, + new KeyValue<>(1, "restored-value")); + + final List<KeyValue<Integer, Integer>> firstRunResults = + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + consumerConfig(), outputStream, 1, Duration.ofSeconds(30).toMillis()); + assertEquals(0, firstRunResults.get(0).value, "First run store check failed"); + + // Stop the streams app + kafkaStreams.close(Duration.ofSeconds(30)); + + // Second run: read from store without writing (verifies restoration) + final VersionedStoreWithHeadersCheckerProcessor processor2 = + new VersionedStoreWithHeadersCheckerProcessor(false); + // Seed the processor's expected data from the first run + processor2.seedExpectedData(1, "restored-value", baseTimestamp, HEADERS1); + + topology = buildTopology(() -> processor2); + kafkaStreams = new KafkaStreams(topology, config); + kafkaStreams.start(); + + // Produce a dummy record to trigger processing + produceDataToTopicWithHeaders(inputStream, baseTimestamp + 5000, EMPTY_HEADERS, + new KeyValue<>(99, "trigger")); + + final List<KeyValue<Integer, Integer>> secondRunResults = + IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived( + consumerConfig(), outputStream, 1, Duration.ofSeconds(30).toMillis()); + assertEquals(0, secondRunResults.get(0).value, "Restoration check failed: headers not preserved"); + } + + private Topology buildTopology( + final ProcessorSupplier<Integer, String, Integer, Integer> processorSupplier) { + final var supplier = Stores.persistentVersionedKeyValueStoreWithHeaders( + STORE_NAME, HISTORY_RETENTION); + + final Topology topology = new Topology(); + topology.addSource("source", new IntegerDeserializer(), + Serdes.String().deserializer(), inputStream); + topology.addProcessor("processor", processorSupplier, "source"); + topology.addStateStore( + Stores.versionedKeyValueStoreBuilderWithHeaders( + supplier, Serdes.Integer(), Serdes.String()), + "processor"); + topology.addSink("sink", outputStream, new IntegerSerializer(), + new IntegerSerializer(), "processor"); + return topology; + } + + private Properties streamsConfiguration() { + final Properties config = new Properties(); + config.put(StreamsConfig.APPLICATION_ID_CONFIG, + "versioned-headers-it-" + safeUniqueTestName(testInfo)); + config.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); + config.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class); + config.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class); + config.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getAbsolutePath()); + config.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000L); + return config; + } + + private Properties consumerConfig() { + final Properties config = new Properties(); + config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); + config.put(ConsumerConfig.GROUP_ID_CONFIG, "versioned-headers-it-consumer"); + config.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); + config.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, IntegerDeserializer.class); + config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, IntegerDeserializer.class); + return config; + } + + @SuppressWarnings("varargs") + @SafeVarargs + private final int produceDataToTopicWithHeaders(final String topic, + final long timestamp, + final Headers headers, + final KeyValue<Integer, String>... keyValues) { + IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp( + topic, + Arrays.asList(keyValues), + TestUtils.producerConfig(CLUSTER.bootstrapServers(), + IntegerSerializer.class, + StringSerializer.class), + headers, + timestamp, + false); + return keyValues.length; + } + + /** + * Processor that stores records in a VersionedKeyValueStoreWithHeaders and validates + * that the store contents match expectations. Forwards the number of failed checks + * as the output value. + */ + private static class VersionedStoreWithHeadersCheckerProcessor + implements Processor<Integer, String, Integer, Integer> { + + private ProcessorContext<Integer, Integer> context; + private VersionedKeyValueStoreWithHeaders<Integer, String> store; + + private final boolean writeToStore; + private final Map<Integer, Optional<VersionedRecordExpectation>> data; + + VersionedStoreWithHeadersCheckerProcessor(final boolean writeToStore) { + this.writeToStore = writeToStore; + this.data = new HashMap<>(); + } + + void seedExpectedData(final int key, final String value, + final long timestamp, final Headers headers) { + data.put(key, Optional.of(new VersionedRecordExpectation(value, timestamp, headers))); + } + + @Override + public void init(final ProcessorContext<Integer, Integer> context) { + this.context = context; + store = context.getStateStore(STORE_NAME); Review Comment: This casts getStateStore(...) to VersionedKeyValueStoreWithHeaders, but the builder returns a plain MeteredVersionedKeyValueStore that doesn't implement it, so this throws ClassCastException. These integration tests can't pass as written — the builder's missing metered wrapper needs fixing first. -- 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]
