Extend the lease by default time

Signed-off-by: Dinu John <[email protected]>
opensearch-project · May 9, 2024 · 5c95c71 · 5c95c71
1 parent f261e3c
commit 5c95c71
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/...n/java/org/opensearch/dataprepper/plugins/mongo/stream/DataStreamPartitionCheckpoint.java b/...n/java/org/opensearch/dataprepper/plugins/mongo/stream/DataStreamPartitionCheckpoint.java
@@ -66,6 +66,11 @@ public void checkpoint(final String resumeToken, final long recordCount) {
         enhancedSourceCoordinator.saveProgressStateForPartition(streamPartition, CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE);
     }
 
+    public void extendLease() {
+        LOG.debug("Extending lease of stream partition for collection {}", streamPartition.getCollection());
+        enhancedSourceCoordinator.saveProgressStateForPartition(streamPartition, CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE);
+    }
+
     /**
      * This method is to reset checkpoint when change stream is invalid. The current thread will give up partition and new thread
      * will take ownership of partition. If change stream is valid then new thread proceeds with processing change stream else the

diff --git a/...in/java/org/opensearch/dataprepper/plugins/mongo/stream/StreamAcknowledgementManager.java b/...in/java/org/opensearch/dataprepper/plugins/mongo/stream/StreamAcknowledgementManager.java
@@ -80,8 +80,8 @@ private void monitorAcknowledgment(final ExecutorService executorService, final
                     }
                 } else {
                     if (System.currentTimeMillis() - lastCheckpointTime >= checkPointIntervalInMs) {
-                        LOG.debug("No records processed. Checkpoint to extend the lease of the worker");
-                        partitionCheckpoint.checkpoint(null, 0);
+                        LOG.debug("No records processed. Extend the lease of the partition worker.");
+                        partitionCheckpoint.extendLease();
                         lastCheckpointTime = System.currentTimeMillis();
                     }
                 }

diff --git a/...va/org/opensearch/dataprepper/plugins/mongo/stream/DataStreamPartitionCheckpointTest.java b/...va/org/opensearch/dataprepper/plugins/mongo/stream/DataStreamPartitionCheckpointTest.java
@@ -62,4 +62,10 @@ public void resetCheckpoint_success() {
         verify(streamProgressState).setLoadedRecords(0);
         verify(streamProgressState).setLastUpdateTimestamp(anyLong());
     }
+
+    @Test
+    public void extendLease_success() {
+        dataStreamPartitionCheckpoint.extendLease();
+        verify(enhancedSourceCoordinator).saveProgressStateForPartition(streamPartition, CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE);
+    }
 }