Lee-W commented on code in PR #66584: URL: https://github.com/apache/airflow/pull/66584#discussion_r3266580601
########## airflow-core/src/airflow/triggers/shared_stream.py: ########## @@ -0,0 +1,276 @@ +# 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. +""" +Shared underlying I/O between :class:`BaseEventTrigger` instances in the triggerer. + +When multiple triggers declare the same non-``None`` +:meth:`~airflow.triggers.base.BaseEventTrigger.shared_stream_key`, the +triggerer routes them through :class:`SharedStreamManager` so that one +underlying poll loop produces raw events that are broadcast to every +participating trigger. Each trigger then runs +:meth:`~airflow.triggers.base.BaseEventTrigger.filter_shared_stream` to +convert the broadcast into its own :class:`~airflow.triggers.base.TriggerEvent` +instances. Triggers that opt out (the default) keep their independent +``run()``-based poll loops untouched. + +Lifecycle invariants +-------------------- + +The manager and groups cooperate to keep a single invariant true at every +``await``-point: + + A key is present in :attr:`SharedStreamManager._groups` only while its + group's poll task is alive and accepting new subscribers. + +This rules out the late-subscriber races that the naive design admits — a +new subscriber for a key whose poll has died or is in the middle of being +torn down always falls through to "create a fresh group" rather than +attaching to a dead one and hanging on an empty queue. The invariant is +maintained synchronously: + +* When ``_poll`` ends for any reason other than cancellation (the upstream + iterator raised, or returned), the group's ``finally`` block evicts the + key from ``_groups`` and broadcasts a terminal sentinel to current + subscribers — all without yielding, so no other coroutine can interleave. +* When the last subscriber leaves, :meth:`SharedStreamManager.unsubscribe` + evicts the key from ``_groups`` *before* awaiting ``group.stop()``, so a + new subscriber arriving while we wait for cancellation creates a fresh + group. +* :meth:`SharedStreamManager.stop_all` clears ``_groups`` in one synchronous + step before awaiting any stop, applying the same rule to shutdown. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Callable, Hashable +from contextlib import suppress +from typing import TYPE_CHECKING, Any + +import structlog + +if TYPE_CHECKING: + from structlog.stdlib import BoundLogger + + from airflow.triggers.base import BaseEventTrigger + +log = structlog.get_logger(__name__) + + +class _PollTerminated(Exception): + """ + Raised inside subscribers when ``open_shared_stream`` returns without yielding more events. + + Implementations are expected to run for the lifetime of the group; an + early return would otherwise leave subscribers waiting forever on an + empty queue. + """ + + +class _PollFailure: + """Sentinel propagated through subscriber queues when the shared poll ends.""" + + __slots__ = ("exc",) + + def __init__(self, exc: BaseException) -> None: + self.exc = exc + + +async def _drain(queue: asyncio.Queue) -> AsyncIterator[Any]: + """ + Yield items from ``queue`` until a poll termination sentinel arrives. + + Subscribers exit either by their consuming task being cancelled + (Airflow's standard idiom — :class:`CancelledError` propagates through + ``queue.get()``) or by the shared poll ending, in which case the + :class:`_PollFailure` sentinel re-raises here. + """ + while True: + item = await queue.get() + if isinstance(item, _PollFailure): + raise item.exc + yield item + + +class _SharedStreamGroup: + """One shared poll loop broadcasting raw events to N subscriber queues.""" + + def __init__( + self, + *, + key: Hashable, + trigger_class: type[BaseEventTrigger], + kwargs: dict[str, Any], + on_poll_terminate: Callable[[_SharedStreamGroup], None], + log: BoundLogger, + ) -> None: + self.key = key + self.trigger_class = trigger_class + self.kwargs = kwargs + self.log = log + self._on_poll_terminate = on_poll_terminate + self._subscribers: dict[int, asyncio.Queue] = {} + self._poll_task: asyncio.Task | None = None + + def start(self) -> None: + """Start the underlying poll loop. Call exactly once per group.""" + if self._poll_task is not None: + raise RuntimeError(f"Shared stream group {self.key!r} already started") + self._poll_task = asyncio.create_task( + self._poll(), + name=f"shared-stream-poll[{self.key!r}]", + ) + + async def _poll(self) -> None: + terminal_exc: BaseException | None = None + try: + async for raw_event in self.trigger_class.open_shared_stream(self.kwargs): + for queue in self._subscribers.values(): + queue.put_nowait(raw_event) + terminal_exc = _PollTerminated( + f"open_shared_stream for {self.key!r} returned without raising; " + "shared streams are expected to run for the lifetime of the group" + ) + except asyncio.CancelledError: + # ``stop()`` initiated this; the manager has already evicted the + # group and is awaiting our exit. Do not run the terminate path. + raise + except Exception as exc: + terminal_exc = exc + self.log.exception("Shared stream poll failed; propagating to subscribers", key=self.key) + finally: + if terminal_exc is not None: + # Synchronous: evict from the manager and broadcast the + # sentinel before returning to the loop, so no coroutine can + # observe ``_groups[key]`` pointing at a dead poll. + self._on_poll_terminate(self) + failure = _PollFailure(terminal_exc) + for queue in self._subscribers.values(): + queue.put_nowait(failure) + + def subscribe(self, trigger_id: int) -> AsyncIterator[Any]: + """Register ``trigger_id`` as a subscriber and return its raw event stream.""" + if trigger_id in self._subscribers: + raise RuntimeError(f"Trigger {trigger_id} already subscribed to shared stream {self.key!r}") + queue: asyncio.Queue = asyncio.Queue() + self._subscribers[trigger_id] = queue + return _drain(queue) + + def unsubscribe(self, trigger_id: int) -> None: + # Active subscribers exit through their consuming task being cancelled + # (Airflow's standard idiom); dropping the queue is enough here. + self._subscribers.pop(trigger_id, None) + + def is_empty(self) -> bool: + return not self._subscribers + + async def stop(self) -> None: + """Cancel the poll task if it is still running and wait for it to exit.""" + if self._poll_task is None or self._poll_task.done(): + return + self._poll_task.cancel() + with suppress(asyncio.CancelledError): + await self._poll_task + + +class SharedStreamManager: Review Comment: yep, kept scoped out of this PR. Manual-commit Kafka, SQS delete-on-process, Pub/Sub ack_id, and Service Bus peek-lock all need a producer-side ack channel that `filter_shared_stream` doesn't provide today, and the right shape for that channel really needs to be designed against a concrete consumer rather than in the abstract. The module docstring, BaseEventTrigger class docstring, and event-scheduling.rst "Suitable upstreams" section all document the scope. I'll open a follow-up issue and link it back here -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
