diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorProcessor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorProcessor.scala index 1d75a0ef376..fc7f7c807cc 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorProcessor.scala +++ b/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) diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/WorkerExecutionCompletedHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/WorkerExecutionCompletedHandler.scala index 5be2fc499f0..ace857183f6 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/WorkerExecutionCompletedHandler.scala +++ b/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() } } diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala index b736cd171c0..e49b659d78e 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala +++ b/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 } diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManagerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManagerSpec.scala index 7ddc7d29678..6e0f1507e28 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManagerSpec.scala +++ b/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()