KAFKA-16505: add source raw bytes in processorContex

Author: Dabz

streams/src/main/java/org/apache/kafka/streams/errors/ErrorHandlerContext.java

@@ -147,4 +147,39 @@ public interface ErrorHandlerContext {
      * @return The timestamp.
      */
     long timestamp();
+
+    /**
+     * Return the non-deserialized byte[] of the input message key if the context has been triggered by a message.
+     *
+     * <p> If this method is invoked within a {@link Punctuator#punctuate(long)
+     * punctuation callback}, or while processing a record that was forwarded by a punctuation
+     * callback, it will return null.
+     *
+     * <p> If this method is invoked in a sub-topology due to a repartition, the returned key would be one sent
+     * to the repartition topic.
+     *
+     * <p> Always returns null if this method is invoked within a
+     * ProductionExceptionHandler.handle(ErrorHandlerContext, ProducerRecord, Exception)
+     *
+     * @return the raw byte of the key of the source message
+     */
+    byte[] sourceRawKey();
+
+
+    /**
+     * Return the non-deserialized byte[] of the input message value if the context has been triggered by a message.
+     *
+     * <p> If this method is invoked within a {@link Punctuator#punctuate(long)
+     * punctuation callback}, or while processing a record that was forwarded by a punctuation
+     * callback, it will return null.
+     *
+     * <p> If this method is invoked in a sub-topology due to a repartition, the returned key would be one sent
+     * to the repartition topic.
+     *
+     * <p> Always returns null if this method is invoked within a
+     * ProductionExceptionHandler.handle(ErrorHandlerContext, ProducerRecord, Exception)
+     *
+     * @return the raw byte of the value of the source message
+     */
+    byte[] sourceRawValue();
 }

streams/src/main/java/org/apache/kafka/streams/errors/internals/DefaultErrorHandlerContext.java

@@ -33,6 +33,8 @@ public class DefaultErrorHandlerContext implements ErrorHandlerContext {
     private final Headers headers;
     private final String processorNodeId;
     private final TaskId taskId;
+    private final byte[] sourceRawKey;
+    private final byte[] sourceRawValue;
 
     private final long timestamp;
     private final ProcessorContext processorContext;
@@ -44,7 +46,9 @@ public DefaultErrorHandlerContext(final ProcessorContext processorContext,
                                       final Headers headers,
                                       final String processorNodeId,
                                       final TaskId taskId,
-                                      final long timestamp) {
+                                      final long timestamp,
+                                      final byte[] sourceRawKey,
+                                      final byte[] sourceRawValue) {
         this.topic = topic;
         this.partition = partition;
         this.offset = offset;
@@ -53,6 +57,8 @@ public DefaultErrorHandlerContext(final ProcessorContext processorContext,
         this.taskId = taskId;
         this.processorContext = processorContext;
         this.timestamp = timestamp;
+        this.sourceRawKey = sourceRawKey;
+        this.sourceRawValue = sourceRawValue;
     }
 
     @Override
@@ -90,6 +96,14 @@ public long timestamp() {
         return timestamp;
     }
 
+    public byte[] sourceRawKey() {
+        return sourceRawKey;
+    }
+
+    public byte[] sourceRawValue() {
+        return sourceRawValue;
+    }
+
     @Override
     public String toString() {
         // we do exclude headers on purpose, to not accidentally log user data

streams/src/main/java/org/apache/kafka/streams/processor/RecordContext.java

@@ -110,4 +110,39 @@ public interface RecordContext {
      */
     Headers headers();
 
+    /**
+     * Return the non-deserialized byte[] of the input message key if the context has been triggered by a message.
+     *
+     * <p> If this method is invoked within a {@link Punctuator#punctuate(long)
+     * punctuation callback}, or while processing a record that was forwarded by a punctuation
+     * callback, it will return null.
+     *
+     * <p> If this method is invoked in a sub-topology due to a repartition, the returned key would be one sent
+     * to the repartition topic.
+     *
+     * <p> Always returns null if this method is invoked within a
+     * ProductionExceptionHandler.handle(ErrorHandlerContext, ProducerRecord, Exception)
+     *
+     * @return the raw byte of the key of the source message
+     */
+    byte[] sourceRawKey();
+
+
+    /**
+     * Return the non-deserialized byte[] of the input message value if the context has been triggered by a message.
+     *
+     * <p> If this method is invoked within a {@link Punctuator#punctuate(long)
+     * punctuation callback}, or while processing a record that was forwarded by a punctuation
+     * callback, it will return null.
+     *
+     * <p> If this method is invoked in a sub-topology due to a repartition, the returned key would be one sent
+     * to the repartition topic.
+     *
+     * <p> Always returns null if this method is invoked within a
+     * ProductionExceptionHandler.handle(ErrorHandlerContext, ProducerRecord, Exception)
+     *
+     * @return the raw byte of the value of the source message
+     */
+    byte[] sourceRawValue();
+
 }

streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorNode.java

@@ -215,7 +215,9 @@ public void process(final Record<KIn, VIn> record) {
                 internalProcessorContext.headers(),
                 internalProcessorContext.currentNode().name(),
                 internalProcessorContext.taskId(),
-                internalProcessorContext.timestamp());
+                internalProcessorContext.timestamp(),
+                internalProcessorContext.recordContext().sourceRawKey(),
+                internalProcessorContext.recordContext().sourceRawValue());
 
             final ProcessingExceptionHandler.ProcessingHandlerResponse response;
             try {

streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java

@@ -37,6 +37,8 @@ public class ProcessorRecordContext implements RecordContext, RecordMetadata {
     private final String topic;
     private final int partition;
     private final Headers headers;
+    private byte[] sourceRawKey;
+    private byte[] sourceRawValue;
 
     public ProcessorRecordContext(final long timestamp,
                                   final long offset,
@@ -48,6 +50,25 @@ public ProcessorRecordContext(final long timestamp,
         this.topic = topic;
         this.partition = partition;
         this.headers = Objects.requireNonNull(headers);
+        this.sourceRawKey = null;
+        this.sourceRawValue = null;
+    }
+
+
+    public ProcessorRecordContext(final long timestamp,
+                                  final long offset,
+                                  final int partition,
+                                  final String topic,
+                                  final Headers headers,
+                                  final byte[] sourceRawKey,
+                                  final byte[] sourceRawValue) {
+        this.timestamp = timestamp;
+        this.offset = offset;
+        this.topic = topic;
+        this.partition = partition;
+        this.headers = Objects.requireNonNull(headers);
+        this.sourceRawKey = sourceRawKey;
+        this.sourceRawValue = sourceRawValue;
     }
 
     @Override
@@ -75,6 +96,16 @@ public Headers headers() {
         return headers;
     }
 
+    @Override
+    public byte[] sourceRawKey() {
+        return sourceRawKey;
+    }
+
+    @Override
+    public byte[] sourceRawValue() {
+        return sourceRawValue;
+    }
+
     public long residentMemorySizeEstimate() {
         long size = 0;
         size += Long.BYTES; // value.context.timestamp
@@ -176,6 +207,12 @@ public static ProcessorRecordContext deserialize(final ByteBuffer buffer) {
         return new ProcessorRecordContext(timestamp, offset, partition, topic, headers);
     }
 
+    public void freeRawRecord() {
+        this.sourceRawKey = null;
+        this.sourceRawValue = null;
+    }
+
+
     @Override
     public boolean equals(final Object o) {
         if (this == o) {

streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java

@@ -258,6 +258,10 @@ public <K, V> void send(final String topic,
 
         final ProducerRecord<byte[], byte[]> serializedRecord = new ProducerRecord<>(topic, partition, timestamp, keyBytes, valBytes, headers);
 
+        // As many records could be in-flight,
+        // freeing raw records in the context to reduce memory pressure
+        freeContext(context);
+
         streamsProducer.send(serializedRecord, (metadata, exception) -> {
             try {
                 // if there's already an exception record, skip logging offsets or new exceptions
@@ -310,6 +314,12 @@ public <K, V> void send(final String topic,
         });
     }
 
+    private static void freeContext(final InternalProcessorContext<Void, Void> context) {
+        if (context != null && context.recordContext() != null) {
+            context.recordContext().freeRawRecord();
+        }
+    }
+
     private <K, V> void handleException(final ProductionExceptionHandler.SerializationExceptionOrigin origin,
                                         final String topic,
                                         final K key,
@@ -387,7 +397,9 @@ private DefaultErrorHandlerContext errorHandlerContext(final InternalProcessorCo
                 recordContext.headers(),
                 processorNodeId,
                 taskId,
-                recordContext.timestamp()
+                recordContext.timestamp(),
+                context.recordContext().sourceRawKey(),
+                context.recordContext().sourceRawValue()
             ) :
             new DefaultErrorHandlerContext(
                 context,
@@ -397,7 +409,9 @@ private DefaultErrorHandlerContext errorHandlerContext(final InternalProcessorCo
                 new RecordHeaders(),
                 processorNodeId,
                 taskId,
-                -1L
+                -1L,
+                null,
+                null
             );
     }
 

streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordDeserializer.java

@@ -95,7 +95,9 @@ public static void handleDeserializationFailure(final DeserializationExceptionHa
             rawRecord.headers(),
             sourceNodeName,
             processorContext.taskId(),
-            rawRecord.timestamp());
+            rawRecord.timestamp(),
+            rawRecord.key(),
+            rawRecord.value());
 
         final DeserializationHandlerResponse response;
         try {

streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordQueue.java

@@ -243,7 +243,7 @@ private void updateHead() {
                 lastCorruptedRecord = raw;
                 continue;
             }
-            headRecord = new StampedRecord(deserialized, timestamp);
+            headRecord = new StampedRecord(deserialized, timestamp, raw.key(), raw.value());
             headRecordSizeInBytes = consumerRecordSizeInBytes(raw);
         }
 

streams/src/main/java/org/apache/kafka/streams/processor/internals/StampedRecord.java

@@ -23,8 +23,22 @@
 
 public class StampedRecord extends Stamped<ConsumerRecord<?, ?>> {
 
+    private final byte[] rawKey;
+    private final byte[] rawValue;
+
     public StampedRecord(final ConsumerRecord<?, ?> record, final long timestamp) {
         super(record, timestamp);
+        this.rawKey = null;
+        this.rawValue = null;
+    }
+
+    public StampedRecord(final ConsumerRecord<?, ?> record,
+                         final long timestamp,
+                         final byte[] rawKey,
+                         final byte[] rawValue) {
+        super(record, timestamp);
+        this.rawKey = rawKey;
+        this.rawValue = rawValue;
     }
 
     public String topic() {
@@ -55,8 +69,26 @@ public Headers headers() {
         return value.headers();
     }
 
+    public byte[] rawKey() {
+        return rawKey;
+    }
+
+    public byte[] rawValue() {
+        return rawValue;
+    }
+
     @Override
     public String toString() {
         return value.toString() + ", timestamp = " + timestamp;
     }
+
+    @Override
+    public boolean equals(final Object other) {
+        return super.equals(other);
+    }
+
+    @Override
+    public int hashCode() {
+        return super.hashCode();
+    }
 }

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java

@@ -853,7 +853,9 @@ private void doProcess(final long wallClockTime) {
             record.offset(),
             record.partition(),
             record.topic(),
-            record.headers()
+            record.headers(),
+            record.rawKey(),
+            record.rawValue()
         );
         updateProcessorContext(currNode, wallClockTime, recordContext);
 
@@ -935,11 +937,12 @@ record = null;
                 recordContext.headers(),
                 node.name(),
                 id(),
-                recordContext.timestamp()
+                recordContext.timestamp(),
+                recordContext.sourceRawKey(),
+                recordContext.sourceRawValue()
             );
 
-            final ProcessingExceptionHandler.ProcessingHandlerResponse response;
-            try {
+            final ProcessingExceptionHandler.ProcessingHandlerResponse response;            try {
                 response = Objects.requireNonNull(
                     processingExceptionHandler.handle(errorHandlerContext, null, processingException),
                     "Invalid ProcessingExceptionHandler response."