KAFKA-20847: Carry record metadata on suppress buffer changelog tombstones

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

@@ -52,11 +52,9 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.NoSuchElementException;
-import java.util.Set;
 import java.util.TreeMap;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
@@ -87,7 +85,8 @@ public final class InMemoryTimeOrderedKeyValueChangeBuffer<K, V, T> implements T
     private final Map<Bytes, BufferKey> index = new HashMap<>();
     private final TreeMap<BufferKey, BufferValue> sortedMap = new TreeMap<>();
 
-    private final Set<Bytes> dirtyKeys = new HashSet<>();
+    // Dirty key -> record context of its row, captured before eviction so tombstones stay self-describing.
+    private final Map<Bytes, ProcessorRecordContext> dirtyKeys = new HashMap<>();
     private final String storeName;
     private final boolean loggingEnabled;
 
@@ -269,13 +268,13 @@ public void close() {
     public void commit(final Map<TopicPartition, Long> changelogOffsets) {
         if (loggingEnabled) {
             // counting on this getting called before the record collector's flush
-            for (final Bytes key : dirtyKeys) {
-
+            for (final Map.Entry<Bytes, ProcessorRecordContext> entry : dirtyKeys.entrySet()) {
+                final Bytes key = entry.getKey();
                 final BufferKey bufferKey = index.get(key);
 
                 if (bufferKey == null) {
-                    // The record was evicted from the buffer. Send a tombstone.
-                    logTombstone(key);
+                    // The record was evicted from the buffer. Send a tombstone with the evicted row's timestamp/headers.
+                    logTombstone(key, entry.getValue());
                 } else {
                     final BufferValue value = sortedMap.get(bufferKey);
 
@@ -338,14 +337,17 @@ private static void addValuePartHeader(final RecordHeaders headers,
         }
     }
 
-    private void logTombstone(final Bytes key) {
+    private void logTombstone(final Bytes key, final ProcessorRecordContext recordContext) {
+        // Fall back to no timestamp/headers if the context is missing.
+        final Headers headers = recordContext == null ? null : recordContext.headers();
+        final Long timestamp = recordContext == null ? null : recordContext.timestamp();
         ((RecordCollector.Supplier) context).recordCollector().send(
             changelogTopic,
             key,
             null,
-            null,
+            headers,
             partition,
-            null,
+            timestamp,
             KEY_SERIALIZER,
             VALUE_SERIALIZER,
             null,
@@ -504,7 +506,8 @@ public void evictWhile(final Supplier<Boolean> predicate,
                 index.remove(next.getKey().key());
 
                 if (loggingEnabled) {
-                    dirtyKeys.add(next.getKey().key());
+                    // Keep the evicted row's context for the tombstone.
+                    dirtyKeys.put(next.getKey().key(), bufferValue.context());
                 }
 
                 memBufferSize -= computeRecordSize(next.getKey().key(), bufferValue);
@@ -726,7 +729,7 @@ public boolean put(final long time,
             new BufferValue(serializedPriorValue, oldValue, newValue, recordContext)
         );
         if (loggingEnabled) {
-            dirtyKeys.add(serializedKey);
+            dirtyKeys.put(serializedKey, recordContext);
         }
         updateBufferMetrics();
         return true;

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

@@ -916,7 +916,13 @@ public void shouldCommit(final String testName, final Function<String, B> buffer
         buffer.init(context, buffer);
         putRecord(buffer, context, 2L, 0L, "asdf", "2093j");
         putRecord(buffer, context, 1L, 1L, "zxcv", "3gon4i");
-        putRecord(buffer, context, 0L, 2L, "deleteme", "deadbeef");
+        // Buffer "deleteme" with headers so we can assert the tombstone preserves the row's timestamp + headers.
+        final RecordHeaders deletemeHeaders =
+            new RecordHeaders(new Header[] {new RecordHeader("h", new byte[] {(byte) 42})});
+        final ProcessorRecordContext deletemeContext =
+            new ProcessorRecordContext(2L, 0, 0, "topic", deletemeHeaders);
+        context.setRecordContext(deletemeContext);
+        buffer.put(0L, new Record<>("deleteme", new Change<>("deadbeef", null), 0L), deletemeContext);
 
         // replace "deleteme" with a tombstone
         buffer.evictWhile(() -> buffer.minTimestamp() < 1, kv -> { });
@@ -956,10 +962,10 @@ public void shouldCommit(final String testName, final Function<String, B> buffer
         assertThat(collected, is(asList(
             new ProducerRecord<>(APP_ID + "-" + testName + "-changelog",
                                  0,   // Producer will assign
-                                 null,
+                                 2L,  // the evicted row's timestamp is carried onto the tombstone
                                  "deleteme",
                                  null,
-                                 new RecordHeaders()
+                                 deletemeHeaders // the evicted row's headers are carried onto the tombstone
             ),
             new ProducerRecord<>(APP_ID + "-" + testName + "-changelog",
                                  0,