From bcb0e5636f1c1e770f8c3bed291bddec86961b23 Mon Sep 17 00:00:00 2001 From: eugenegujing Date: Tue, 21 Jul 2026 14:59:25 -0700 Subject: [PATCH 1/6] fix(pyamber): disable data sub-queues registered after disable_data Fix: when put() registers a new data channel while _queue_state is non-empty, the new sub-queue now starts disabled (done under the same lock as disable_data/enable_data, and before the first element is enqueued). Control channels are never disabled, and enable_data already re-enables such channels on resume. Also removes the unreachable InternalMarker branch in put() and merges the two identical dispatch branches. Adds 8 regression tests to test_internal_queue.py covering the late-channel leak, release on resume, stacked disables, control channels staying unblocked, existing baselines, and two threaded stress tests. --- .../main/python/core/models/internal_queue.py | 19 +- .../python/core/models/test_internal_queue.py | 195 ++++++++++++++++++ 2 files changed, 208 insertions(+), 6 deletions(-) diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index abc1793ff6c..bed02dd88a8 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -22,7 +22,6 @@ from threading import RLock from typing import TypeVar, Set -from core.models.internal_marker import InternalMarker from core.models.payload import DataPayload from core.util.customized_queue.linked_blocking_multi_queue import ( LinkedBlockingMultiQueue, @@ -77,11 +76,19 @@ def get(self) -> T: def put(self, item: T) -> None: if isinstance(item, InternalQueueElement): if item.tag not in self._queue_ids: - self._queue.add_sub_queue(item.tag, 1 if item.tag.is_control else 2) - self._queue_ids.add(item.tag) - if isinstance(item, (DataElement, InternalMarker, ECMElement)): - self._queue.put(item.tag, item) - elif isinstance(item, DCMElement): + # registration must not interleave with disable_data/enable_data + with self._lock: + if item.tag not in self._queue_ids: + self._queue.add_sub_queue( + item.tag, 1 if item.tag.is_control else 2 + ) + # while data is disabled, a new data sub-queue must + # start disabled too (before its first element is + # enqueued), or it would leak data during the pause + if not item.tag.is_control and self._queue_state: + self._queue.disable(item.tag) + self._queue_ids.add(item.tag) + if isinstance(item, (DataElement, ECMElement, DCMElement)): self._queue.put(item.tag, item) else: raise ValueError(f"item {item} is not recognized by internal queue") 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..94033aaacfb 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import threading from dataclasses import dataclass import pytest @@ -364,3 +365,197 @@ def test_it_can_disable_and_enable_a_single_data_channel( queue.enable(data_channel) assert queue.get() is blocked assert queue.is_empty() + + # Regression tests below: data channels whose sub-queue is created lazily + # (on the channel's first put) AFTER disable_data has been called must + # come up disabled — a paused or backpressured worker must not be able to + # dequeue data from them, and is_data_enabled() must not flip back to + # True just because a new channel delivered its first message. + + def test_channel_registered_after_disable_comes_up_disabled( + self, queue, data_channel + ): + # the main regression: disable first, then the channel's FIRST put + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + queue.put(self.data_element(data_channel)) + assert not queue.is_data_enabled() + # the element stays queued but must not be dequeuable + assert queue.size_data() == 1 + assert queue._queue.peek() is None + + @pytest.mark.timeout(2) + def test_enable_data_releases_a_channel_registered_mid_disable( + self, queue, data_channel + ): + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + data = self.data_element(data_channel) + queue.put(data) + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert queue.is_data_enabled() + assert queue._queue.peek() is data + assert queue.get() is data + assert queue.is_empty() + + @pytest.mark.timeout(2) + def test_channel_registered_under_stacked_disables_stays_disabled( + self, queue, data_channel + ): + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) + data = self.data_element(data_channel) + queue.put(data) + # releasing only one of the two reasons must not open the channel + assert not queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert not queue.is_data_enabled() + assert queue._queue.peek() is None + # releasing the remaining reason makes the element dequeuable + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) + assert queue.is_data_enabled() + assert queue.get() is data + + @pytest.mark.timeout(2) + def test_control_channel_registered_mid_disable_is_never_blocked( + self, queue, control_channel, data_channel + ): + # register a data channel first so is_data_enabled() is meaningful + queue.put(self.data_element(data_channel)) + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + # the control channel's FIRST put happens while data is disabled + dcm = self.dcm_element(control_channel) + queue.put(dcm) + # control must flow immediately, and data must stay disabled + assert queue._queue.peek() is dcm + assert queue.get() is dcm + assert not queue.is_data_enabled() + assert queue.size_data() == 1 + + @pytest.mark.timeout(2) + def test_channel_registered_before_disable_is_disabled_and_reenabled( + self, queue, data_channel + ): + # baseline: the pre-existing behavior for eagerly-registered channels + data = self.data_element(data_channel) + queue.put(data) + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert not queue.is_data_enabled() + assert queue._queue.peek() is None + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert queue.is_data_enabled() + assert queue.get() is data + + @pytest.mark.timeout(2) + def test_channel_registered_while_enabled_behaves_normally( + self, queue, second_data_channel + ): + data = self.data_element(second_data_channel) + queue.put(data) + assert queue.is_data_enabled() + assert queue._queue.peek() is data + assert queue.get() is data + assert queue.is_empty() + + @pytest.mark.timeout(10) + def test_concurrent_first_time_puts_while_toggling_disable(self, queue): + # concurrency smoke test: receiver threads register brand-new data + # channels while the DP thread toggles disable_data/enable_data; + # only the final state is asserted, deterministically. + n_threads = 8 + elements_per_thread = 25 + start_barrier = threading.Barrier(n_threads + 1) + errors = [] + + def producer(thread_index): + channel = ChannelIdentity( + ActorVirtualIdentity(f"upstream_{thread_index}"), + ActorVirtualIdentity("dummy_worker_id"), + False, + ) + try: + start_barrier.wait() + for _ in range(elements_per_thread): + queue.put(self.data_element(channel)) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + threads = [ + threading.Thread(target=producer, args=(i,)) for i in range(n_threads) + ] + for thread in threads: + thread.start() + start_barrier.wait() + for _ in range(5): + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + for thread in threads: + thread.join() + # one last full cycle after all puts settled: every channel must be + # disabled, then re-enabled with its count added back exactly once + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert not queue.is_data_enabled() + assert queue._queue.peek() is None + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + + assert not errors + total = n_threads * elements_per_thread + assert queue.is_data_enabled() + assert queue.size_data() == total + # size() is the getable total_count: a mismatch with size_data() + # means an element was double-counted or lost by a toggle race + assert queue.size() == total + drained = 0 + while queue._queue.peek() is not None: + queue.get() + drained += 1 + assert drained == total + assert queue.is_empty() + + @pytest.mark.timeout(30) + def test_concurrent_first_time_puts_racing_disable_enable_toggles(self): + # Receiver threads deliver first-ever messages on distinct new data + # channels while the DP-thread side toggles pause on and off. Only + # the final state is asserted (deterministic): with the queue left + # disabled, nothing is dequeuable; after the final enable_data every + # element is dequeuable exactly once, so total_count stayed exact. + threads, channels_per_thread, toggles = 4, 10, 10 + for _ in range(5): + queue = InternalQueue() + errors = [] + start = threading.Barrier(threads + 1) + + def producer(thread_id): + try: + start.wait() + for i in range(channels_per_thread): + channel = ChannelIdentity( + ActorVirtualIdentity(f"upstream-{thread_id}-{i}"), + ActorVirtualIdentity("dummy_worker_id"), + False, + ) + queue.put(self.data_element(channel)) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + producers = [ + threading.Thread(target=producer, args=(t,)) for t in range(threads) + ] + for producer_thread in producers: + producer_thread.start() + start.wait() + for _ in range(toggles): + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + for producer_thread in producers: + producer_thread.join() + + assert errors == [] + total = threads * channels_per_thread + assert queue.size_data() == total + assert queue._queue.peek() is None + assert not queue.is_data_enabled() + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + dequeued = 0 + while queue._queue.peek() is not None: + queue.get() + dequeued += 1 + assert dequeued == total From 205b9b78c3b918589e2f06934cd4275f7cd789a1 Mon Sep 17 00:00:00 2001 From: eugenegujing Date: Tue, 21 Jul 2026 15:56:49 -0700 Subject: [PATCH 2/6] chore: retrigger CI From 9dbbec219ea59a650aed63da74507aa4615c0ded Mon Sep 17 00:00:00 2001 From: eugenegujing Date: Mon, 10 Aug 2026 15:55:23 -0700 Subject: [PATCH 3/6] fix(pyamber): iterate a snapshot of _queue_ids in category query methods - The six category query methods (is_control_empty, is_data_empty, size_control, size_data, in_mem_size, is_data_enabled) now iterate a tuple() snapshot of _queue_ids instead of the live set, so a first-time put() registering a new channel concurrently can no longer raise "Set changed size during iteration" and kill the DP thread. - Also reword the put() comment to mention backpressure alongside pause. - Add 6 parametrized regression tests and a threaded stress test; all fail without the snapshot fix. --- .../main/python/core/models/internal_queue.py | 16 +-- .../python/core/models/test_internal_queue.py | 100 ++++++++++++++++++ 2 files changed, 109 insertions(+), 7 deletions(-) diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index bed02dd88a8..d90833473cb 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -84,7 +84,7 @@ def put(self, item: T) -> None: ) # while data is disabled, a new data sub-queue must # start disabled too (before its first element is - # enqueued), or it would leak data during the pause + # enqueued), or it would leak data during pause/backpressure if not item.tag.is_control and self._queue_state: self._queue.disable(item.tag) self._queue_ids.add(item.tag) @@ -102,16 +102,18 @@ def enable(self, channel_id: ChannelIdentity) -> None: self._queue.enable(channel_id) def is_control_empty(self) -> bool: + # snapshot: put() may add channels concurrently, and iterating the + # live set while it grows raises RuntimeError return all( self.is_empty(queue_id) - for queue_id in self._queue_ids + for queue_id in tuple(self._queue_ids) if queue_id.is_control ) def is_data_empty(self) -> bool: return all( self.is_empty(queue_id) - for queue_id in self._queue_ids + for queue_id in tuple(self._queue_ids) if not queue_id.is_control ) @@ -124,14 +126,14 @@ def size(self) -> int: def size_control(self) -> int: return sum( self._queue.size(queue_id) - for queue_id in self._queue_ids + for queue_id in tuple(self._queue_ids) if queue_id.is_control ) def size_data(self) -> int: return sum( self._queue.size(queue_id) - for queue_id in self._queue_ids + for queue_id in tuple(self._queue_ids) if not queue_id.is_control ) @@ -156,13 +158,13 @@ def disable_data(self, disable_type: DisableType) -> None: def in_mem_size(self) -> int: return sum( self._queue.in_mem_size(queue_id) - for queue_id in self._queue_ids + for queue_id in tuple(self._queue_ids) if not queue_id.is_control ) def is_data_enabled(self) -> bool: return any( self._queue.is_enabled(queue_id) - for queue_id in self._queue_ids + for queue_id in tuple(self._queue_ids) if not queue_id.is_control ) 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 94033aaacfb..c4a0736c69e 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -559,3 +559,103 @@ def producer(thread_id): queue.get() dequeued += 1 assert dequeued == total + + # Regression tests below: the per-category query methods iterate + # _queue_ids, which put() grows on a channel's first message. Iterating + # the live set while another thread grows it raises RuntimeError + # ("Set changed size during iteration"), killing the calling thread — + # e.g. the DP thread polling is_data_enabled() in the main loop — so the + # queries must iterate a snapshot of the set instead. + + @pytest.mark.parametrize( + "query, expected", + [ + ("is_control_empty", True), + ("is_data_empty", True), + ("size_control", 0), + ("size_data", 0), + ("in_mem_size", 0), + ("is_data_enabled", False), + ], + ) + def test_queries_survive_a_channel_registration_mid_iteration( + self, query, expected + ): + # A key whose is_control access (evaluated inside the query's + # iteration over _queue_ids) delivers the first-ever message of a + # brand-new data channel, interleaving a registration into the + # iteration exactly like a concurrent Flight reader thread would. + queue = InternalQueue() + outer = self + + class RegisteringKey: + def __init__(self): + self.fired = 0 + + @property + def is_control(self): + self.fired += 1 + late_channel = ChannelIdentity( + ActorVirtualIdentity(f"late_upstream_{self.fired}"), + ActorVirtualIdentity("dummy_worker_id"), + False, + ) + queue.put(outer.data_element(late_channel)) + return False + + registering_key = RegisteringKey() + queue._queue.add_sub_queue(registering_key, 2) + # keep this sub-queue disabled and empty so no query short-circuits + # on its yielded value: each one must advance the iteration past the + # mid-iteration registration, which raises RuntimeError on the live + # set and must not raise on a snapshot + queue._queue.disable(registering_key) + queue._queue_ids.add(registering_key) + + assert getattr(queue, query)() == expected + assert registering_key.fired == 1 + + @pytest.mark.timeout(20) + def test_queries_survive_concurrent_first_time_registrations(self): + # realistic race: reader threads deliver first-ever messages on new + # data channels while the DP-thread side polls the category queries, + # as main_loop's _check_and_process_control does + queue = InternalQueue() + n_threads, channels_per_thread = 4, 200 + start_barrier = threading.Barrier(n_threads + 1) + errors = [] + + def producer(thread_id): + try: + start_barrier.wait() + for i in range(channels_per_thread): + channel = ChannelIdentity( + ActorVirtualIdentity(f"upstream_{thread_id}_{i}"), + ActorVirtualIdentity("dummy_worker_id"), + False, + ) + queue.put(self.data_element(channel)) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + producers = [ + threading.Thread(target=producer, args=(t,)) for t in range(n_threads) + ] + for producer_thread in producers: + producer_thread.start() + start_barrier.wait() + # a RuntimeError from any query fails the test right here + while any(producer_thread.is_alive() for producer_thread in producers): + queue.is_control_empty() + queue.is_data_empty() + queue.size_control() + queue.size_data() + queue.in_mem_size() + queue.is_data_enabled() + for producer_thread in producers: + producer_thread.join() + + assert errors == [] + assert queue.size_data() == n_threads * channels_per_thread + assert queue.size_control() == 0 + assert queue.is_data_enabled() From 92f4dc3e6202653043c8f51e5ed3aa738ef18aa6 Mon Sep 17 00:00:00 2001 From: eugenegujing Date: Tue, 11 Aug 2026 17:00:45 -0700 Subject: [PATCH 4/6] refactor(pyamber): move the _queue_ids snapshot into helper methods Per review, the six category query methods now go through _control_queue_ids() / _data_queue_ids(), which take the tuple() snapshot and filter it, so the snapshot is enforced by the API instead of being a convention each query method has to remember. The snapshot rationale comment moves to the helper docstring. The helpers copy the live set first and filter the private copy, so no user-level code runs while the live set is being iterated. Also reword a test comment per review ("Registers a key whose ..."). --- .../main/python/core/models/internal_queue.py | 51 ++++++++----------- .../python/core/models/test_internal_queue.py | 2 +- 2 files changed, 23 insertions(+), 30 deletions(-) diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index d90833473cb..5900db1cf10 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 Tuple, TypeVar, Set from core.models.payload import DataPayload from core.util.customized_queue.linked_blocking_multi_queue import ( @@ -101,21 +101,26 @@ def disable(self, channel_id: ChannelIdentity) -> None: def enable(self, channel_id: ChannelIdentity) -> None: self._queue.enable(channel_id) + def _control_queue_ids(self) -> Tuple[ChannelIdentity, ...]: + """Snapshot of the registered control channels. + + put() can grow _queue_ids from another thread, and iterating the + live set while it grows raises RuntimeError, so queries must iterate + a snapshot taken through these helpers. + """ + snapshot = tuple(self._queue_ids) + return tuple(queue_id for queue_id in snapshot if queue_id.is_control) + + def _data_queue_ids(self) -> Tuple[ChannelIdentity, ...]: + """Snapshot of the registered data channels; see _control_queue_ids.""" + snapshot = tuple(self._queue_ids) + return tuple(queue_id for queue_id in snapshot if not queue_id.is_control) + def is_control_empty(self) -> bool: - # snapshot: put() may add channels concurrently, and iterating the - # live set while it grows raises RuntimeError - return all( - self.is_empty(queue_id) - for queue_id in tuple(self._queue_ids) - if queue_id.is_control - ) + return all(self.is_empty(queue_id) for queue_id in self._control_queue_ids()) def is_data_empty(self) -> bool: - return all( - self.is_empty(queue_id) - for queue_id in tuple(self._queue_ids) - if not queue_id.is_control - ) + return all(self.is_empty(queue_id) for queue_id in self._data_queue_ids()) def __len__(self) -> int: return self.size() @@ -124,18 +129,10 @@ def size(self) -> int: return self._queue.size() def size_control(self) -> int: - return sum( - self._queue.size(queue_id) - for queue_id in tuple(self._queue_ids) - if queue_id.is_control - ) + return sum(self._queue.size(queue_id) for queue_id in self._control_queue_ids()) def size_data(self) -> int: - return sum( - self._queue.size(queue_id) - for queue_id in tuple(self._queue_ids) - if not queue_id.is_control - ) + return sum(self._queue.size(queue_id) for queue_id in self._data_queue_ids()) def enable_data(self, disable_type: DisableType) -> bool: with self._lock: @@ -157,14 +154,10 @@ def disable_data(self, disable_type: DisableType) -> None: def in_mem_size(self) -> int: return sum( - self._queue.in_mem_size(queue_id) - for queue_id in tuple(self._queue_ids) - if not queue_id.is_control + self._queue.in_mem_size(queue_id) for queue_id in self._data_queue_ids() ) def is_data_enabled(self) -> bool: return any( - self._queue.is_enabled(queue_id) - for queue_id in tuple(self._queue_ids) - if not queue_id.is_control + self._queue.is_enabled(queue_id) for queue_id in self._data_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 c4a0736c69e..1583d7c513d 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -581,7 +581,7 @@ def producer(thread_id): def test_queries_survive_a_channel_registration_mid_iteration( self, query, expected ): - # A key whose is_control access (evaluated inside the query's + # Registers a key whose is_control access (evaluated inside the query's # iteration over _queue_ids) delivers the first-ever message of a # brand-new data channel, interleaving a registration into the # iteration exactly like a concurrent Flight reader thread would. From 1a65167ba49f280df7cfeaf195adb6c87a3689b7 Mon Sep 17 00:00:00 2001 From: eugenegujing Date: Thu, 13 Aug 2026 16:41:07 -0700 Subject: [PATCH 5/6] fix(pyamber): withhold leaked data at dequeue instead of disabling new channels - Registering a data channel disabled during a disable window also blocked ECMs, which ride data channels, so reconfiguring a paused worker never completed. - Channels now register enabled as on main and get() withholds instead: a DataElement taken while _queue_state is non-empty has its channel disabled and is pushed back to its sub-queue's head, for enable_data to release. - Adds put_first to LinkedBlockingMultiQueue for the push-back, reports is_data_enabled() as False while any disable reason is active, and rewrites the regression tests for the new semantics. --- .../main/python/core/models/internal_queue.py | 49 ++- .../linked_blocking_multi_queue.py | 41 +++ .../python/core/models/test_internal_queue.py | 343 +++++++++++++++--- 3 files changed, 383 insertions(+), 50 deletions(-) diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index 5900db1cf10..8f3121d0bff 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -71,22 +71,52 @@ def is_empty(self, key=None) -> bool: return self._queue.is_empty(key) def get(self) -> T: - return self._queue.get() + """Blocking get of the next available element. + + Data channels register enabled even during a disable window, because + ECMs ride data channels and one swallowed by a channel that came up + disabled would never be acked. A DataElement arriving here during + such a window is withheld instead: its channel is closed and the + element goes back to its sub-queue's head, for enable_data() to + release. An ECM queued behind it on the same channel is therefore + delayed until resume, which is unavoidable without unbounded + buffering, and is what main does for channels disable_data() closed. + + Assumes a single consumer; only PauseManager and BackpressureHandler + toggle the disable state, and only on the input queue. + """ + while True: + item = self._queue.get() + # a control-tagged DataElement cannot occur today, but the tag is + # wire-derived: withholding one would close a control sub-queue, + # which enable_data() never reopens + if ( + not isinstance(item, DataElement) + or item.tag.is_control + or not self._queue_state + ): + return item + with self._lock: + # enable_data() may have cleared the last reason since the + # check above; closing the channel now would strand it + if not self._queue_state: + return item + # disable first: the element must never be dequeuable in between + self._queue.disable(item.tag) + self._queue.put_first(item.tag, item) def put(self, item: T) -> None: if isinstance(item, InternalQueueElement): if item.tag not in self._queue_ids: - # registration must not interleave with disable_data/enable_data + # both the lock and the re-check are load-bearing: + # disable_data/enable_data iterate _queue_ids live, and a + # second add_sub_queue for the same channel would replace its + # sub-queue with an empty one, dropping whatever it holds with self._lock: if item.tag not in self._queue_ids: self._queue.add_sub_queue( item.tag, 1 if item.tag.is_control else 2 ) - # while data is disabled, a new data sub-queue must - # start disabled too (before its first element is - # enqueued), or it would leak data during pause/backpressure - if not item.tag.is_control and self._queue_state: - self._queue.disable(item.tag) self._queue_ids.add(item.tag) if isinstance(item, (DataElement, ECMElement, DCMElement)): self._queue.put(item.tag, item) @@ -158,6 +188,11 @@ def in_mem_size(self) -> int: ) def is_data_enabled(self) -> bool: + # channels registered mid-disable come up enabled (see get()), so + # per-channel state alone would report data as enabled during a pause + # and let main_loop's wait-loop exit and resume processing + if self._queue_state: + return False return any( self._queue.is_enabled(queue_id) for queue_id in self._data_queue_ids() ) diff --git a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py index fcf175d6bcd..48036757599 100644 --- a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py +++ b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py @@ -106,6 +106,14 @@ def enqueue(self, node: LinkedBlockingMultiQueue.Node[T]) -> None: self.last = node self.in_mem_size.inc(node.in_mem_size) + def enqueue_first(self, node: LinkedBlockingMultiQueue.Node[T]) -> None: + # head is a sentinel, so the first real element is head.next + node.next = self.head.next + self.head.next = node + if self.last is self.head: + self.last = node + self.in_mem_size.inc(node.in_mem_size) + def dequeue(self) -> T: assert self.size() > 0 h = self.head @@ -149,6 +157,26 @@ def put(self, obj: T) -> None: if old_size == 0: self.owner._signal_not_empty() + def put_first(self, obj: T) -> None: + # Same accounting as put(), but the element goes to the head of + # the SubQueue. It takes both locks because it mutates head.next, + # which dequeue() also mutates (same reasoning as remove()). + if obj is None: + raise ValueError("Does not support NoneType.") + old_size = -1 + node = LinkedBlockingMultiQueue.Node(obj) + self.fully_lock() + try: + self.enqueue_first(node) + self.count.inc() + if self.enabled: + old_size = self.owner.total_count.get_and_inc() + finally: + self.fully_unlock() + + if old_size == 0: + self.owner._signal_not_empty() + def remove(self, obj: T) -> bool: if obj is None: return False @@ -299,6 +327,19 @@ def put(self, key: K, item: T) -> None: """ self.get_sub_queue(key).put(item) + def put_first(self, key: K, item: T) -> None: + """ + Put one item at the head of the SubQueue specified by the key, so that + it is the next item that SubQueue hands out. Used to return an item + that was taken but must not be consumed yet. + + :param key: the identifier of a SubQueue. + :param item: Any instance. + :raises KeyError for non-existing keys. + :return: None + """ + self.get_sub_queue(key).put_first(item) + def get(self) -> T: """ Blocking get the next available item from the queue. 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 1583d7c513d..5b7975ad357 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -89,6 +89,24 @@ def dcm_element(channel): def ecm_element(channel): return ECMElement(tag=channel, payload=EmbeddedControlMessage()) + @staticmethod + def start_consumer(queue): + """Start a helper thread doing one blocking queue.get(). + + Returns the thread and the list it appends the taken element to, so a + test can assert that nothing is handed out (thread still alive, list + empty) without depending on which sub-queue the selection strategy + happens to visit first. + + The thread is a daemon: a failing assertion can leave it blocked in + get() forever, and a non-daemon thread would then hang interpreter + shutdown instead of letting the run report the failure. + """ + taken = [] + thread = threading.Thread(target=lambda: taken.append(queue.get()), daemon=True) + thread.start() + return thread, taken + def test_it_can_init(self, queue): assert queue.is_empty() assert queue.is_control_empty() @@ -366,22 +384,223 @@ def test_it_can_disable_and_enable_a_single_data_channel( assert queue.get() is blocked assert queue.is_empty() - # Regression tests below: data channels whose sub-queue is created lazily - # (on the channel's first put) AFTER disable_data has been called must - # come up disabled — a paused or backpressured worker must not be able to - # dequeue data from them, and is_data_enabled() must not flip back to - # True just because a new channel delivered its first message. + # Regression tests below: a data channel whose sub-queue is created lazily + # (on the channel's first put) while disable_data is in effect comes up + # ENABLED, because ECMs ride data channels and an ECM landing first on + # such a channel must still be delivered. DataElements are instead + # withheld on the way out of get(): the channel is closed and the element + # is pushed back to the head of its own sub-queue, so a paused or + # backpressured worker never consumes data, nothing is lost or reordered, + # and is_data_enabled() stays False for the whole disabled period. - def test_channel_registered_after_disable_comes_up_disabled( - self, queue, data_channel + @pytest.mark.timeout(2) + @pytest.mark.parametrize( + "disable_type", + [ + InternalQueue.DisableType.DISABLE_BY_PAUSE, + InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE, + ], + ) + def test_ecm_first_on_a_channel_registered_mid_disable_is_delivered( + self, queue, data_channel, disable_type + ): + # The must-fix: ECMs travel on data channels, so a reconfiguration ECM + # that is the FIRST-EVER message of a channel registered mid-pause has + # to come out; otherwise it is never acked and the coordinator's await + # expires. The timeout turns a regression into a failure, not a hang. + queue.disable_data(disable_type) + ecm = self.ecm_element(data_channel) + queue.put(ecm) + assert queue.get() is ecm + + @pytest.mark.timeout(10) + @pytest.mark.parametrize( + "first_element_kind, expected_delivered", + [ + ("data", False), + ("ecm", True), + ("dcm", True), + ], + ) + def test_first_element_kind_decides_delivery_mid_disable( + self, + queue, + data_channel, + first_element_kind, + expected_delivered, ): - # the main regression: disable first, then the channel's FIRST put + # The matrix dimension that matters is the ELEMENT kind, not just the + # channel kind: on a data channel registered mid-disable, only a + # DataElement is withheld; control-carrying elements flow. The dcm case + # is a type gate rather than a real message shape, since a DCMElement + # is never tagged with a data channel in production. + first = { + "data": self.data_element, + "ecm": self.ecm_element, + "dcm": self.dcm_element, + }[first_element_kind](data_channel) queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - queue.put(self.data_element(data_channel)) + queue.put(first) # the channel's first-ever message, mid-disable + + consumer, taken = self.start_consumer(queue) + consumer.join(1) + if expected_delivered: + assert not consumer.is_alive() + assert taken[0] is first + else: + # withheld: nothing is handed out, and the channel is closed + assert consumer.is_alive() + assert taken == [] + assert not queue._queue.is_enabled(data_channel) + assert queue.size_data() == 1 + # released only on resume + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + consumer.join(5) + assert not consumer.is_alive() + assert taken[0] is first + + @pytest.mark.timeout(10) + def test_data_first_on_a_channel_registered_mid_disable_is_withheld( + self, queue, control_channel, data_channel + ): + # A DataElement handed to get() while data is disabled must be put + # back, not consumed: the channel closes, the element stays queued and + # keeps its place, and everything queued behind it follows in FIFO + # order once data is re-enabled. + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + data_elements = [self.data_element(data_channel) for _ in range(3)] + dcm = self.dcm_element(control_channel) + queue.put(data_elements[0]) + queue.put(dcm) + in_mem_size_before = queue.in_mem_size() + + # only the control traffic is handed out; this get() is already where + # the data element is withheld, closing its channel and leaving it + # queued in place before the DCM is returned + assert queue.get() is dcm + # so a consumer coming back for more now gets nothing + consumer, taken = self.start_consumer(queue) + consumer.join(1) + assert consumer.is_alive() + assert taken == [] + assert not queue._queue.is_enabled(data_channel) assert not queue.is_data_enabled() - # the element stays queued but must not be dequeuable assert queue.size_data() == 1 - assert queue._queue.peek() is None + assert queue.in_mem_size() == in_mem_size_before + + # more data arrives on the now-closed channel and queues up behind it + queue.put(data_elements[1]) + queue.put(data_elements[2]) + assert queue.size_data() == 3 + assert consumer.is_alive() + + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert queue.is_data_enabled() + consumer.join(5) + assert not consumer.is_alive() + # FIFO is preserved across the withhold: the released element is the + # one that was put back, and the later ones follow it + results = taken + [queue.get() for _ in range(2)] + assert all(got is put for got, put in zip(results, data_elements)) + assert queue.is_empty() + assert queue.in_mem_size() == 0 + + @pytest.mark.timeout(2) + def test_an_ecm_queued_behind_withheld_data_is_delayed_until_resume( + self, queue, control_channel, data_channel + ): + # KNOWN, PRE-EXISTING LIMITATION, asserted so nobody "fixes" it + # silently: withholding a DataElement closes its channel, which also + # holds back an ECM queued behind it on that SAME channel. Per-channel + # FIFO, "no data while paused" and "deliver ECMs immediately" cannot + # all hold once data comes first, short of unbounded buffering that + # would defeat backpressure. main behaves the same way for a channel + # already disabled by disable_data. + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + data = self.data_element(data_channel) + ecm = self.ecm_element(data_channel) + dcm = self.dcm_element(control_channel) + queue.put(data) + queue.put(ecm) + queue.put(dcm) + assert queue.get() is dcm + # the ECM does not overtake the withheld data element in front of it + consumer, taken = self.start_consumer(queue) + consumer.join(1) + assert consumer.is_alive() + assert taken == [] + assert not queue._queue.is_enabled(data_channel) + assert queue.size_data() == 2 + # both are released, in order, on resume + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + consumer.join(5) + assert not consumer.is_alive() + assert taken[0] is data + assert queue.get() is ecm + + @pytest.mark.timeout(2) + def test_is_data_enabled_stays_false_while_a_disable_reason_is_active( + self, queue, data_channel, second_data_channel + ): + # main_loop's pause wait-loop spins while `not is_control_empty() or + # not is_data_enabled()`, so a channel registering mid-pause must not + # make is_data_enabled() flip back to True and let the loop exit. + queue.put(self.data_element(data_channel)) + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert not queue.is_data_enabled() + # a brand-new channel's first-ever message arrives mid-pause + queue.put(self.data_element(second_data_channel)) + assert not queue.is_data_enabled() + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) + # one reason cleared, one still active + assert not queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert not queue.is_data_enabled() + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) + assert queue.is_data_enabled() + + @pytest.mark.timeout(5) + def test_a_resume_racing_a_withhold_does_not_strand_the_channel( + self, queue, data_channel + ): + # An element taken while data is disabled, with the last disable + # reason cleared before the withhold takes effect, must still be + # handed out. Closing the channel at that point would leave it closed + # with no reason left for enable_data to clear, stranding that channel + # for good. The patched get() places the resume exactly in the window + # between the element leaving the multi-queue and the withhold + # acquiring the lock, which is too narrow to hit reliably by racing + # real threads. + data = self.data_element(data_channel) + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + queue.put(data) # registers the channel, so put() needs no lock later + + class ResumingLock: + """Resumes once, when get() takes the lock to withhold.""" + + def __init__(self, inner): + self.inner = inner + self.fired = False + + def __enter__(self): + if not self.fired: + self.fired = True + queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + return self.inner.__enter__() + + def __exit__(self, *exc_info): + return self.inner.__exit__(*exc_info) + + real_lock = queue._lock + queue._lock = ResumingLock(real_lock) + try: + assert queue.get() is data + finally: + queue._lock = real_lock + # the channel is still open afterwards + later = self.data_element(data_channel) + queue.put(later) + assert queue.get() is later + assert queue.is_empty() @pytest.mark.timeout(2) def test_enable_data_releases_a_channel_registered_mid_disable( @@ -396,22 +615,36 @@ def test_enable_data_releases_a_channel_registered_mid_disable( assert queue.get() is data assert queue.is_empty() - @pytest.mark.timeout(2) - def test_channel_registered_under_stacked_disables_stays_disabled( - self, queue, data_channel + @pytest.mark.timeout(10) + def test_channel_registered_under_stacked_disables_stays_withheld( + self, queue, control_channel, data_channel ): queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) data = self.data_element(data_channel) + dcm = self.dcm_element(control_channel) queue.put(data) + queue.put(dcm) + # this get() withholds the data element and closes its channel before + # returning the control element + assert queue.get() is dcm + # so nothing further is handed out while both reasons are active + consumer, taken = self.start_consumer(queue) + consumer.join(1) + assert consumer.is_alive() + assert queue._queue.peek() is None # releasing only one of the two reasons must not open the channel assert not queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) assert not queue.is_data_enabled() assert queue._queue.peek() is None + assert consumer.is_alive() + assert taken == [] # releasing the remaining reason makes the element dequeuable assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) assert queue.is_data_enabled() - assert queue.get() is data + consumer.join(5) + assert not consumer.is_alive() + assert taken[0] is data @pytest.mark.timeout(2) def test_control_channel_registered_mid_disable_is_never_blocked( @@ -454,15 +687,20 @@ def test_channel_registered_while_enabled_behaves_normally( assert queue.get() is data assert queue.is_empty() - @pytest.mark.timeout(10) - def test_concurrent_first_time_puts_while_toggling_disable(self, queue): - # concurrency smoke test: receiver threads register brand-new data - # channels while the DP thread toggles disable_data/enable_data; - # only the final state is asserted, deterministically. + @pytest.mark.timeout(30) + def test_concurrent_consumption_while_toggling_disable_loses_nothing(self, queue): + # Receiver threads register brand-new data channels and keep putting + # while a consumer thread drains through get() and the DP-thread side + # toggles disable_data/enable_data. Every element must come out + # exactly once: the withhold path must neither drop an element nor + # hand the same one out twice, and an element withheld just as the + # last disable reason clears must not be stranded in a closed channel. n_threads = 8 elements_per_thread = 25 + total = n_threads * elements_per_thread start_barrier = threading.Barrier(n_threads + 1) errors = [] + consumed = [] def producer(thread_index): channel = ChannelIdentity( @@ -477,49 +715,53 @@ def producer(thread_index): except Exception as exc: # pragma: no cover - failure path errors.append(exc) + def consumer(): + try: + while len(consumed) < total: + consumed.append(queue.get()) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + threads = [ threading.Thread(target=producer, args=(i,)) for i in range(n_threads) ] + consumer_thread = threading.Thread(target=consumer, daemon=True) for thread in threads: thread.start() + consumer_thread.start() start_barrier.wait() for _ in range(5): queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) for thread in threads: thread.join() - # one last full cycle after all puts settled: every channel must be - # disabled, then re-enabled with its count added back exactly once - queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - assert not queue.is_data_enabled() - assert queue._queue.peek() is None - assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + # release whatever the last toggle left withheld + queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + consumer_thread.join(20) assert not errors - total = n_threads * elements_per_thread - assert queue.is_data_enabled() - assert queue.size_data() == total - # size() is the getable total_count: a mismatch with size_data() - # means an element was double-counted or lost by a toggle race - assert queue.size() == total - drained = 0 - while queue._queue.peek() is not None: - queue.get() - drained += 1 - assert drained == total + assert not consumer_thread.is_alive() + assert len(consumed) == total + # identity, not equality: the elements are equal by value, so only + # identity can prove none was handed out twice + assert len({id(element) for element in consumed}) == total assert queue.is_empty() + assert queue.size_data() == 0 + assert queue.in_mem_size() == 0 @pytest.mark.timeout(30) def test_concurrent_first_time_puts_racing_disable_enable_toggles(self): # Receiver threads deliver first-ever messages on distinct new data - # channels while the DP-thread side toggles pause on and off. Only - # the final state is asserted (deterministic): with the queue left - # disabled, nothing is dequeuable; after the final enable_data every - # element is dequeuable exactly once, so total_count stayed exact. + # channels while the DP-thread side toggles pause on and off. Only the + # final state is asserted (deterministic): with a disable reason still + # active a consumer must not obtain anything, and after the final + # enable_data every element comes out exactly once, so the withhold + # path kept total_count and the per-channel accounting exact. threads, channels_per_thread, toggles = 4, 10, 10 for _ in range(5): queue = InternalQueue() errors = [] + consumed = [] start = threading.Barrier(threads + 1) def producer(thread_id): @@ -551,14 +793,29 @@ def producer(thread_id): assert errors == [] total = threads * channels_per_thread assert queue.size_data() == total - assert queue._queue.peek() is None assert not queue.is_data_enabled() + + # a consumer started while the queue is disabled withholds every + # data element it is offered and then blocks + consumer_thread = threading.Thread( + target=lambda: consumed.append(queue.get()), daemon=True + ) + consumer_thread.start() + consumer_thread.join(0.5) + assert consumer_thread.is_alive() + assert consumed == [] + # nothing was lost by the withholding + assert queue.size_data() == total + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - dequeued = 0 + consumer_thread.join(5) + assert not consumer_thread.is_alive() + dequeued = len(consumed) while queue._queue.peek() is not None: queue.get() dequeued += 1 assert dequeued == total + assert errors == [] # Regression tests below: the per-category query methods iterate # _queue_ids, which put() grows on a channel's first message. Iterating From 87939215dd06e6b6710791c1691a0a74b17f901c Mon Sep 17 00:00:00 2001 From: eugenegujing Date: Mon, 17 Aug 2026 22:28:42 -0700 Subject: [PATCH 6/6] fix(pyamber): restore disabled-at-registration and fix the reconfiguration await order This reverts the dequeue-side withhold (1a65167ba): per review, "no ECM while paused" is the engine's intended semantics, and the JVM DPThread already refuses all data-channel traffic while paused. A data sub-queue registered while _queue_state is non-empty starts disabled again, which closes the pause/backpressure leak; enable_data() re-enables such channels on resume, so an ECM arriving mid-pause is delayed until resume rather than dropped, and a new regression test pins that semantics. The 30s ReconfigurationIntegrationSpec timeouts were the harness, not the engine: TestUtils.shouldReconfigure awaited reconfigureWorkflow while still paused, an order production never uses because reconfigurations only take effect on resume. shouldReconfigure now dispatches the reconfiguration without awaiting, awaits the resume ack, and only then awaits the reconfiguration ack, which also removes the test's dependence on source timing. The resume ack alone establishes that resume took effect, since ResumeHandler completes it only after every worker acked; the engine emits no RUNNING state event to wait on. --- .../main/python/core/models/internal_queue.py | 49 +-- .../linked_blocking_multi_queue.py | 41 -- .../python/core/models/test_internal_queue.py | 373 ++++-------------- .../texera/amber/engine/e2e/TestUtils.scala | 22 +- 4 files changed, 94 insertions(+), 391 deletions(-) diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index 8f3121d0bff..5900db1cf10 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -71,52 +71,22 @@ def is_empty(self, key=None) -> bool: return self._queue.is_empty(key) def get(self) -> T: - """Blocking get of the next available element. - - Data channels register enabled even during a disable window, because - ECMs ride data channels and one swallowed by a channel that came up - disabled would never be acked. A DataElement arriving here during - such a window is withheld instead: its channel is closed and the - element goes back to its sub-queue's head, for enable_data() to - release. An ECM queued behind it on the same channel is therefore - delayed until resume, which is unavoidable without unbounded - buffering, and is what main does for channels disable_data() closed. - - Assumes a single consumer; only PauseManager and BackpressureHandler - toggle the disable state, and only on the input queue. - """ - while True: - item = self._queue.get() - # a control-tagged DataElement cannot occur today, but the tag is - # wire-derived: withholding one would close a control sub-queue, - # which enable_data() never reopens - if ( - not isinstance(item, DataElement) - or item.tag.is_control - or not self._queue_state - ): - return item - with self._lock: - # enable_data() may have cleared the last reason since the - # check above; closing the channel now would strand it - if not self._queue_state: - return item - # disable first: the element must never be dequeuable in between - self._queue.disable(item.tag) - self._queue.put_first(item.tag, item) + return self._queue.get() def put(self, item: T) -> None: if isinstance(item, InternalQueueElement): if item.tag not in self._queue_ids: - # both the lock and the re-check are load-bearing: - # disable_data/enable_data iterate _queue_ids live, and a - # second add_sub_queue for the same channel would replace its - # sub-queue with an empty one, dropping whatever it holds + # registration must not interleave with disable_data/enable_data with self._lock: if item.tag not in self._queue_ids: self._queue.add_sub_queue( item.tag, 1 if item.tag.is_control else 2 ) + # while data is disabled, a new data sub-queue must + # start disabled too (before its first element is + # enqueued), or it would leak data during pause/backpressure + if not item.tag.is_control and self._queue_state: + self._queue.disable(item.tag) self._queue_ids.add(item.tag) if isinstance(item, (DataElement, ECMElement, DCMElement)): self._queue.put(item.tag, item) @@ -188,11 +158,6 @@ def in_mem_size(self) -> int: ) def is_data_enabled(self) -> bool: - # channels registered mid-disable come up enabled (see get()), so - # per-channel state alone would report data as enabled during a pause - # and let main_loop's wait-loop exit and resume processing - if self._queue_state: - return False return any( self._queue.is_enabled(queue_id) for queue_id in self._data_queue_ids() ) diff --git a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py index 48036757599..fcf175d6bcd 100644 --- a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py +++ b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py @@ -106,14 +106,6 @@ def enqueue(self, node: LinkedBlockingMultiQueue.Node[T]) -> None: self.last = node self.in_mem_size.inc(node.in_mem_size) - def enqueue_first(self, node: LinkedBlockingMultiQueue.Node[T]) -> None: - # head is a sentinel, so the first real element is head.next - node.next = self.head.next - self.head.next = node - if self.last is self.head: - self.last = node - self.in_mem_size.inc(node.in_mem_size) - def dequeue(self) -> T: assert self.size() > 0 h = self.head @@ -157,26 +149,6 @@ def put(self, obj: T) -> None: if old_size == 0: self.owner._signal_not_empty() - def put_first(self, obj: T) -> None: - # Same accounting as put(), but the element goes to the head of - # the SubQueue. It takes both locks because it mutates head.next, - # which dequeue() also mutates (same reasoning as remove()). - if obj is None: - raise ValueError("Does not support NoneType.") - old_size = -1 - node = LinkedBlockingMultiQueue.Node(obj) - self.fully_lock() - try: - self.enqueue_first(node) - self.count.inc() - if self.enabled: - old_size = self.owner.total_count.get_and_inc() - finally: - self.fully_unlock() - - if old_size == 0: - self.owner._signal_not_empty() - def remove(self, obj: T) -> bool: if obj is None: return False @@ -327,19 +299,6 @@ def put(self, key: K, item: T) -> None: """ self.get_sub_queue(key).put(item) - def put_first(self, key: K, item: T) -> None: - """ - Put one item at the head of the SubQueue specified by the key, so that - it is the next item that SubQueue hands out. Used to return an item - that was taken but must not be consumed yet. - - :param key: the identifier of a SubQueue. - :param item: Any instance. - :raises KeyError for non-existing keys. - :return: None - """ - self.get_sub_queue(key).put_first(item) - def get(self) -> T: """ Blocking get the next available item from the queue. 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 5b7975ad357..6cedc31847d 100644 --- a/amber/src/test/python/core/models/test_internal_queue.py +++ b/amber/src/test/python/core/models/test_internal_queue.py @@ -89,24 +89,6 @@ def dcm_element(channel): def ecm_element(channel): return ECMElement(tag=channel, payload=EmbeddedControlMessage()) - @staticmethod - def start_consumer(queue): - """Start a helper thread doing one blocking queue.get(). - - Returns the thread and the list it appends the taken element to, so a - test can assert that nothing is handed out (thread still alive, list - empty) without depending on which sub-queue the selection strategy - happens to visit first. - - The thread is a daemon: a failing assertion can leave it blocked in - get() forever, and a non-daemon thread would then hang interpreter - shutdown instead of letting the run report the failure. - """ - taken = [] - thread = threading.Thread(target=lambda: taken.append(queue.get()), daemon=True) - thread.start() - return thread, taken - def test_it_can_init(self, queue): assert queue.is_empty() assert queue.is_control_empty() @@ -384,223 +366,22 @@ def test_it_can_disable_and_enable_a_single_data_channel( assert queue.get() is blocked assert queue.is_empty() - # Regression tests below: a data channel whose sub-queue is created lazily - # (on the channel's first put) while disable_data is in effect comes up - # ENABLED, because ECMs ride data channels and an ECM landing first on - # such a channel must still be delivered. DataElements are instead - # withheld on the way out of get(): the channel is closed and the element - # is pushed back to the head of its own sub-queue, so a paused or - # backpressured worker never consumes data, nothing is lost or reordered, - # and is_data_enabled() stays False for the whole disabled period. - - @pytest.mark.timeout(2) - @pytest.mark.parametrize( - "disable_type", - [ - InternalQueue.DisableType.DISABLE_BY_PAUSE, - InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE, - ], - ) - def test_ecm_first_on_a_channel_registered_mid_disable_is_delivered( - self, queue, data_channel, disable_type - ): - # The must-fix: ECMs travel on data channels, so a reconfiguration ECM - # that is the FIRST-EVER message of a channel registered mid-pause has - # to come out; otherwise it is never acked and the coordinator's await - # expires. The timeout turns a regression into a failure, not a hang. - queue.disable_data(disable_type) - ecm = self.ecm_element(data_channel) - queue.put(ecm) - assert queue.get() is ecm - - @pytest.mark.timeout(10) - @pytest.mark.parametrize( - "first_element_kind, expected_delivered", - [ - ("data", False), - ("ecm", True), - ("dcm", True), - ], - ) - def test_first_element_kind_decides_delivery_mid_disable( - self, - queue, - data_channel, - first_element_kind, - expected_delivered, - ): - # The matrix dimension that matters is the ELEMENT kind, not just the - # channel kind: on a data channel registered mid-disable, only a - # DataElement is withheld; control-carrying elements flow. The dcm case - # is a type gate rather than a real message shape, since a DCMElement - # is never tagged with a data channel in production. - first = { - "data": self.data_element, - "ecm": self.ecm_element, - "dcm": self.dcm_element, - }[first_element_kind](data_channel) - queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - queue.put(first) # the channel's first-ever message, mid-disable - - consumer, taken = self.start_consumer(queue) - consumer.join(1) - if expected_delivered: - assert not consumer.is_alive() - assert taken[0] is first - else: - # withheld: nothing is handed out, and the channel is closed - assert consumer.is_alive() - assert taken == [] - assert not queue._queue.is_enabled(data_channel) - assert queue.size_data() == 1 - # released only on resume - assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - consumer.join(5) - assert not consumer.is_alive() - assert taken[0] is first + # Regression tests below: data channels whose sub-queue is created lazily + # (on the channel's first put) AFTER disable_data has been called must + # come up disabled — a paused or backpressured worker must not be able to + # dequeue data from them, and is_data_enabled() must not flip back to + # True just because a new channel delivered its first message. - @pytest.mark.timeout(10) - def test_data_first_on_a_channel_registered_mid_disable_is_withheld( - self, queue, control_channel, data_channel - ): - # A DataElement handed to get() while data is disabled must be put - # back, not consumed: the channel closes, the element stays queued and - # keeps its place, and everything queued behind it follows in FIFO - # order once data is re-enabled. - queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - data_elements = [self.data_element(data_channel) for _ in range(3)] - dcm = self.dcm_element(control_channel) - queue.put(data_elements[0]) - queue.put(dcm) - in_mem_size_before = queue.in_mem_size() - - # only the control traffic is handed out; this get() is already where - # the data element is withheld, closing its channel and leaving it - # queued in place before the DCM is returned - assert queue.get() is dcm - # so a consumer coming back for more now gets nothing - consumer, taken = self.start_consumer(queue) - consumer.join(1) - assert consumer.is_alive() - assert taken == [] - assert not queue._queue.is_enabled(data_channel) - assert not queue.is_data_enabled() - assert queue.size_data() == 1 - assert queue.in_mem_size() == in_mem_size_before - - # more data arrives on the now-closed channel and queues up behind it - queue.put(data_elements[1]) - queue.put(data_elements[2]) - assert queue.size_data() == 3 - assert consumer.is_alive() - - assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - assert queue.is_data_enabled() - consumer.join(5) - assert not consumer.is_alive() - # FIFO is preserved across the withhold: the released element is the - # one that was put back, and the later ones follow it - results = taken + [queue.get() for _ in range(2)] - assert all(got is put for got, put in zip(results, data_elements)) - assert queue.is_empty() - assert queue.in_mem_size() == 0 - - @pytest.mark.timeout(2) - def test_an_ecm_queued_behind_withheld_data_is_delayed_until_resume( - self, queue, control_channel, data_channel + def test_channel_registered_after_disable_comes_up_disabled( + self, queue, data_channel ): - # KNOWN, PRE-EXISTING LIMITATION, asserted so nobody "fixes" it - # silently: withholding a DataElement closes its channel, which also - # holds back an ECM queued behind it on that SAME channel. Per-channel - # FIFO, "no data while paused" and "deliver ECMs immediately" cannot - # all hold once data comes first, short of unbounded buffering that - # would defeat backpressure. main behaves the same way for a channel - # already disabled by disable_data. + # the main regression: disable first, then the channel's FIRST put queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - data = self.data_element(data_channel) - ecm = self.ecm_element(data_channel) - dcm = self.dcm_element(control_channel) - queue.put(data) - queue.put(ecm) - queue.put(dcm) - assert queue.get() is dcm - # the ECM does not overtake the withheld data element in front of it - consumer, taken = self.start_consumer(queue) - consumer.join(1) - assert consumer.is_alive() - assert taken == [] - assert not queue._queue.is_enabled(data_channel) - assert queue.size_data() == 2 - # both are released, in order, on resume - assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - consumer.join(5) - assert not consumer.is_alive() - assert taken[0] is data - assert queue.get() is ecm - - @pytest.mark.timeout(2) - def test_is_data_enabled_stays_false_while_a_disable_reason_is_active( - self, queue, data_channel, second_data_channel - ): - # main_loop's pause wait-loop spins while `not is_control_empty() or - # not is_data_enabled()`, so a channel registering mid-pause must not - # make is_data_enabled() flip back to True and let the loop exit. queue.put(self.data_element(data_channel)) - queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - assert not queue.is_data_enabled() - # a brand-new channel's first-ever message arrives mid-pause - queue.put(self.data_element(second_data_channel)) assert not queue.is_data_enabled() - queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) - # one reason cleared, one still active - assert not queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - assert not queue.is_data_enabled() - assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) - assert queue.is_data_enabled() - - @pytest.mark.timeout(5) - def test_a_resume_racing_a_withhold_does_not_strand_the_channel( - self, queue, data_channel - ): - # An element taken while data is disabled, with the last disable - # reason cleared before the withhold takes effect, must still be - # handed out. Closing the channel at that point would leave it closed - # with no reason left for enable_data to clear, stranding that channel - # for good. The patched get() places the resume exactly in the window - # between the element leaving the multi-queue and the withhold - # acquiring the lock, which is too narrow to hit reliably by racing - # real threads. - data = self.data_element(data_channel) - queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - queue.put(data) # registers the channel, so put() needs no lock later - - class ResumingLock: - """Resumes once, when get() takes the lock to withhold.""" - - def __init__(self, inner): - self.inner = inner - self.fired = False - - def __enter__(self): - if not self.fired: - self.fired = True - queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - return self.inner.__enter__() - - def __exit__(self, *exc_info): - return self.inner.__exit__(*exc_info) - - real_lock = queue._lock - queue._lock = ResumingLock(real_lock) - try: - assert queue.get() is data - finally: - queue._lock = real_lock - # the channel is still open afterwards - later = self.data_element(data_channel) - queue.put(later) - assert queue.get() is later - assert queue.is_empty() + # the element stays queued but must not be dequeuable + assert queue.size_data() == 1 + assert queue._queue.peek() is None @pytest.mark.timeout(2) def test_enable_data_releases_a_channel_registered_mid_disable( @@ -615,36 +396,22 @@ def test_enable_data_releases_a_channel_registered_mid_disable( assert queue.get() is data assert queue.is_empty() - @pytest.mark.timeout(10) - def test_channel_registered_under_stacked_disables_stays_withheld( - self, queue, control_channel, data_channel + @pytest.mark.timeout(2) + def test_channel_registered_under_stacked_disables_stays_disabled( + self, queue, data_channel ): queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) data = self.data_element(data_channel) - dcm = self.dcm_element(control_channel) queue.put(data) - queue.put(dcm) - # this get() withholds the data element and closes its channel before - # returning the control element - assert queue.get() is dcm - # so nothing further is handed out while both reasons are active - consumer, taken = self.start_consumer(queue) - consumer.join(1) - assert consumer.is_alive() - assert queue._queue.peek() is None # releasing only one of the two reasons must not open the channel assert not queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) assert not queue.is_data_enabled() assert queue._queue.peek() is None - assert consumer.is_alive() - assert taken == [] # releasing the remaining reason makes the element dequeuable assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE) assert queue.is_data_enabled() - consumer.join(5) - assert not consumer.is_alive() - assert taken[0] is data + assert queue.get() is data @pytest.mark.timeout(2) def test_control_channel_registered_mid_disable_is_never_blocked( @@ -687,20 +454,15 @@ def test_channel_registered_while_enabled_behaves_normally( assert queue.get() is data assert queue.is_empty() - @pytest.mark.timeout(30) - def test_concurrent_consumption_while_toggling_disable_loses_nothing(self, queue): - # Receiver threads register brand-new data channels and keep putting - # while a consumer thread drains through get() and the DP-thread side - # toggles disable_data/enable_data. Every element must come out - # exactly once: the withhold path must neither drop an element nor - # hand the same one out twice, and an element withheld just as the - # last disable reason clears must not be stranded in a closed channel. + @pytest.mark.timeout(10) + def test_concurrent_first_time_puts_while_toggling_disable(self, queue): + # concurrency smoke test: receiver threads register brand-new data + # channels while the DP thread toggles disable_data/enable_data; + # only the final state is asserted, deterministically. n_threads = 8 elements_per_thread = 25 - total = n_threads * elements_per_thread start_barrier = threading.Barrier(n_threads + 1) errors = [] - consumed = [] def producer(thread_index): channel = ChannelIdentity( @@ -715,53 +477,49 @@ def producer(thread_index): except Exception as exc: # pragma: no cover - failure path errors.append(exc) - def consumer(): - try: - while len(consumed) < total: - consumed.append(queue.get()) - except Exception as exc: # pragma: no cover - failure path - errors.append(exc) - threads = [ threading.Thread(target=producer, args=(i,)) for i in range(n_threads) ] - consumer_thread = threading.Thread(target=consumer, daemon=True) for thread in threads: thread.start() - consumer_thread.start() start_barrier.wait() for _ in range(5): queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) for thread in threads: thread.join() - # release whatever the last toggle left withheld - queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - consumer_thread.join(20) + # one last full cycle after all puts settled: every channel must be + # disabled, then re-enabled with its count added back exactly once + queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) + assert not queue.is_data_enabled() + assert queue._queue.peek() is None + assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) assert not errors - assert not consumer_thread.is_alive() - assert len(consumed) == total - # identity, not equality: the elements are equal by value, so only - # identity can prove none was handed out twice - assert len({id(element) for element in consumed}) == total + total = n_threads * elements_per_thread + assert queue.is_data_enabled() + assert queue.size_data() == total + # size() is the getable total_count: a mismatch with size_data() + # means an element was double-counted or lost by a toggle race + assert queue.size() == total + drained = 0 + while queue._queue.peek() is not None: + queue.get() + drained += 1 + assert drained == total assert queue.is_empty() - assert queue.size_data() == 0 - assert queue.in_mem_size() == 0 @pytest.mark.timeout(30) def test_concurrent_first_time_puts_racing_disable_enable_toggles(self): # Receiver threads deliver first-ever messages on distinct new data - # channels while the DP-thread side toggles pause on and off. Only the - # final state is asserted (deterministic): with a disable reason still - # active a consumer must not obtain anything, and after the final - # enable_data every element comes out exactly once, so the withhold - # path kept total_count and the per-channel accounting exact. + # channels while the DP-thread side toggles pause on and off. Only + # the final state is asserted (deterministic): with the queue left + # disabled, nothing is dequeuable; after the final enable_data every + # element is dequeuable exactly once, so total_count stayed exact. threads, channels_per_thread, toggles = 4, 10, 10 for _ in range(5): queue = InternalQueue() errors = [] - consumed = [] start = threading.Barrier(threads + 1) def producer(thread_id): @@ -793,29 +551,44 @@ def producer(thread_id): assert errors == [] total = threads * channels_per_thread assert queue.size_data() == total + assert queue._queue.peek() is None assert not queue.is_data_enabled() - - # a consumer started while the queue is disabled withholds every - # data element it is offered and then blocks - consumer_thread = threading.Thread( - target=lambda: consumed.append(queue.get()), daemon=True - ) - consumer_thread.start() - consumer_thread.join(0.5) - assert consumer_thread.is_alive() - assert consumed == [] - # nothing was lost by the withholding - assert queue.size_data() == total - assert queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE) - consumer_thread.join(5) - assert not consumer_thread.is_alive() - dequeued = len(consumed) + dequeued = 0 while queue._queue.peek() is not None: queue.get() dequeued += 1 assert dequeued == total - assert errors == [] + + @pytest.mark.timeout(2) + @pytest.mark.parametrize( + "disable_type", + [ + InternalQueue.DisableType.DISABLE_BY_PAUSE, + InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE, + ], + ) + def test_ecm_first_on_a_channel_registered_mid_disable_is_delayed_until_resume( + self, queue, data_channel, disable_type + ): + # ECMs ride data channels, so an ECM arriving as the first-ever + # message of a channel registered mid-disable is held back with the + # channel. This is the engine's intended semantics: reconfigurations + # submitted while paused take effect on resume + # (ExecutionReconfigurationService), matching the JVM DPThread, which + # refuses ALL data-channel traffic — ECMs included — while paused. + # The ECM is delayed, not dropped. Misreading exactly this behavior + # as an engine deadlock once cost a full redesign of this queue, + # hence this pin. + queue.disable_data(disable_type) + ecm = self.ecm_element(data_channel) + queue.put(ecm) # the channel's first-ever message + assert not queue.is_data_enabled() + assert queue._queue.peek() is None + # the ECM sits in a data sub-queue, so it counts towards size_data + assert queue.size_data() == 1 + assert queue.enable_data(disable_type) + assert queue.get() is ecm # Regression tests below: the per-category query methods iterate # _queue_ids, which put() grows on a channel's first message. Iterating diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala index f5c4657b8e6..737c1fcfb9e 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala @@ -371,20 +371,26 @@ object TestUtils { val physicalOps = targetOps.flatMap(op => workflow.physicalPlan.getPhysicalOpsOfLogicalOp(op.operatorIdentifier) ) - Await.result( - client.coordinatorInterface.reconfigureWorkflow( - WorkflowReconfigureRequest( - reconfiguration = physicalOps.map(op => UpdateExecutorRequest(op.id, newOpExecInitInfo)), - reconfigurationId = "test-reconfigure-1" - ), - () + // Production dispatches the reconfiguration without awaiting its ack and it + // only takes effect on resume (see ExecutionReconfigurationService), so the + // harness must not await the ack while still paused — that await is what + // used to deadlock for the full 30s command timeout. The reconfigure ack is + // awaited only after the resume ack, which ResumeHandler completes once + // every worker has acknowledged the resume. (There is no RUNNING event to + // wait for: the engine only pushes ExecutionStateUpdate to the client for + // PAUSED and terminal states.) + val reconfigured = client.coordinatorInterface.reconfigureWorkflow( + WorkflowReconfigureRequest( + reconfiguration = physicalOps.map(op => UpdateExecutorRequest(op.id, newOpExecInitInfo)), + reconfigurationId = "test-reconfigure-1" ), - commandTimeout + () ) Await.result( client.coordinatorInterface.resumeWorkflow(EmptyRequest(), ()), commandTimeout ) + Await.result(reconfigured, commandTimeout) Await.result(completion, Duration.fromMinutes(1)) result }