diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py index abc1793ff6c..5900db1cf10 100644 --- a/amber/src/main/python/core/models/internal_queue.py +++ b/amber/src/main/python/core/models/internal_queue.py @@ -20,9 +20,8 @@ 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.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 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) else: raise ValueError(f"item {item} is not recognized by internal queue") @@ -94,19 +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: - return all( - self.is_empty(queue_id) - for queue_id in 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 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() @@ -115,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 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 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: @@ -148,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 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 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 663f95d89a8..6cedc31847d 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,327 @@ 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 + + @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 + # 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 + ): + # 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. + 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() 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 }