abstract recordHeartbeat method

Signed-off-by: Sean Kao <[email protected]>
opensearch-project · Jun 14, 2024 · 346e344 · 346e344
1 parent 9664595
commit 346e344
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 5 deletions.
diff --git a/...t-core/src/main/scala/org/opensearch/flint/core/metadata/log/FlintMetadataLogService.java b/...t-core/src/main/scala/org/opensearch/flint/core/metadata/log/FlintMetadataLogService.java
@@ -39,4 +39,11 @@ default <T> OptimisticTransaction<T> startTransaction(String indexName) {
    * @return optional metadata log
    */
   Optional<FlintMetadataLog<FlintMetadataLogEntry>> getIndexMetadataLog(String indexName);
+
+  /**
+   * Record heartbeat timestamp for index streaming job.
+   *
+   * @param indexName index name
+   */
+  void recordHeartbeat(String indexName);
 }
diff --git a/...e/src/main/scala/org/opensearch/flint/core/storage/FlintOpenSearchMetadataLogService.java b/...e/src/main/scala/org/opensearch/flint/core/storage/FlintOpenSearchMetadataLogService.java
@@ -17,6 +17,7 @@
 import org.opensearch.flint.core.metadata.log.DefaultOptimisticTransaction;
 import org.opensearch.flint.core.metadata.log.FlintMetadataLog;
 import org.opensearch.flint.core.metadata.log.FlintMetadataLogEntry;
+import org.opensearch.flint.core.metadata.log.FlintMetadataLogEntry.IndexState$;
 import org.opensearch.flint.core.metadata.log.FlintMetadataLogService;
 import org.opensearch.flint.core.metadata.log.OptimisticTransaction;
 
@@ -55,6 +56,14 @@ public Optional<FlintMetadataLog<FlintMetadataLogEntry>> getIndexMetadataLog(Str
     return getIndexMetadataLog(indexName, false);
   }
 
+  @Override
+  public void recordHeartbeat(String indexName) {
+    startTransaction(indexName)
+        .initialLog(latest -> latest.state() == IndexState$.MODULE$.REFRESHING())
+        .finalLog(latest -> latest) // timestamp will update automatically
+        .commit(latest -> null);
+  }
+
   private Optional<FlintMetadataLog<FlintMetadataLogEntry>> getIndexMetadataLog(String indexName, boolean initIfNotExist) {
     LOG.info("Getting metadata log for index " + indexName + " and data source " + dataSourceName);
     try (IRestHighLevelClient client = createOpenSearchClient()) {

diff --git a/...-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexMonitor.scala b/...-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexMonitor.scala
@@ -156,11 +156,7 @@ class FlintSparkIndexMonitor(
       try {
         if (isStreamingJobActive(indexName)) {
           logInfo("Streaming job is still active")
-          flintMetadataLogService
-            .startTransaction(indexName)
-            .initialLog(latest => latest.state == REFRESHING)
-            .finalLog(latest => latest) // timestamp will update automatically
-            .commit(_ => {})
+          flintMetadataLogService.recordHeartbeat(indexName)
         } else {
           logError("Streaming job is not active. Cancelling monitor task")
           stopMonitor(indexName)

diff --git a/integ-test/src/test/scala/org/opensearch/flint/core/FlintMetadataLogITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/core/FlintMetadataLogITSuite.scala
@@ -78,4 +78,23 @@ class FlintMetadataLogITSuite extends OpenSearchTransactionSuite with Matchers {
     val metadataLog = flintMetadataLogService.getIndexMetadataLog(testFlintIndex)
     metadataLog.isPresent shouldBe false
   }
+
+  test("should update timestamp when record heartbeat") {
+    val refreshingLogEntry = flintMetadataLogEntry.copy(state = REFRESHING)
+    createLatestLogEntry(refreshingLogEntry)
+    val updateTimeBeforeHeartbeat =
+      latestLogEntry(testLatestId).get("lastUpdateTime").get.asInstanceOf[Long]
+    flintMetadataLogService.recordHeartbeat(testFlintIndex)
+    latestLogEntry(testLatestId)
+      .get("lastUpdateTime")
+      .get
+      .asInstanceOf[Long] should be > updateTimeBeforeHeartbeat
+  }
+
+  test("should fail when record heartbeat if index not refreshing") {
+    createLatestLogEntry(flintMetadataLogEntry)
+    the[IllegalStateException] thrownBy {
+      flintMetadataLogService.recordHeartbeat(testFlintIndex)
+    }
+  }
 }