fix(amber): skip statistics queries to workers of a terminating region

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.

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)
 

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)
+    )
+  }
+
 }

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(

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()