fix(amber): tolerate in-flight RPC acks when EndWorker terminates a worker

amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py

@@ -18,10 +18,12 @@
 from loguru import logger
 
 from core.architecture.handlers.control.control_handler_base import ControlHandler
-from core.models.internal_queue import InternalQueue
+from core.models.internal_queue import DCMElement, InternalQueue
+from core.util.proto import get_one_of
 from proto.org.apache.texera.amber.engine.architecture.rpc import (
     EmptyReturn,
     EmptyRequest,
+    ReturnInvocation,
 )
 
 
@@ -37,19 +39,45 @@ async def end_worker(self, req: EmptyRequest) -> EmptyReturn:
         has finished not only the data processing logic, but also the processing
         of all the control messages.
         """
-        # Ensure this is really the last message. Read the queued count once (InternalQueue
-        # exposes size(); the base IQueue interface does not) and branch on it.
+        # Ensure this is really the last message that carries work. Read the
+        # queued count once (InternalQueue exposes size(); the base IQueue
+        # interface does not) and branch on it.
         input_queue: InternalQueue = self.context.input_queue
-        queued_count = input_queue.size()
-        if queued_count > 0:
+        if input_queue.size() == 0:
+            # Now we can safely acknowledge that this worker can be terminated.
+            return EmptyReturn()
+
+        # RPC acks (ReturnInvocations for this worker's own fire-and-forget
+        # calls, e.g. worker_execution_completed) race with EndWorker by
+        # design: the coordinator decides to end the worker while its acks are
+        # still in flight. The worker never awaits those acks, so an ack-only
+        # backlog is safe to drop at termination. Anything else fails loudly
+        # AND is put back on the queue so the coordinator's retried EndWorker
+        # finds it processed (mirrors the Scala EndHandler).
+        pending = []
+        while not input_queue.is_empty():
+            pending.append(input_queue.get())
+
+        def is_ack(element) -> bool:
+            return isinstance(element, DCMElement) and isinstance(
+                get_one_of(element.payload, sealed=False), ReturnInvocation
+            )
+
+        if all(is_ack(element) for element in pending):
             logger.warning(
-                f"Received EndWorker before all {queued_count} queued "
-                f"message(s) were processed; failing the RPC so a later "
-                f"coordinator retry succeeds once the queue has drained."
+                f"Received EndWorker with only RPC acks left in the queue; "
+                f"proceeding with termination. Pending acks: {pending}"
             )
-            # Fail this RPC (the counterpart of the Scala EndHandler's
-            # Future.exception) so a later coordinator retry succeeds once
-            # the queue has drained, instead of dropping the pending message.
-            raise RuntimeError("worker still has unprocessed messages")
-        # Now we can safely acknowledge that this worker can be terminated.
-        return EmptyReturn()
+            return EmptyReturn()
+
+        for element in pending:
+            input_queue.put(element)
+        logger.warning(
+            f"Received EndWorker before all {len(pending)} queued "
+            f"message(s) were processed; failing the RPC so a later "
+            f"coordinator retry succeeds once the queue has drained."
+        )
+        # Fail this RPC (the counterpart of the Scala EndHandler's
+        # Future.exception) so a later coordinator retry succeeds once
+        # the queue has drained, instead of dropping the pending message.
+        raise RuntimeError("worker still has unprocessed messages")

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

@@ -24,8 +24,14 @@ 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.FIFOMessageElement
+
+import scala.jdk.CollectionConverters.CollectionHasAsScala
 
 /**
   * The EndWorker control messages is needed to ensure all the other control messages in a worker
@@ -42,14 +48,33 @@ trait EndHandler {
       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 carries work. RPC acks
+    // (ReturnInvocations for this worker's own fire-and-forget calls, e.g.
+    // workerExecutionCompleted / portCompleted) race with EndWorker by design:
+    // the coordinator decides to end the worker while its acks are still in
+    // flight. The worker never awaits those acks, so an ack-only backlog is
+    // safe to leave unprocessed at termination -- failing on it only makes
+    // region termination retry and CI flaky. Anything else still fails loudly
+    // so the region execution manager retries the kill instead of dropping
+    // real work.
+    val pending = dp.inputManager.inputMessageQueue.asScala.toList
+    val ackOnly = pending.forall {
+      case FIFOMessageElement(msg) => msg.payload.isInstanceOf[ReturnInvocation]
+      case _                       => false
+    }
+    if (pending.nonEmpty && !ackOnly) {
       logger.warn(
         s"Received EndHandler before all messages are processed. Unprocessed messages: " +
-          s"${dp.inputManager.inputMessageQueue.peek()}"
+          s"$pending"
       )
       return Future.exception(new IllegalStateException("worker still has unprocessed messages"))
     }
+    if (pending.nonEmpty) {
+      logger.warn(
+        s"Received EndHandler with only RPC acks left in the queue; proceeding with " +
+          s"termination. Pending acks: $pending"
+      )
+    }
     // Now we can safely acknowledge that this worker can be terminated.
     EmptyReturn()
   }

amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py

@@ -15,17 +15,31 @@
 # specific language governing permissions and limitations
 # under the License.
 
+"""EndWorker is the coordinator's acknowledgement point before it stops the
+worker actor. A successful reply promises no queued work is dropped.
+
+RPC acks (ReturnInvocations for this worker's own fire-and-forget calls, e.g.
+worker_execution_completed) race with EndWorker by design: the coordinator
+decides to end the worker while its acks are still in flight. An ack-only
+backlog carries no work and must not fail the kill; anything else must fail
+loudly AND stay in the queue so the coordinator's retry lets the worker drain
+it.
+"""
+
 import asyncio
 from types import SimpleNamespace
 
 import pytest
 
 from core.architecture.handlers.control.end_worker_handler import EndWorkerHandler
 from core.models.internal_queue import DCMElement, InternalQueue
+from core.util.proto import set_one_of
 from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
 from proto.org.apache.texera.amber.engine.architecture.rpc import (
+    ControlReturn,
     EmptyRequest,
     EmptyReturn,
+    ReturnInvocation,
 )
 from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2
 
@@ -48,6 +62,26 @@ def make_control_message(seq: int) -> DCMElement:
         )
         return DCMElement(tag=channel, payload=DirectControlMessagePayloadV2())
 
+    @staticmethod
+    def make_rpc_ack(command_id: int) -> DCMElement:
+        """A ReturnInvocation ack for one of this worker's own fire-and-forget
+        RPCs -- the message kind that legitimately races with EndWorker."""
+        channel = ChannelIdentity(
+            ActorVirtualIdentity(name="COORDINATOR"),
+            ActorVirtualIdentity(name="worker-0"),
+            is_control=True,
+        )
+        return DCMElement(
+            tag=channel,
+            payload=set_one_of(
+                DirectControlMessagePayloadV2,
+                ReturnInvocation(
+                    command_id=command_id,
+                    return_value=set_one_of(ControlReturn, EmptyReturn()),
+                ),
+            ),
+        )
+
     def test_acknowledges_end_worker_on_empty_queue(self, handler):
         result = asyncio.run(handler.end_worker(EmptyRequest()))
         assert isinstance(result, EmptyReturn)
@@ -91,3 +125,32 @@ def test_succeeds_after_queue_drains(self, handler, input_queue):
         input_queue.get()
         result = asyncio.run(handler.end_worker(EmptyRequest()))
         assert isinstance(result, EmptyReturn)
+
+    def test_succeeds_when_only_rpc_acks_are_queued(self, handler, input_queue):
+        # An ack-only backlog carries no work (the worker never awaits its
+        # fire-and-forget acks), so it must not fail the kill -- failing here
+        # only makes region termination retry and loop e2e tests flaky.
+        input_queue.put(self.make_rpc_ack(command_id=1))
+        input_queue.put(self.make_rpc_ack(command_id=2))
+
+        result = asyncio.run(handler.end_worker(EmptyRequest()))
+
+        assert isinstance(result, EmptyReturn)
+
+    def test_fails_and_preserves_real_work_mixed_with_acks(self, handler, input_queue):
+        # Leniency is strictly ack-only: real work queued alongside an ack
+        # must fail the call AND survive in the queue for the retry to drain.
+        input_queue.put(self.make_rpc_ack(command_id=1))
+        work = self.make_control_message(0)
+        input_queue.put(work)
+
+        with pytest.raises(RuntimeError, match="unprocessed messages"):
+            asyncio.run(handler.end_worker(EmptyRequest()))
+
+        remaining = []
+        while not input_queue.is_empty():
+            remaining.append(input_queue.get())
+        assert work in remaining, (
+            "the queued control message must survive a failed EndWorker "
+            f"so the retry can drain it; queue held: {remaining}"
+        )

amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala

@@ -25,7 +25,10 @@ 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.rpc.workerservice.WorkerServiceGrpc.METHOD_QUERY_STATISTICS
 import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
   ActorCommandElement,
@@ -93,6 +96,15 @@ class EndHandlerSpec extends AnyFlatSpec {
     queue
   }
 
+  private def ackElement(sequenceNumber: Long, commandId: Long): DPInputQueueElement =
+    FIFOMessageElement(
+      WorkflowFIFOMessage(
+        ChannelIdentity(COORDINATOR, workerId, isControl = true),
+        sequenceNumber,
+        ReturnInvocation(commandId, EmptyReturn())
+      )
+    )
+
   "EndHandler" should "reply successfully when there are no unprocessed messages" in {
     val handler = createEndHandlerForQueue(new LinkedBlockingQueue[DPInputQueueElement]())
 
@@ -110,4 +122,29 @@ class EndHandlerSpec extends AnyFlatSpec {
 
     assertEndWorkerFails(handler)
   }
+
+  it should "reply successfully when only RPC acks (ReturnInvocations) are queued" in {
+    // Race at region termination: the coordinator's ReturnInvocation ack for
+    // this worker's own fire-and-forget RPC (e.g. workerExecutionCompleted)
+    // can still be in flight when EndWorker arrives. Such acks carry no work
+    // -- the worker never awaits them -- so an ack-only backlog must not fail
+    // the kill (it previously did, making region termination retry and loop
+    // e2e tests flaky).
+    val queue = new LinkedBlockingQueue[DPInputQueueElement]()
+    queue.put(ackElement(sequenceNumber = 0, commandId = 1))
+    queue.put(ackElement(sequenceNumber = 1, commandId = 2))
+    val handler = createEndHandlerForQueue(queue)
+
+    assert(await(handler.endWorker(EmptyRequest(), rpcContext)) == EmptyReturn())
+  }
+
+  it should "fail when an RPC ack is queued together with real work" in {
+    // Leniency is strictly ack-only: any non-ack message (here a control
+    // invocation) still fails the kill so the retry lets the worker drain it.
+    val queue = queueWithFifoControlMessage()
+    queue.put(ackElement(sequenceNumber = 1, commandId = 2))
+    val handler = createEndHandlerForQueue(queue)
+
+    assertEndWorkerFails(handler)
+  }
 }