Introduce safe offset record header (#140)

What and why? For every record consumed from li-apache-kafka-client consumer, the record will be annotated with a safe offset header such that systems like brooklin aren't forced to use LiKafkaConsumer interface (which has #safeOffsets() api) but instead can just work with Consumer interface and the headers.
linkedin · Sep 6, 2019 · 313dbc9 · 313dbc9
1 parent 00cf733
commit 313dbc9
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 4 deletions.
diff --git a/...ients/src/main/java/com/linkedin/kafka/clients/largemessage/ConsumerRecordsProcessor.java b/...ients/src/main/java/com/linkedin/kafka/clients/largemessage/ConsumerRecordsProcessor.java
@@ -10,6 +10,7 @@
 import com.linkedin.kafka.clients.largemessage.errors.SkippableException;
 import com.linkedin.kafka.clients.utils.Constants;
 import com.linkedin.kafka.clients.utils.LiKafkaClientsUtils;
+import com.linkedin.kafka.clients.utils.PrimitiveEncoderDecoder;
 import java.util.Collections;
 import java.util.Objects;
 import java.util.Set;
@@ -438,6 +439,13 @@ private ConsumerRecord<K, V> handleConsumerRecord(ConsumerRecord<byte[], byte[]>
         headers.add(Constants.LARGE_MESSAGE_HEADER, LargeMessageHeaderValue.toBytes(largeMessageHeaderValue));
       }
 
+      // Introduce a safe offset header if and only if the safe offset is not the same as the high watermark
+      // Note: Safe offset is the last consumed record offset + 1
+      Long safeOffset = safeOffset(new TopicPartition(consumerRecord.topic(), consumerRecord.partition()));
+      if (safeOffset != null && safeOffset != consumerRecord.offset() + 1) {
+        headers.add(Constants.SAFE_OFFSET_HEADER, PrimitiveEncoderDecoder.encodeLong(safeOffset));
+      }
+
       handledRecord = new ConsumerRecord<>(
           consumerRecord.topic(),
           consumerRecord.partition(),

diff --git a/li-apache-kafka-clients/src/main/java/com/linkedin/kafka/clients/utils/Constants.java b/li-apache-kafka-clients/src/main/java/com/linkedin/kafka/clients/utils/Constants.java
@@ -8,6 +8,7 @@ public class Constants {
   // The variables reserved by kafka for auditing purpose
   public static final String TIMESTAMP_HEADER = "_t";
   public static final String LARGE_MESSAGE_HEADER = "_lm";
+  public static final String SAFE_OFFSET_HEADER = "_so";
 
   /**
    * Avoid instantiating the constants class

diff --git a/...s/src/test/java/com/linkedin/kafka/clients/largemessage/ConsumerRecordsProcessorTest.java b/...s/src/test/java/com/linkedin/kafka/clients/largemessage/ConsumerRecordsProcessorTest.java
@@ -7,8 +7,10 @@
 import com.linkedin.kafka.clients.largemessage.errors.OffsetNotTrackedException;
 import com.linkedin.kafka.clients.largemessage.errors.SkippableException;
 import com.linkedin.kafka.clients.producer.UUIDFactory;
+import com.linkedin.kafka.clients.utils.Constants;
 import com.linkedin.kafka.clients.utils.LiKafkaClientsUtils;
 import com.linkedin.kafka.clients.utils.LiKafkaClientsTestUtils;
+import com.linkedin.kafka.clients.utils.PrimitiveEncoderDecoder;
 import java.util.Collections;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -84,10 +86,25 @@ public void testCorrectness() {
     ConsumerRecords<String, String> processedRecords = consumerRecordsProcessor.process(getConsumerRecords()).consumerRecords();
     assertEquals(processedRecords.count(), 4, "There should be 4 records");
     Iterator<ConsumerRecord<String, String>> iter = processedRecords.iterator();
-    assertEquals(iter.next().offset(), 0, "Message offset should b 0");
-    assertEquals(iter.next().offset(), 2, "Message offset should b 2");
-    assertEquals(iter.next().offset(), 4, "Message offset should b 4");
-    assertEquals(iter.next().offset(), 5, "Message offset should b 5");
+    ConsumerRecord<String, String> record = iter.next();
+    assertEquals(record.offset(), 0, "Message offset should b 0");
+    assertNull(record.headers().lastHeader(Constants.SAFE_OFFSET_HEADER));
+
+    record = iter.next();
+    assertEquals(record.offset(), 2, "Message offset should b 2");
+    assertEquals(
+        PrimitiveEncoderDecoder.decodeLong(record.headers().lastHeader(Constants.SAFE_OFFSET_HEADER).value(), 0),
+        1L
+    );
+    record = iter.next();
+    assertEquals(record.offset(), 4, "Message offset should b 4");
+    assertEquals(
+        PrimitiveEncoderDecoder.decodeLong(record.headers().lastHeader(Constants.SAFE_OFFSET_HEADER).value(), 0),
+        1L
+    );
+    record = iter.next();
+    assertEquals(record.offset(), 5, "Message offset should b 5");
+    assertNull(record.headers().lastHeader(Constants.SAFE_OFFSET_HEADER));
   }
 
   @Test