diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala index 0b32b74101a..74849cf0286 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala @@ -185,7 +185,9 @@ class RegionExecutionManager( } private def terminateWorkers(regionExecution: RegionExecution) = { - // 1. Send EndWorkers to every worker + implicit val timer: Timer = new JavaTimer(true) + val killTimeout = com.twitter.util.Duration.fromMilliseconds(1000) + // 1. Send EndWorkers with timeout val endWorkerRequests = regionExecution.getAllOperatorExecutions.flatMap { case (_, opExec) => @@ -196,9 +198,12 @@ class RegionExecutionManager( }.toSeq val endWorkerFuture: Future[Unit] = - Future.collect(endWorkerRequests).unit + Future + .collect(endWorkerRequests) + .within(killTimeout) + .unit - // 2. Send GracefulStops only after 1 has finished + // 2. Send GracefulStops with timeout val gracefulStopRequests: Future[Unit] = endWorkerFuture.flatMap { _ => val gracefulStops = @@ -206,20 +211,17 @@ class RegionExecutionManager( case (_, opExec) => opExec.getWorkerIds.map { workerId => val actorRef = actorRefService.getActorRef(workerId) - // Remove the actorRef so that no other actors can find the worker and send messages. - actorRefService.removeActorRef(workerId) - // Restarted regions reuse actorId. Remove stale control channels so the - // coordinator does not reuse old control-message sequence numbers for new workers. - asyncRPCClient.inputGateway.removeControlChannel(workerId) - asyncRPCClient.outputGateway.removeControlChannel(workerId) gracefulStop(actorRef, ScalaDuration(5, TimeUnit.SECONDS)).asTwitter() } }.toSeq - Future.collect(gracefulStops).unit + Future + .collect(gracefulStops) + .within(killTimeout) + .unit } - // 3. Log whether the kills were successful + // 3. Cleanup only after all gracefulStops succeed gracefulStopRequests.transform { case Return(_) => logger.debug(s"Region ${region.id.id} successfully terminated.") @@ -227,6 +229,12 @@ class RegionExecutionManager( case (_, opExec) => opExec.getWorkerIds.foreach { workerId => opExec.getWorkerExecution(workerId).forceTerminate() + // Remove the actorRef after successful termination so other actors cannot reach the worker. + actorRefService.removeActorRef(workerId) + // Restarted regions reuse actorId. Remove stale control channels so the + // coordinator does not reuse old control-message sequence numbers for new workers. + asyncRPCClient.inputGateway.removeControlChannel(workerId) + asyncRPCClient.outputGateway.removeControlChannel(workerId) } } Future.Unit // propagate success