fix(amber): advance region executions in a later coordinator round, not inside portCompleted

amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/coordinatorservice.proto

@@ -45,6 +45,7 @@ service CoordinatorService {
   rpc JumpToOperatorRegion(JumpToOperatorRegionRequest) returns (EmptyReturn);
   rpc LinkWorkers(LinkWorkersRequest) returns (EmptyReturn);
   rpc CoordinatorInitiateQueryStatistics(QueryStatisticsRequest) returns (EmptyReturn);
+  rpc CoordinatorInitiateAdvanceRegionExecutions(EmptyRequest) returns (EmptyReturn);
   rpc RetryWorkflow(RetryWorkflowRequest) returns (EmptyReturn);
   rpc ReconfigureWorkflow(WorkflowReconfigureRequest) returns (EmptyReturn);
 }

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

@@ -41,6 +41,7 @@ class CoordinatorAsyncRPCHandlerInitializer(
     with ResumeHandler
     with StartWorkflowHandler
     with PortCompletedHandler
+    with AdvanceRegionExecutionsHandler
     with ConsoleMessageHandler
     with RetryWorkflowHandler
     with EvaluatePythonExpressionHandler

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

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.coordinator.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.engine.architecture.coordinator.{
+  CoordinatorAsyncRPCHandlerInitializer,
+  FatalError
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+  AsyncRPCContext,
+  EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+
+/** Advance the region executions of this workflow.
+  *
+  * A handler that needs region executions advanced sends this to itself rather than advancing
+  * inline (see `PortCompletedHandler`), so that the advance — which terminates completed regions
+  * and therefore sends `EndWorker` to their workers — runs in its own control round, after every
+  * reply the requesting round owed.
+  *
+  * possible sender: coordinator (itself)
+  */
+trait AdvanceRegionExecutionsHandler {
+  this: CoordinatorAsyncRPCHandlerInitializer =>
+
+  override def coordinatorInitiateAdvanceRegionExecutions(
+      request: EmptyRequest,
+      ctx: AsyncRPCContext
+  ): Future[EmptyReturn] = {
+    cp.workflowExecutionManager
+      .advanceRegionExecutions(cp.actorService)
+      // The requester is the coordinator itself and discards this reply, so a failure has no
+      // caller to propagate to. A fatal error is sent to the client, indicating that the region
+      // cannot be scheduled.
+      .onFailure {
+        case err: WorkflowRuntimeException =>
+          sendToClient(FatalError(err, err.relatedWorkerId))
+        case other =>
+          sendToClient(FatalError(other, None))
+      }
+    EmptyReturn()
+  }
+
+}

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

@@ -20,14 +20,11 @@
 package org.apache.texera.amber.engine.architecture.coordinator.promisehandlers
 
 import com.twitter.util.Future
-import org.apache.texera.amber.core.WorkflowRuntimeException
 import org.apache.texera.amber.core.workflow.GlobalPortIdentity
-import org.apache.texera.amber.engine.architecture.coordinator.{
-  CoordinatorAsyncRPCHandlerInitializer,
-  FatalError
-}
+import org.apache.texera.amber.engine.architecture.coordinator.CoordinatorAsyncRPCHandlerInitializer
 import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
   AsyncRPCContext,
+  EmptyRequest,
   PortCompletedRequest,
   QueryStatisticsRequest,
   StatisticsUpdateTarget
@@ -81,16 +78,17 @@ trait PortCompletedHandler {
               else operatorExecution.isOutputPortCompleted(msg.portId)
 
             if (isPortCompleted) {
-              cp.workflowExecutionManager
-                .advanceRegionExecutions(cp.actorService)
-                // Since this message is sent from a worker, any exception from the above code will be returned to that worker.
-                // Additionally, a fatal error is sent to the client, indicating that the region cannot be scheduled.
-                .onFailure {
-                  case err: WorkflowRuntimeException =>
-                    sendToClient(FatalError(err, err.relatedWorkerId))
-                  case other =>
-                    sendToClient(FatalError(other, None))
-                }
+              // Advance region executions in a later control round instead of here. Advancing
+              // inline terminates the completed region and sends `EndWorker` to this very sender
+              // before this handler's own reply, on the same control channel — the worker would
+              // then process `EndWorker` with the reply still queued behind it and reject the
+              // termination (see `EndHandler`). A message the coordinator addresses to itself is
+              // transmitted and received before it is handled, so the advance lands behind the
+              // reply below.
+              coordinatorInterface.coordinatorInitiateAdvanceRegionExecutions(
+                EmptyRequest(),
+                COORDINATOR
+              )
             }
           case None => // currently "start" and "end" ports are not part of a region, thus no region can be found.
           // do nothing.

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

@@ -134,8 +134,16 @@ class RegionExecutionManager(
     *
     * Additionally, this method will also terminate all the workers of this region:
     *
-    * 1.  An `EndWorker` control message is first sent to all the workers. This will be the last message each worker
-    * receives. We wait for all workers have replied to indicate they have finished processing all control messages.
+    * 1.  An `EndWorker` control message is first sent to all the workers. We wait for all workers to reply that they
+    * have finished processing all control messages; a worker that has not fails the request, and
+    * `terminateWorkersWithRetry` re-sends `EndWorker` after `killRetryDelay`.
+    *
+    * Because a worker rejects `EndWorker` while work is still queued for it, termination must not be triggered from
+    * inside the handler of a request sent by one of these workers — the `EndWorker` would be emitted before that
+    * handler's own reply and overtake it on their shared FIFO control channel. Such handlers therefore send themselves
+    * a `CoordinatorInitiateAdvanceRegionExecutions` instead of advancing inline, which defers the advance that reaches
+    * here to a later coordinator round (see `PortCompletedHandler`). That orders the replies already produced; for one
+    * the coordinator has not produced yet, `EndHandler` does not count a queued reply as work.
     *
     * 2. Only after all workers have processed all control messages do we send a `gracefulStop` (pekko message) to each
     * worker. JVM workers will be terminated by `gracefulStop`. Python proxy workes will also be terminated by

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

@@ -96,6 +96,10 @@ class WorkflowExecutionManager(
     * in `Completed` status (phase).
     *
     * After the syncs, if there are no running region(s), it will start new regions (if available).
+    *
+    * Callers handling a worker-initiated request must not call this directly; they send themselves a
+    * `CoordinatorInitiateAdvanceRegionExecutions` (see `PortCompletedHandler`) so the resulting
+    * `EndWorker` cannot overtake the reply that request still owes.
     */
   def advanceRegionExecutions(actorService: PekkoActorService): Future[Unit] = {
     val unfinishedRegionManagers =

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala

@@ -24,8 +24,15 @@ import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
   AsyncRPCContext,
   EmptyRequest
 }
-import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+  EmptyReturn,
+  ReturnInvocation
+}
 import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+  DPInputQueueElement,
+  FIFOMessageElement
+}
 
 /**
   * The EndWorker control messages is needed to ensure all the other control messages in a worker
@@ -37,20 +44,66 @@ trait EndHandler {
   /**
     * The response of endWorker to the coordinator indicates that this worker has finished not only
     * the data processing logic, but also the processing of all the control messages.
+    *
+    * A queued reply to one of this worker's own requests does not count: it carries no work (see
+    * `findUnprocessedWork`), and the coordinator cannot order all of its replies before `EndWorker`.
+    * The coordinator defers region advancement to a later control round precisely so that
+    * `EndWorker` follows the replies it owes (see `PortCompletedHandler`), but that only orders
+    * requests it has already handled. A request of this worker's that is still queued at the
+    * coordinator when the deferred advance runs — `workerExecutionCompleted`, which the worker
+    * emits right after its last `portCompleted` — is replied to afterwards, and the coordinator
+    * picks its input channels out of a `HashMap` (`NetworkInputGateway.tryPickControlChannel`), so
+    * there is no cross-channel order to rely on.
     */
   override def endWorker(
       request: EmptyRequest,
       ctx: AsyncRPCContext
   ): Future[EmptyReturn] = {
-    // Ensure this is really the last message.
-    if (!dp.inputManager.inputMessageQueue.isEmpty) {
+    // Ensure this is really the last message that asks this worker to do anything.
+    val pendingWork = findUnprocessedWork
+    if (pendingWork.isDefined) {
       logger.warn(
-        s"Received EndHandler before all messages are processed. Unprocessed messages: " +
-          s"${dp.inputManager.inputMessageQueue.peek()}"
+        s"Received EndHandler before all messages are processed. Unprocessed message: " +
+          s"${describe(pendingWork.get)}"
       )
       return Future.exception(new IllegalStateException("worker still has unprocessed messages"))
     }
     // Now we can safely acknowledge that this worker can be terminated.
     EmptyReturn()
   }
+
+  /**
+    * The first queued element that represents work, if any.
+    *
+    * A `ReturnInvocation` is excluded: processing one only fulfills a promise for a request this
+    * worker already issued (`AmberProcessor.processDCM`), and every worker-to-coordinator call
+    * discards its future, so nothing is pending on it. Everything else — control invocations, data,
+    * embedded control messages, timer-based controls, actor commands — still blocks termination, so
+    * the coordinator retries and this worker drains it first.
+    *
+    * This is the worker's own arrival queue rather than its actor mailbox, and at termination it
+    * holds at most a couple of elements.
+    */
+  private def findUnprocessedWork: Option[DPInputQueueElement] = {
+    val iterator = dp.inputManager.inputMessageQueue.iterator()
+    while (iterator.hasNext) {
+      val element = iterator.next()
+      val isWork = element match {
+        case FIFOMessageElement(message) => !message.payload.isInstanceOf[ReturnInvocation]
+        case _                           => true
+      }
+      if (isWork) {
+        return Some(element)
+      }
+    }
+    None
+  }
+
+  /** Identifies a queued element without logging payload contents. */
+  private def describe(element: DPInputQueueElement): String =
+    element match {
+      case FIFOMessageElement(message) =>
+        s"${message.payload.getClass.getSimpleName} on ${message.channelId}"
+      case other => other.getClass.getSimpleName
+    }
 }