KAFKA-20538: Make RocksDBTimeOrderedKeyValueBuffer headers-aware (#22166)

Problem: `RocksDBTimeOrderedKeyValueBuffer` (the persistent buffer used
by stream-table joins with grace periods) was missed when headers were
propagated through Streams serdes. Two bugs:

1. evictWhile() deserializes with the wrong headers. Both key and value
deserializers are called with internalContext.headers() -- the headers of
whatever record the processor is currently handling, not the record
being evicted. Since records can sit in the buffer for the entire grace
period, these almost never match.

2. put() does not persist the record's headers. It stores the
caller-supplied recordContext directly in BufferValue, so
record.headers() is never copied into the persisted
`ProcessorRecordContext` and is unrecoverable at eviction time.

This PR fixes the followings:

 - put() builds a fresh ProcessorRecordContext carrying
record.headers() and stores that in BufferValue, so the headers are
persisted with the record

- evictWhile() uses bufferValue.context().headers() for both key and
value deserialization, so each evicted record is deserialized with the
headers it was buffered with.

Reviewers: Matthias J. Sax <matthias@confluent.io>

Author: aliehsaeedii

streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedKeyValueBuffer.java

@@ -231,7 +231,7 @@ public void evictWhile(final Supplier<Boolean> predicate, final Consumer<Evictio
 
                     final BufferValue bufferValue = BufferValue.deserialize(ByteBuffer.wrap(keyValue.value));
                     final K key = keySerde.deserializer().deserialize(topic,
-                        internalContext.headers(),
+                        bufferValue.context().headers(),
                         PrefixedWindowKeySchemas.TimeFirstWindowKeySchema.extractStoreKeyBytes(keyValue.key.get()));
 
                     if (bufferValue.context().timestamp() < minTimestamp && minValid) {
@@ -243,7 +243,7 @@ public void evictWhile(final Supplier<Boolean> predicate, final Consumer<Evictio
                     minTimestamp = bufferValue.context().timestamp();
                     minValid = true;
 
-                    final V value = valueSerde.deserializer().deserialize(topic, internalContext.headers(), bufferValue.newValue());
+                    final V value = valueSerde.deserializer().deserialize(topic, bufferValue.context().headers(), bufferValue.newValue());
 
                     callback.accept(new Eviction<>(key, value, bufferValue.context()));
 

streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKTableJoinTest.java

@@ -17,8 +17,15 @@
 package org.apache.kafka.streams.kstream.internals;
 
 import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.IntegerSerializer;
+import org.apache.kafka.common.serialization.Serde;
 import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.common.utils.LogCaptureAppender;
 import org.apache.kafka.streams.KeyValue;
@@ -35,6 +42,7 @@
 import org.apache.kafka.streams.kstream.KTable;
 import org.apache.kafka.streams.kstream.Materialized;
 import org.apache.kafka.streams.state.Stores;
+import org.apache.kafka.streams.test.TestRecord;
 import org.apache.kafka.test.MockApiProcessor;
 import org.apache.kafka.test.MockApiProcessorSupplier;
 import org.apache.kafka.test.MockValueJoiner;
@@ -44,6 +52,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Collection;
@@ -311,6 +320,82 @@ public void shouldHandleLateJoinsWithGracePeriod(final boolean withHeaders) {
             new KeyValueTimestamp<>(0, "X0+Y0", 0));
     }
 
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    public void shouldDeserializeBufferedValueWithPutTimeHeadersDuringEviction(final boolean withHeaders) {
+        builder = new StreamsBuilder();
+        final Consumed<Integer, String> consumed = Consumed.with(Serdes.Integer(), Serdes.String());
+        final KStream<Integer, String> stream = builder.stream(streamTopic, consumed);
+        final KTable<Integer, String> table = builder.table("tableTopic2", consumed, Materialized.as(
+            Stores.persistentVersionedKeyValueStore("V-grace", Duration.ofMinutes(5))));
+        stream.join(table,
+            MockValueJoiner.TOSTRING_JOINER,
+            // The built-in leaf serdes (String, primitives, …) ignore the `headers` argument, so with any of them the test would pass identically
+            // with and without the fix (a false green). The HeaderValueAppendingSerde is the minimal serde whose deserializer reflects the `headers`
+            // argument into the value
+            Joined.with(Serdes.Integer(), new HeaderValueAppendingSerde(), Serdes.String(), "Grace", Duration.ofMillis(2))
+        ).process(supplier);
+        final Properties props = StreamsTestUtils.getStreamsConfig(Serdes.Integer(), Serdes.String());
+        StreamsTestUtils.maybeSetDslStoreFormatHeaders(props, withHeaders);
+        driver = new TopologyTestDriver(builder.build(), props);
+        inputStreamTopic = driver.createInputTopic(streamTopic, new IntegerSerializer(), new StringSerializer(), Instant.ofEpochMilli(0L), Duration.ZERO);
+        inputTableTopic = driver.createInputTopic("tableTopic2", new IntegerSerializer(), new StringSerializer(), Instant.ofEpochMilli(0L), Duration.ZERO);
+        processor = supplier.theCapturedProcessor();
+
+        // Table entry so the buffered stream record finds a join match when it is evicted.
+        inputTableTopic.pipeInput(0, "Y0", 0L);
+
+        // Record A is buffered (grace not yet elapsed) with header v=first at t=0. Nothing is emitted yet.
+        inputStreamTopic.pipeInput(new TestRecord<>(0, "X0", headers("first"), 0L));
+        processor.checkAndClearProcessResult(EMPTY);
+
+        // Record B carries header v=second and its t=10 timestamp advances stream time past A's grace window, which
+        // triggers A's eviction while the live processor context still holds B's headers (v=second). B itself is
+        // buffered, not evicted, so it produces no output here.
+        inputStreamTopic.pipeInput(new TestRecord<>(1, "X1", headers("second"), 10L));
+
+        // A must be deserialized with its own put-time header (first), NOT B's live-context header (second). Before
+        // the fix this asserted "X0[second]+Y0".
+        processor.checkAndClearProcessResult(new KeyValueTimestamp<>(0, "X0[first]+Y0", 0));
+    }
+
+    private static Headers headers(final String value) {
+        return new RecordHeaders(new Header[]{new RecordHeader("v", value.getBytes(StandardCharsets.UTF_8))});
+    }
+
+    /**
+     * A stream-side value {@link Serde} whose <em>deserializer</em> output depends on the headers it is handed: it
+     * appends the value of the {@code "v"} header to the deserialized string. Serialization is a plain string encode
+     * and ignores headers. This fixture exists so tests can observe <em>which</em> headers reached the buffer's value
+     * deserializer during eviction -- a plain serde would ignore the headers argument and hide the difference.
+     */
+    private static final class HeaderValueAppendingSerde implements Serde<String> {
+        @Override
+        public Serializer<String> serializer() {
+            return new StringSerializer();
+        }
+
+        @Override
+        public Deserializer<String> deserializer() {
+            return new Deserializer<>() {
+                @Override
+                public String deserialize(final String topic, final byte[] data) {
+                    return data == null ? null : new String(data, StandardCharsets.UTF_8);
+                }
+
+                @Override
+                public String deserialize(final String topic, final Headers headers, final byte[] data) {
+                    final String base = deserialize(topic, data);
+                    final Header header = headers == null ? null : headers.lastHeader("v");
+                    if (base == null || header == null) {
+                        return base;
+                    }
+                    return base + "[" + new String(header.value(), StandardCharsets.UTF_8) + "]";
+                }
+            };
+        }
+    }
+
     @ParameterizedTest
     @ValueSource(booleans = {false, true})
     public void shouldReuseRepartitionTopicWithGeneratedName(final boolean withHeaders) {

streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedKeyValueBufferTest.java

@@ -16,11 +16,18 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.Serde;
 import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.streams.processor.TaskId;
 import org.apache.kafka.streams.processor.api.Record;
@@ -40,11 +47,17 @@
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
 
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.everyItem;
+import static org.hamcrest.Matchers.hasItem;
+import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
@@ -205,6 +218,100 @@ public void shouldHandleCollidingKeys() {
         assertNumSizeAndTimestamp(buffer, 1, 7, 42);
     }
 
+    @Test
+    public void shouldDeserializeWithPutTimeHeadersEvenAfterContextMutation() {
+        final HeaderCapturingSerde serde = new HeaderCapturingSerde();
+        createBuffer(Duration.ZERO, serde);
+        final RecordHeaders putHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("at-put", "first".getBytes(StandardCharsets.UTF_8))
+        });
+        // Give the record its own headers, distinct from the context headers, so the assertions below prove that
+        // eviction deserialization reads from the stored recordContext snapshot rather than from record.headers().
+        final RecordHeaders recordHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("on-record", "rec".getBytes(StandardCharsets.UTF_8))
+        });
+        context.setRecordContext(new ProcessorRecordContext(0L, offset++, 0, "testing", putHeaders));
+        buffer.put(0L, new Record<>("k", "v", 0L, recordHeaders), context.recordContext());
+
+        // Simulate the processor moving on to another record with different headers before eviction runs.
+        final RecordHeaders laterHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("at-evict", "second".getBytes(StandardCharsets.UTF_8))
+        });
+        context.setRecordContext(new ProcessorRecordContext(0L, offset++, 0, "testing", laterHeaders));
+
+        final List<TimeOrderedKeyValueBuffer.Eviction<String, String>> evicted = new ArrayList<>();
+        buffer.evictWhile(() -> buffer.numRecords() > 0, evicted::add);
+
+        assertThat(evicted.size(), is(1));
+        // The key/value deserializers must see the headers captured at put time, not the mutated context headers.
+        assertThat(serde.capturedHeaders, hasItem(putHeaders));
+        assertThat(serde.capturedHeaders, everyItem(is(putHeaders)));
+        assertThat(evicted.get(0).recordContext().headers(), is(putHeaders));
+    }
+
+    @Test
+    public void shouldNotBeAffectedByProcessorContextHeaderMutationBetweenPutAndEvict() {
+        final HeaderCapturingSerde serde = new HeaderCapturingSerde();
+        createBuffer(Duration.ofMillis(1), serde);
+        final RecordHeaders putHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("at-put", "first".getBytes(StandardCharsets.UTF_8))
+        });
+        // Give the record its own headers, distinct from the context headers, so the assertions below prove that
+        // eviction deserialization reads from the stored recordContext snapshot rather than from record.headers().
+        final RecordHeaders recordHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("on-record", "rec".getBytes(StandardCharsets.UTF_8))
+        });
+        context.setRecordContext(new ProcessorRecordContext(0L, offset++, 0, "testing", putHeaders));
+        buffer.put(0L, new Record<>("k", "v", 0L, recordHeaders), context.recordContext());
+
+        // Simulate the processor moving on to handle a different record with different headers
+        // before the grace period expires and eviction runs.
+        final RecordHeaders laterHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("at-evict", "second".getBytes(StandardCharsets.UTF_8))
+        });
+        final RecordHeaders triggerRecordHeaders = new RecordHeaders(new Header[]{
+            new RecordHeader("on-trigger-record", "trig".getBytes(StandardCharsets.UTF_8))
+        });
+        context.setRecordContext(new ProcessorRecordContext(10L, offset++, 0, "testing", laterHeaders));
+        // Advance stream time past the grace period for the original record.
+        buffer.put(10L, new Record<>("trigger", "v", 10L, triggerRecordHeaders), context.recordContext());
+
+        final List<TimeOrderedKeyValueBuffer.Eviction<String, String>> evicted = new ArrayList<>();
+        buffer.evictWhile(() -> true, evicted::add);
+
+        // Only the original "k" record at t=0 falls outside the grace window of t=10.
+        assertThat(evicted.size(), is(1));
+        assertThat(evicted.get(0).key(), is("k"));
+        // The deserializers for the evicted record must see its put-time headers, not the later context headers.
+        assertThat(serde.capturedHeaders, hasItem(putHeaders));
+        assertThat(serde.capturedHeaders, everyItem(is(putHeaders)));
+        assertThat(evicted.get(0).recordContext().headers(), is(putHeaders));
+    }
+
+    /**
+     * A {@link Serde} whose deserializer records the {@link Headers} it is handed on each call, so tests can assert
+     * which headers reached the key/value deserializers during eviction. Serialization behaves like a plain String serde.
+     */
+    private static final class HeaderCapturingSerde implements Serde<String> {
+        private final List<Headers> capturedHeaders = new ArrayList<>();
+
+        @Override
+        public Serializer<String> serializer() {
+            return new StringSerializer();
+        }
+
+        @Override
+        public Deserializer<String> deserializer() {
+            return new StringDeserializer() {
+                @Override
+                public String deserialize(final String topic, final Headers headers, final byte[] data) {
+                    capturedHeaders.add(headers);
+                    return super.deserialize(topic, data);
+                }
+            };
+        }
+    }
+
     private void assertNumSizeAndTimestamp(final TimeOrderedKeyValueBuffer<String, String, String> buffer,
                                            final int num,
                                            final long time,