Fix KafkaBuffer isEmpty

Signed-off-by: Krishna Kondaka <[email protected]>
kkondaka · Feb 29, 2024 · 0ad52c9 · 0ad52c9
1 parent 337eb71
commit 0ad52c9
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/AbstractBuffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/AbstractBuffer.java
@@ -145,12 +145,14 @@ public Map.Entry<Collection<T>, CheckpointState> read(int timeoutInMillis) {
     @Override
     public void checkpoint(final CheckpointState checkpointState) {
         checkpointTimer.record(() -> doCheckpoint(checkpointState));
-        final int numRecordsToBeChecked = checkpointState.getNumRecordsToBeChecked();
-        recordsInFlight.addAndGet(-numRecordsToBeChecked);
-        recordsProcessedCounter.increment(numRecordsToBeChecked);
+        if (!isByteBuffer()) {
+            final int numRecordsToBeChecked = checkpointState.getNumRecordsToBeChecked();
+            recordsInFlight.addAndGet(-numRecordsToBeChecked);
+            recordsProcessedCounter.increment(numRecordsToBeChecked);
+        }
     }
 
-    protected int getRecordsInFlight() {
+    public int getRecordsInFlight() {
         return recordsInFlight.intValue();
     }
 

diff --git a/...c/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java b/...c/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java
@@ -150,6 +150,11 @@ void write_and_read() throws TimeoutException {
         // TODO: The metadata is not included. It needs to be included in the Buffer, though not in the Sink. This may be something we make configurable in the consumer/producer - whether to serialize the metadata or not.
         //assertThat(onlyResult.getData().getMetadata(), equalTo(record.getData().getMetadata()));
         assertThat(onlyResult.getData().toMap(), equalTo(record.getData().toMap()));
+        assertThat(objectUnderTest.getRecordsInFlight(), equalTo(0));
+        assertThat(objectUnderTest.getInnerBufferRecordsInFlight(), equalTo(1));
+        objectUnderTest.checkpoint(readResult.getValue());
+        assertThat(objectUnderTest.getRecordsInFlight(), equalTo(0));
+        assertThat(objectUnderTest.getInnerBufferRecordsInFlight(), equalTo(0));
     }
 
     @Test

diff --git a/...ka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java b/...ka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java
@@ -162,7 +162,7 @@ public void postProcess(final Long recordsInBuffer) {
     public void doCheckpoint(CheckpointState checkpointState) {
         try {
             setMdc();
-            innerBuffer.doCheckpoint(checkpointState);
+            innerBuffer.checkpoint(checkpointState);
         } finally {
             resetMdc();
         }
@@ -188,6 +188,10 @@ public boolean isWrittenOffHeapOnly() {
         return true;
     }
 
+    int getInnerBufferRecordsInFlight() {
+        return innerBuffer.getRecordsInFlight();
+    }
+
     @Override
     public void shutdown() {
         try {