diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/QueryWorkerStatisticsHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/QueryWorkerStatisticsHandler.scala index 82cb7cc7997..2d48ef05e04 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/QueryWorkerStatisticsHandler.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/QueryWorkerStatisticsHandler.scala @@ -142,6 +142,13 @@ trait QueryWorkerStatisticsHandler { // Skip operators not included in the filtered subset (if any) if (opFilter.nonEmpty && !opFilter.contains(opId)) { Seq.empty + } else if (cp.workflowExecutionManager.isRegionTerminating(opId)) { + // The COMPLETED check below does not cover this: termination starts once every port of + // the region is booked complete (`RegionExecution.isCompleted`), whereas an operator + // reads as COMPLETED only once its workers' states say so, and a worker sends + // `portCompleted` before `workerExecutionCompleted`. Between the two, the region is + // being torn down while its operators still aggregate as RUNNING. + Seq.empty } else { cp.workflowExecution.getLatestOperatorExecutionOption(opId) match { // Operator region has not been initialized yet; skip in this polling round. 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..363f025ca9d 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 @@ -23,7 +23,7 @@ import org.apache.pekko.pattern.gracefulStop import com.twitter.util.{Future, JavaTimer, Return, Throw, Timer} import org.apache.texera.amber.core.state.State import org.apache.texera.amber.core.storage.{DocumentFactory, VFSURIFactory} -import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity +import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, PhysicalOpIdentity} import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink, PhysicalOp} import org.apache.texera.amber.engine.architecture.common.{ PekkoActorRefMappingService, @@ -59,7 +59,7 @@ import org.apache.texera.web.SessionState import org.apache.texera.web.model.websocket.event.RegionStateEvent import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.AtomicReference +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} import scala.concurrent.duration.{Duration => ScalaDuration} object RegionExecutionManager { @@ -132,6 +132,16 @@ class RegionExecutionManager( ) private val terminationFutureRef: AtomicReference[Future[Unit]] = new AtomicReference(null) + /** + * Set once `EndWorker` has gone out to this region's workers. From that point a worker refuses + * any further request, so the coordinator must send none — see `isTerminating`. + * + * Deliberately not derived from `terminationFutureRef`: that reference is CAS-ed only after + * `terminateWorkersWithRetry` has been constructed, and constructing it already runs + * `terminateWorkers`, so it would flip after the fan-out rather than before it. + */ + private val endWorkerSentRef: AtomicBoolean = new AtomicBoolean(false) + /** * Sync the status of `RegionExecution` and transition this manager's phase to `Completed` only when the * manager is currently in `ExecutingNonDependeePortsPhase`, all the ports of this region are completed, and @@ -185,6 +195,10 @@ class RegionExecutionManager( } private def terminateWorkers(regionExecution: RegionExecution) = { + // From here on these workers refuse anything else the coordinator might send them, so report + // the region as terminating before the fan-out rather than after it. + endWorkerSentRef.set(true) + // 1. Send EndWorkers to every worker val endWorkerRequests = regionExecution.getAllOperatorExecutions.flatMap { @@ -276,9 +290,23 @@ class RegionExecutionManager( def isCompleted: Boolean = currentPhaseRef.get == Completed + /** + * True once this region's workers have been sent `EndWorker`. A request arriving after it is work + * the worker must do, so `EndHandler` refuses to terminate and this region pays a retry — callers + * that would otherwise address a worker of this region must consult this first. + */ + def isTerminating: Boolean = endWorkerSentRef.get + + /** Whether `opId` is one of this region's operators. */ + def containsPhysicalOp(opId: PhysicalOpIdentity): Boolean = + region.getOperators.exists(_.id == opId) + /** * Returns the region termination future if termination has been initiated. * This is only set by `tryCompleteRegionExecution()`. + * + * Not a substitute for `isTerminating`: this is CAS-ed after the future is constructed, and + * constructing it already begins the `EndWorker` fan-out. */ def getTerminationFutureOpt: Option[Future[Unit]] = Option(terminationFutureRef.get) 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 0b63d1f17cc..7909639560a 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 @@ -22,6 +22,7 @@ package org.apache.texera.amber.engine.architecture.scheduling import com.twitter.util.Future import com.typesafe.scalalogging.LazyLogging import org.apache.texera.amber.core.storage.VFSURIFactory +import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink} import org.apache.texera.amber.engine.architecture.scheduling.config.InputPortConfig import org.apache.texera.amber.engine.architecture.common.{ @@ -179,4 +180,19 @@ class WorkflowExecutionManager( regionExecutionManagers.values.exists(!_.isCompleted) } + /** + * Whether `opId` belongs to a region whose workers have already been sent `EndWorker`, and which + * therefore must not be sent anything further. + * + * Call this at the point of use rather than caching the answer: a caller that spans several + * coordinator rounds — `QueryWorkerStatisticsHandler` awaits one operator layer before emitting + * the next — would otherwise act on an answer from before a region started tearing down, which is + * exactly the window this guards. + */ + def isRegionTerminating(opId: PhysicalOpIdentity): Boolean = { + regionExecutionManagers.values.exists(manager => + manager.isTerminating && manager.containsPhysicalOp(opId) + ) + } + } diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManagerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManagerSpec.scala index 147a8e153ac..d305d23790b 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManagerSpec.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManagerSpec.scala @@ -141,6 +141,37 @@ class RegionExecutionManagerSpec assert(workerState(fixture) == WorkerState.TERMINATED) } + // A request that lands after `EndWorker` is work the worker must do, so `EndHandler` refuses to + // terminate and the region pays a retry. `isTerminating` is what lets the rest of the coordinator + // avoid that — notably `QueryWorkerStatisticsHandler`, whose layered traversal spans enough + // coordinator rounds to walk into a region that started tearing down after the traversal began. + it should "report itself as terminating from the moment EndWorker is sent" in { + val terminatingWhenEndWorkerSent = new atomic.AtomicBoolean(false) + lazy val fixture: SingleRegionFixture = createSingleRegionFixture(endWorkerResponse = _ => { + terminatingWhenEndWorkerSent.set(fixture.manager.isTerminating) + Some(EmptyReturn()) + }) + + assert(!fixture.manager.isTerminating) + launchRegion(fixture.manager) + // Still false while the region is merely running: nothing has been sent yet. + assert(!fixture.manager.isTerminating) + + await(requestRegionCompletion(fixture.manager)) + + // The flag has to be observable by the time EndWorker goes out, not merely afterwards. + assert(terminatingWhenEndWorkerSent.get) + assert(fixture.manager.isTerminating) + } + + it should "recognise only its own operators as belonging to the region" in { + val fixture = createSingleRegionFixture(endWorkerResponse = _ => Some(EmptyReturn())) + val foreignOp = createSourceOp("other-op") + + assert(fixture.manager.containsPhysicalOp(fixture.physicalOp.id)) + assert(!fixture.manager.containsPhysicalOp(foreignOp.id)) + } + it should "give up with a descriptive error once the EndWorker retry budget is exhausted" in { // EndWorker always fails: a worker that never finishes draining. val fixture = createSingleRegionFixture( 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..83183bdbf4e 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 @@ -31,7 +31,10 @@ import org.apache.texera.amber.core.virtualidentity.{ import org.apache.texera.amber.core.workflow.PhysicalOp import org.apache.texera.amber.engine.architecture.coordinator.CoordinatorConfig import org.apache.texera.amber.engine.architecture.coordinator.execution.WorkflowExecution -import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + WorkflowAggregatedState +} import org.apache.texera.amber.engine.architecture.scheduling.RegionExecutionManagerTestSupport._ import org.apache.texera.amber.engine.common.AmberRuntime import org.scalatest.BeforeAndAfterAll @@ -150,6 +153,72 @@ class WorkflowExecutionManagerSpec assert(rpcProbe.startedWorkers.contains(secondWorkerId)) } + /** + * A request that reaches a worker after its `EndWorker` is work the worker must do, so + * `EndHandler` refuses to terminate and the region pays a retry (#6891). + * `QueryWorkerStatisticsHandler` consults `isRegionTerminating` to stay out of that window, + * which its layered traversal is otherwise wide open to: it awaits each operator layer before + * emitting the next, so a region can start tearing down between two layers of one query. + * + * The window this asserts is not covered by the handler's existing completed-operator skip. + * Termination begins once every port of the region is booked complete, whereas an operator only + * reads as COMPLETED once its workers' states say so — and here the worker has answered + * `startWorker` with RUNNING and nothing since, exactly as a real worker looks while its + * `workerExecutionCompleted` is still queued at the coordinator. + */ + it should "report a region as terminating from the moment its EndWorker is sent" in { + val firstOp = createSourceOp("first-op") + val firstWorkerId = createWorkerId(firstOp) + val firstRegion = createSingleWorkerRegion(1, firstOp, firstWorkerId) + + val secondOp = createSourceOp("second-op") + val secondWorkerId = createWorkerId(secondOp) + val secondRegion = createSingleWorkerRegion(2, secondOp, secondWorkerId) + + val workflowExecution = WorkflowExecution() + seedReusableWorkerExecution(workflowExecution, seedRegionId = 101, firstOp, firstWorkerId) + seedReusableWorkerExecution(workflowExecution, seedRegionId = 102, secondOp, secondWorkerId) + + // Hold the first region's endWorker pending, so it stays mid-teardown for the assertions. + val rpcProbe = new CoordinatorRpcProbe( + endWorkerResponse = call => if (call.receiver == firstWorkerId) None else Some(EmptyReturn()) + ) + val coordinator = createCoordinatorHarness() + registerLiveWorker(coordinator.actorRefService, firstWorkerId) + registerLiveWorker(coordinator.actorRefService, secondWorkerId) + + val workflowManager = new WorkflowExecutionManager( + workflowExecution, + CoordinatorConfig(None, None, None, None), + rpcProbe.asyncRPCClient + ) + workflowManager.schedule = Schedule(Map(0 -> Set(firstRegion), 1 -> Set(secondRegion))) + workflowManager.setupActorRefService(coordinator.actorRefService) + + await(workflowManager.advanceRegionExecutions(coordinator.actorService)) + // Running, nothing sent yet. + assert(!workflowManager.isRegionTerminating(firstOp.id)) + + val advanceFuture = workflowManager.advanceRegionExecutions(coordinator.actorService) + waitUntil(rpcProbe.endWorkerCalls.size == 1) + + // EndWorker is on the wire and unanswered: the region must now be off limits. + assert(workflowManager.isRegionTerminating(firstOp.id)) + // The operator still aggregates as RUNNING, so the completed-operator skip would not fire here. + assert( + workflowExecution.getLatestOperatorExecutionOption(firstOp.id).get.getState == + WorkflowAggregatedState.RUNNING + ) + // A region that has not begun terminating is unaffected. + assert(!workflowManager.isRegionTerminating(secondOp.id)) + + rpcProbe.fulfill(rpcProbe.onlyEndWorkerCall, EmptyReturn()) + await(advanceFuture) + + // Still terminating after the fact: the workers are gone, so nothing may address them again. + assert(workflowManager.isRegionTerminating(firstOp.id)) + } + "Jumping to an operator's region" should "make the next scheduled region contain the target operator's region" in { val (first, second, _, schedule) = threeLevelSchedule()