From b352eb101c27978d88afd91ed1e26677f907cfac Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Fri, 17 Jul 2026 20:57:23 -0700 Subject: [PATCH 1/3] fix+test(pyamber): stop EndWorker from consuming straggler messages EndWorkerHandler logged its "unprocessed messages" warning through input_queue.get(), a destructive read: with exactly one straggler the message was silently dropped and EndWorker acknowledged success; with more, one message was still destroyed and the RPC failed with a bare AssertionError, so each coordinator retry that hit the guard ate another queued message. Mirror the Scala EndHandler instead: log via a non-destructive peek() (added to InternalQueue) and raise, so the RPC layer replies with a ControlError and the coordinator's terminateWorkersWithRetry retries once the queue has drained, with no messages lost. Tests live under amber/src/test/python per the current layout: a new test_end_worker_handler.py (the Python analogue of EndHandlerSpec) and peek() cases appended to the existing test_internal_queue.py. --- .../handlers/control/end_worker_handler.py | 11 ++- .../main/python/core/models/internal_queue.py | 10 +- .../control/test_end_worker_handler.py | 97 +++++++++++++++++++ .../python/core/models/test_internal_queue.py | 31 ++++++ 4 files changed, 144 insertions(+), 5 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..f7325458e45 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,7 +18,7 @@ from loguru import logger from core.architecture.handlers.control.control_handler_base import ControlHandler -from core.util import IQueue +from core.models.internal_queue import InternalQueue from proto.org.apache.texera.amber.engine.architecture.rpc import ( EmptyReturn, EmptyRequest, @@ -38,12 +38,15 @@ async def end_worker(self, req: EmptyRequest) -> EmptyReturn: of all the control messages. """ # Ensure this is really the last message. - input_queue: IQueue = self.context.input_queue + input_queue: InternalQueue = self.context.input_queue if not input_queue.is_empty(): logger.warning( f"Received EndHandler before all messages are " - f"processed. Unprocessed messages: {input_queue.get()}" + f"processed. Unprocessed messages: {input_queue.peek()}" ) - assert input_queue.is_empty() + # Fail this RPC (the counterpart of the Scala EndHandler's + # Future.exception) so the coordinator retries 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() diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index abc1793ff6c..20f991baef6 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -20,7 +20,7 @@ from dataclasses import dataclass from enum import Enum from threading import RLock -from typing import TypeVar, Set +from typing import Optional, TypeVar, Set from core.models.internal_marker import InternalMarker from core.models.payload import DataPayload @@ -74,6 +74,14 @@ def is_empty(self, key=None) -> bool: def get(self) -> T: return self._queue.get() + def peek(self) -> Optional[T]: + """ + Peek the element that get() would return next, without removing it. + + :return: Optional[T], the next available element, or None if empty. + """ + return self._queue.peek() + def put(self, item: T) -> None: if isinstance(item, InternalQueueElement): if item.tag not in self._queue_ids: 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..6472738b010 --- /dev/null +++ b/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py @@ -0,0 +1,97 @@ +# 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. + +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 proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity +from proto.org.apache.texera.amber.engine.architecture.rpc import ( + EmptyRequest, + EmptyReturn, +) +from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2 + + +class TestEndWorkerHandler: + @pytest.fixture + def input_queue(self): + return InternalQueue() + + @pytest.fixture + def handler(self, input_queue): + # ControlHandler.__init__ just stashes context; bypass the protobuf + # base class' __init__ by constructing via __new__. + instance = EndWorkerHandler.__new__(EndWorkerHandler) + instance.context = SimpleNamespace(input_queue=input_queue) + return instance + + @staticmethod + def make_control_message(seq: int) -> DCMElement: + channel = ChannelIdentity( + ActorVirtualIdentity(name="CONTROLLER"), + ActorVirtualIdentity(name=f"worker-{seq}"), + is_control=True, + ) + return DCMElement(tag=channel, payload=DirectControlMessagePayloadV2()) + + def test_acknowledges_end_worker_on_empty_queue(self, handler): + result = asyncio.run(handler.end_worker(EmptyRequest())) + assert isinstance(result, EmptyReturn) + + def test_rejects_end_worker_with_one_unprocessed_message( + self, handler, input_queue + ): + # The controller resolves endWorker's reply as a failed future and + # retries after the queue drains (terminateWorkersWithRetry), so a + # straggler message must fail the call — not be dropped with a + # successful acknowledgement. + input_queue.put(self.make_control_message(0)) + with pytest.raises(RuntimeError, match="unprocessed messages"): + asyncio.run(handler.end_worker(EmptyRequest())) + + @pytest.mark.timeout(2) + def test_does_not_consume_the_unprocessed_message(self, handler, input_queue): + straggler = self.make_control_message(0) + input_queue.put(straggler) + with pytest.raises(RuntimeError): + asyncio.run(handler.end_worker(EmptyRequest())) + # The straggler must survive the failed call so the main loop can + # still process it before the controller's retry. + assert input_queue.size() == 1 + assert input_queue.get() is straggler + + def test_keeps_all_messages_with_multiple_stragglers(self, handler, input_queue): + input_queue.put(self.make_control_message(0)) + input_queue.put(self.make_control_message(1)) + with pytest.raises(RuntimeError): + asyncio.run(handler.end_worker(EmptyRequest())) + assert input_queue.size() == 2 + + @pytest.mark.timeout(2) + def test_succeeds_after_queue_drains(self, handler, input_queue): + # The retry protocol end-to-end: fail while a message is pending, + # acknowledge once the queue has drained. + input_queue.put(self.make_control_message(0)) + with pytest.raises(RuntimeError): + asyncio.run(handler.end_worker(EmptyRequest())) + input_queue.get() + result = asyncio.run(handler.end_worker(EmptyRequest())) + assert isinstance(result, EmptyReturn) diff --git a/amber/src/test/python/core/models/test_internal_queue.py b/amber/src/test/python/core/models/test_internal_queue.py index 663f95d89a8..90504661f67 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -342,6 +342,37 @@ def test_it_tracks_in_mem_size_of_data_channels_only( queue.get() assert queue.in_mem_size() == 0 + def test_peek_on_empty_queue_returns_none(self, queue): + assert queue.peek() is None + + @pytest.mark.timeout(2) + def test_peek_does_not_remove_the_element(self, queue, control_channel): + dcm = self.dcm_element(control_channel) + queue.put(dcm) + # repeated peeks see the same element and nothing is consumed + assert queue.peek() is dcm + assert queue.peek() is dcm + assert queue.size() == 1 + assert queue.get() is dcm + assert queue.peek() is None + + @pytest.mark.timeout(2) + def test_peek_is_consistent_with_get_across_sub_queues( + self, queue, control_channel, data_channel + ): + dcm = self.dcm_element(control_channel) + data = self.data_element(data_channel) + queue.put(dcm) + queue.put(data) + # peek must surface exactly the element the next get() returns, + # whatever order the sub-queue selection strategy imposes + first_peeked = queue.peek() + assert queue.get() is first_peeked + second_peeked = queue.peek() + assert second_peeked is not first_peeked + assert queue.get() is second_peeked + assert queue.peek() is None + @pytest.mark.timeout(2) def test_it_can_disable_and_enable_a_single_data_channel( self, queue, control_channel, data_channel, second_data_channel From 87d618b0d3721e4a7b2ab1f4b5175fd690c59525 Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Sat, 18 Jul 2026 15:55:39 -0700 Subject: [PATCH 2/3] =?UTF-8?q?refactor(pyamber):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20drop=20InternalQueue.peek,=20keep=20IQueue?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback: keep InternalQueue non-peekable (parity with queue.Queue) and keep the handler typed as IQueue. The EndWorker guard does not need the next message's content — it logs the pending count via the IQueue interface (len()/is_empty()) and raises, so no InternalQueue internals are exposed. Also fix the warning text to say EndWorker rather than EndHandler. - end_worker_handler.py: revert annotation to IQueue; the warning now references EndWorker and reports the pending count; the guard no longer calls peek() or get(). - internal_queue.py: remove InternalQueue.peek() and the now-unused Optional import. - test_internal_queue.py: drop the three peek tests. --- .../handlers/control/end_worker_handler.py | 9 +++--- .../main/python/core/models/internal_queue.py | 10 +----- .../python/core/models/test_internal_queue.py | 31 ------------------- 3 files changed, 6 insertions(+), 44 deletions(-) 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 f7325458e45..2ef3906c568 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,7 +18,7 @@ from loguru import logger from core.architecture.handlers.control.control_handler_base import ControlHandler -from core.models.internal_queue import InternalQueue +from core.util import IQueue from proto.org.apache.texera.amber.engine.architecture.rpc import ( EmptyReturn, EmptyRequest, @@ -38,11 +38,12 @@ async def end_worker(self, req: EmptyRequest) -> EmptyReturn: of all the control messages. """ # Ensure this is really the last message. - input_queue: InternalQueue = self.context.input_queue + input_queue: IQueue = self.context.input_queue if not input_queue.is_empty(): logger.warning( - f"Received EndHandler before all messages are " - f"processed. Unprocessed messages: {input_queue.peek()}" + f"Received EndWorker before all {len(input_queue)} queued " + f"message(s) were processed; failing the RPC so the coordinator " + f"retries once the queue has drained." ) # Fail this RPC (the counterpart of the Scala EndHandler's # Future.exception) so the coordinator retries once the queue diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index 20f991baef6..abc1793ff6c 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -20,7 +20,7 @@ from dataclasses import dataclass from enum import Enum from threading import RLock -from typing import Optional, TypeVar, Set +from typing import TypeVar, Set from core.models.internal_marker import InternalMarker from core.models.payload import DataPayload @@ -74,14 +74,6 @@ def is_empty(self, key=None) -> bool: def get(self) -> T: return self._queue.get() - def peek(self) -> Optional[T]: - """ - Peek the element that get() would return next, without removing it. - - :return: Optional[T], the next available element, or None if empty. - """ - return self._queue.peek() - def put(self, item: T) -> None: if isinstance(item, InternalQueueElement): if item.tag not in self._queue_ids: diff --git a/amber/src/test/python/core/models/test_internal_queue.py b/amber/src/test/python/core/models/test_internal_queue.py index 90504661f67..663f95d89a8 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -342,37 +342,6 @@ def test_it_tracks_in_mem_size_of_data_channels_only( queue.get() assert queue.in_mem_size() == 0 - def test_peek_on_empty_queue_returns_none(self, queue): - assert queue.peek() is None - - @pytest.mark.timeout(2) - def test_peek_does_not_remove_the_element(self, queue, control_channel): - dcm = self.dcm_element(control_channel) - queue.put(dcm) - # repeated peeks see the same element and nothing is consumed - assert queue.peek() is dcm - assert queue.peek() is dcm - assert queue.size() == 1 - assert queue.get() is dcm - assert queue.peek() is None - - @pytest.mark.timeout(2) - def test_peek_is_consistent_with_get_across_sub_queues( - self, queue, control_channel, data_channel - ): - dcm = self.dcm_element(control_channel) - data = self.data_element(data_channel) - queue.put(dcm) - queue.put(data) - # peek must surface exactly the element the next get() returns, - # whatever order the sub-queue selection strategy imposes - first_peeked = queue.peek() - assert queue.get() is first_peeked - second_peeked = queue.peek() - assert second_peeked is not first_peeked - assert queue.get() is second_peeked - assert queue.peek() is None - @pytest.mark.timeout(2) def test_it_can_disable_and_enable_a_single_data_channel( self, queue, control_channel, data_channel, second_data_channel From bbbb779150ef5c63a0afb86aa3167fde02e2ff32 Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Sun, 19 Jul 2026 23:01:26 -0700 Subject: [PATCH 3/3] fix(pyamber): type EndWorker queue as InternalQueue and read the count once The local was annotated IQueue (which declares neither __len__ nor size), so len(input_queue) tripped static analysis. Type it as the actual InternalQueue and read queued_count = size() once, branching on it instead of the redundant is_empty()+len() double-check. --- .../handlers/control/end_worker_handler.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) 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 2ef3906c568..932bc7174ff 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,7 +18,7 @@ from loguru import logger from core.architecture.handlers.control.control_handler_base import ControlHandler -from core.util import IQueue +from core.models.internal_queue import InternalQueue from proto.org.apache.texera.amber.engine.architecture.rpc import ( EmptyReturn, EmptyRequest, @@ -37,11 +37,13 @@ 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. - input_queue: IQueue = self.context.input_queue - if not input_queue.is_empty(): + # 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. + input_queue: InternalQueue = self.context.input_queue + queued_count = input_queue.size() + if queued_count > 0: logger.warning( - f"Received EndWorker before all {len(input_queue)} queued " + f"Received EndWorker before all {queued_count} queued " f"message(s) were processed; failing the RPC so the coordinator " f"retries once the queue has drained." )