fix(amber): run completion cleanup

amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorProcessor.scala

@@ -46,7 +46,12 @@ class CoordinatorProcessor(
   val workflowExecutionManager: WorkflowExecutionManager = new WorkflowExecutionManager(
     workflowExecution,
     coordinatorConfig,
-    asyncRPCClient
+    asyncRPCClient,
+    onWorkflowCompleted = () =>
+      Option(coordinatorTimerService).foreach { timerService =>
+        timerService.disableStatusUpdate()
+        timerService.disableRuntimeStatisticsCollection()
+      }
   )
 
   private val initializer = new CoordinatorAsyncRPCHandlerInitializer(this)

amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/WorkerExecutionCompletedHandler.scala

@@ -20,10 +20,7 @@
 package org.apache.texera.amber.engine.architecture.coordinator.promisehandlers
 
 import com.twitter.util.Future
-import org.apache.texera.amber.engine.architecture.coordinator.{
-  CoordinatorAsyncRPCHandlerInitializer,
-  ExecutionStateUpdate
-}
+import org.apache.texera.amber.engine.architecture.coordinator.CoordinatorAsyncRPCHandlerInitializer
 import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
   AsyncRPCContext,
   EmptyRequest,
@@ -51,27 +48,10 @@ trait WorkerExecutionCompletedHandler {
     // after worker execution is completed, query statistics immediately one last time
     // because the worker might be killed before the next query statistics interval
     // and the user sees the last update before completion
-    val statsRequest =
-      coordinatorInterface.coordinatorInitiateQueryStatistics(
-        QueryStatisticsRequest(Seq(ctx.sender), StatisticsUpdateTarget.BOTH_UI_AND_PERSISTENCE),
-        mkContext(SELF)
-      )
-
-    Future
-      .collect(Seq(statsRequest))
-      .flatMap(_ => {
-        // if entire workflow is completed, clean up
-        val isWorkflowTerminal =
-          cp.workflowExecution.isCompleted &&
-            !cp.workflowScheduler.hasPendingRegions &&
-            !cp.workflowExecutionManager.hasUnfinishedRegionManagers
-        if (isWorkflowTerminal) {
-          // after query result come back: send completed event, cleanup ,and kill workflow
-          sendToClient(ExecutionStateUpdate(cp.workflowExecution.getState))
-          cp.coordinatorTimerService.disableStatusUpdate()
-          cp.coordinatorTimerService.disableRuntimeStatisticsCollection()
-        }
-      })
+    coordinatorInterface.coordinatorInitiateQueryStatistics(
+      QueryStatisticsRequest(Seq(ctx.sender), StatisticsUpdateTarget.BOTH_UI_AND_PERSISTENCE),
+      mkContext(SELF)
+    )
     EmptyReturn()
   }
 }

amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala

@@ -38,7 +38,8 @@ import scala.collection.mutable
 class WorkflowExecutionManager(
     workflowExecution: WorkflowExecution,
     coordinatorConfig: CoordinatorConfig,
-    asyncRPCClient: AsyncRPCClient
+    asyncRPCClient: AsyncRPCClient,
+    onWorkflowCompleted: () => Unit = () => ()
 ) extends LazyLogging {
 
   var schedule: Schedule = Schedule(Map.empty)
@@ -132,6 +133,7 @@ class WorkflowExecutionManager(
     if (nextRegions.isEmpty) {
       if (workflowExecution.isCompleted && completionNotified.compareAndSet(false, true)) {
         asyncRPCClient.sendToClient(ExecutionStateUpdate(workflowExecution.getState))
+        onWorkflowCompleted()
       }
       return Future.Unit
     }

amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManagerSpec.scala

@@ -150,6 +150,41 @@ class WorkflowExecutionManagerSpec
     assert(rpcProbe.startedWorkers.contains(secondWorkerId))
   }
 
+  it should "run workflow-completion cleanup exactly once after the workflow becomes terminal" in {
+    val workflowExecution = WorkflowExecution()
+    workflowExecution.initRegionExecution(jumpRegion(1, "completed"))
+    assert(workflowExecution.isCompleted)
+
+    val rpcProbe = new CoordinatorRpcProbe(_ => Some(EmptyReturn()))
+    var cleanupCalls = 0
+    val workflowManager = new WorkflowExecutionManager(
+      workflowExecution,
+      CoordinatorConfig(None, None, None, None),
+      rpcProbe.asyncRPCClient,
+      onWorkflowCompleted = () => cleanupCalls += 1
+    )
+
+    await(workflowManager.advanceRegionExecutions(null))
+    await(workflowManager.advanceRegionExecutions(null))
+
+    assert(cleanupCalls == 1)
+  }
+
+  it should "not run workflow-completion cleanup for an uninitialized workflow" in {
+    val rpcProbe = new CoordinatorRpcProbe(_ => Some(EmptyReturn()))
+    var cleanupCalls = 0
+    val workflowManager = new WorkflowExecutionManager(
+      WorkflowExecution(),
+      CoordinatorConfig(None, None, None, None),
+      rpcProbe.asyncRPCClient,
+      onWorkflowCompleted = () => cleanupCalls += 1
+    )
+
+    await(workflowManager.advanceRegionExecutions(null))
+
+    assert(cleanupCalls == 0)
+  }
+
   "Jumping to an operator's region" should
     "make the next scheduled region contain the target operator's region" in {
     val (first, second, _, schedule) = threeLevelSchedule()