From dc535e1b1f622fb09f73384b96c677232e57f40b Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Sat, 25 Jul 2026 01:01:57 -0700 Subject: [PATCH] fix(amber): tolerate in-flight RPC acks when EndWorker terminates a worker Loop e2e tests failed intermittently (~1 in 10 runs): region termination raced with the coordinator's ReturnInvocation acks for the worker's own fire-and-forget RPCs (workerExecutionCompleted / portCompleted). The very RPC that makes the coordinator decide to end the worker is the one whose ack is still in flight, so the race is inherent -- but the worker never awaits those acks, so an ack-only backlog carries no work. Loops amplify the race because every iteration terminates and re-executes regions, and the suite's single retry does not save runs where it strikes twice (same fail-fast as #5614's DataProcessingSpec flake). Scala EndHandler: succeed with a warning when everything queued is a ReturnInvocation; keep the fail-fast for anything else (control invocations, data, ECMs, actor commands) so the region execution manager's retry lets the worker drain real work. Python EndWorkerHandler: mirror the same semantics, and fix two latent defects -- the old handler consumed one message as a side effect of logging it (silently swallowing real work on the failure path) and its bare assert crashed on a two-ack backlog. Non-ack messages are now put back on the queue before failing. Tests: EndHandlerSpec gains ack-only (succeeds) and ack+work (still fails) cases; new test_end_worker_handler.py pins the same two behaviors plus that real work survives a failed EndWorker. --- .../handlers/control/end_worker_handler.py | 42 +++++- .../worker/promisehandlers/EndHandler.scala | 33 ++++- .../control/test_end_worker_handler.py | 126 ++++++++++++++++++ .../promisehandlers/EndHandlerSpec.scala | 39 +++++- 4 files changed, 228 insertions(+), 12 deletions(-) create mode 100644 amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py diff --git a/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py b/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py index 56baf1f345c..fd2ceb804e6 100644 --- a/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py +++ b/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py @@ -18,10 +18,13 @@ from loguru import logger from core.architecture.handlers.control.control_handler_base import ControlHandler +from core.models.internal_queue import DCMElement from core.util import IQueue +from core.util.proto import get_one_of from proto.org.apache.texera.amber.engine.architecture.rpc import ( EmptyReturn, EmptyRequest, + ReturnInvocation, ) @@ -37,13 +40,38 @@ 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. + # Ensure this is really the last message that carries work. 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). input_queue: IQueue = self.context.input_queue - if not input_queue.is_empty(): + pending = [] + while not input_queue.is_empty(): + pending.append(input_queue.get()) + if not pending: + # Now we can safely acknowledge that this worker can be terminated. + return EmptyReturn() + + 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 EndHandler before all messages are " - f"processed. Unprocessed messages: {input_queue.get()}" + f"Received EndHandler with only RPC acks left in the queue; " + f"proceeding with termination. Pending acks: {pending}" ) - assert input_queue.is_empty() - # 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 EndHandler before all messages are processed. " + f"Unprocessed messages: {pending}" + ) + raise RuntimeError("worker still has unprocessed messages") diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala index f91a28c627b..db10e2bb633 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala +++ b/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() } diff --git a/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py b/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py new file mode 100644 index 00000000000..402f13dc93f --- /dev/null +++ b/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py @@ -0,0 +1,126 @@ +# 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. + +"""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 (the old handler consumed one message as a side effect of logging 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 ( + AsyncRpcContext, + ControlInvocation, + ControlRequest, + ControlReturn, + EmptyRequest, + EmptyReturn, + ReturnInvocation, +) +from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2 + + +_CONTROL_CHANNEL = ChannelIdentity( + from_worker_id=ActorVirtualIdentity(name="COORDINATOR"), + to_worker_id=ActorVirtualIdentity(name="Worker:WF1-test-op-main-0"), + is_control=True, +) + + +def _ack_element(command_id: int) -> DCMElement: + return DCMElement( + tag=_CONTROL_CHANNEL, + payload=set_one_of( + DirectControlMessagePayloadV2, + ReturnInvocation( + command_id=command_id, + return_value=set_one_of(ControlReturn, EmptyReturn()), + ), + ), + ) + + +def _control_invocation_element(command_id: int) -> DCMElement: + return DCMElement( + tag=_CONTROL_CHANNEL, + payload=set_one_of( + DirectControlMessagePayloadV2, + ControlInvocation( + method_name="PauseWorker", + command=set_one_of(ControlRequest, EmptyRequest()), + context=AsyncRpcContext(), + command_id=command_id, + ), + ), + ) + + +def _make_handler(input_queue: InternalQueue) -> EndWorkerHandler: + return EndWorkerHandler(SimpleNamespace(input_queue=input_queue)) + + +class TestEndWorkerHandler: + def test_succeeds_on_empty_queue(self): + handler = _make_handler(InternalQueue()) + result = asyncio.run(handler.end_worker(EmptyRequest())) + assert isinstance(result, EmptyReturn) + + def test_succeeds_when_only_rpc_acks_are_queued(self): + # Two queued acks (more than the old handler's accidental + # consume-one-in-the-log leniency could absorb) must not fail the kill. + queue = InternalQueue() + queue.put(_ack_element(command_id=1)) + queue.put(_ack_element(command_id=2)) + handler = _make_handler(queue) + + result = asyncio.run(handler.end_worker(EmptyRequest())) + + assert isinstance(result, EmptyReturn) + + def test_fails_and_preserves_real_work(self): + # A queued control invocation is real work: the handler must fail so + # the coordinator retries, and the message must STILL be in the queue + # afterwards (the old handler swallowed one message while logging it). + queue = InternalQueue() + queue.put(_ack_element(command_id=1)) + work = _control_invocation_element(command_id=2) + queue.put(work) + handler = _make_handler(queue) + + with pytest.raises(Exception, match="unprocessed messages"): + asyncio.run(handler.end_worker(EmptyRequest())) + + remaining = [] + while not queue.is_empty(): + remaining.append(queue.get()) + assert work in remaining, ( + "the queued control invocation must survive a failed EndWorker " + f"so the retry can drain it; queue held: {remaining}" + ) diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala index 3a3c0ecb9ed..c3577bffd0f 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala +++ b/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) + } }