diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index e9fcbded..1df27f40 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -195,6 +195,64 @@ jobs: name: python-${{ matrix.python-version }} fail_ci_if_error: false + # The lock file pins one FastStream release, so the `test` job never sees the + # bottom of the range pyproject.toml allows. This job runs the suite against + # the lowest supported FastStream, on the lowest supported Python. Keep the + # version here in step with the lower bound in sdk/pyproject.toml. + test-min-faststream: + runs-on: ubuntu-latest + name: Test SDK (lowest FastStream, Python 3.10) + env: + MIN_FASTSTREAM: '0.6.4' + + services: + redis: + image: redis:7-alpine + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 10s + --health-timeout 5s + --health-retries 5 + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Set up Python 3.10 + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: '3.10' + + - name: Cache Poetry dependencies + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: | + ~/.cache/pypoetry + ~/.cache/pip + # Own prefix: this job's venv has FastStream 0.6.4 installed over the + # lock, which must not be restored into the regular test jobs. + key: poetry-min-faststream-${{ runner.os }}-py3.10-${{ hashFiles('sdk/poetry.lock') }} + restore-keys: | + poetry-min-faststream-${{ runner.os }}-py3.10- + + - name: Install Poetry + run: pipx install poetry==2.4.3 + + - name: Copy README to SDK + run: cp README.md sdk/ + + - name: Install dependencies + working-directory: sdk + run: | + poetry install --all-extras + poetry run pip install "faststream[kafka,redis]==${MIN_FASTSTREAM}" + poetry run python -c "from importlib.metadata import version; v = version('faststream'); assert v == '${MIN_FASTSTREAM}', v" + + - name: Run tests + working-directory: sdk + run: poetry run pytest -v + build: runs-on: ubuntu-latest name: Build and verify SDK package @@ -263,7 +321,7 @@ jobs: # Summary job to check if all required jobs passed all-checks-passed: runs-on: ubuntu-latest - needs: [lint, changelog, test, build] + needs: [lint, changelog, test, test-min-faststream, build] if: always() steps: @@ -272,6 +330,7 @@ jobs: CHANGELOG_RESULT: ${{ needs.changelog.result }} LINT_RESULT: ${{ needs.lint.result }} TEST_RESULT: ${{ needs.test.result }} + TEST_MIN_FASTSTREAM_RESULT: ${{ needs.test-min-faststream.result }} BUILD_RESULT: ${{ needs.build.result }} run: | # Changelog is optional (skipped on push to main or when no code changes) @@ -282,6 +341,7 @@ jobs: if [[ "$LINT_RESULT" != "success" || \ "$TEST_RESULT" != "success" || \ + "$TEST_MIN_FASTSTREAM_RESULT" != "success" || \ "$BUILD_RESULT" != "success" ]]; then echo "One or more required checks failed" exit 1 diff --git a/docs/docs/sdk/redis-transport.md b/docs/docs/sdk/redis-transport.md index c6e385e4..e40d40a0 100644 --- a/docs/docs/sdk/redis-transport.md +++ b/docs/docs/sdk/redis-transport.md @@ -323,10 +323,49 @@ async def handle_order(message): Guidelines: -- `retry_on_idle_ms` should be comfortably longer than your handler's expected worst-case execution time to avoid false positives. +- `retry_on_idle_ms` should be comfortably longer than your handler's expected worst-case execution time to avoid false positives — or enable [`renew_lease`](#long-running-handlers-lease-renewal), which decouples the two. - `retry_reclaim_interval_s` controls how often the background reclaimer wakes up. Lower values increase Redis load; 15 s is a sensible default for most workloads. - `max_retries` prevents poison messages from looping forever. Set to `None` for unlimited retries (no DLQ). +### Long-Running Handlers: Lease Renewal + +Redis resets a PEL entry's idle time only on (re)delivery, **not while a handler works on it**. A handler that runs longer than `retry_on_idle_ms` (a multi-minute LLM call, say) is therefore reclaimed and redelivered *while it is still running*: processed twice, in parallel. A larger `retry_on_idle_ms` avoids that, but delays recovery from a crashed worker as much. + +`renew_lease=True` removes the trade-off: while a handler runs, its entry is renewed so the reclaimer skips it. A crashed worker stops renewing, and its entries are reclaimed `retry_on_idle_ms` after the last renewal. + +```python +@agent.subscribe( + channel=requests, + retry_on_idle_ms=60_000, # crash recovery after ~1 min of silence + renew_lease=True, # ...however long a healthy handler runs + max_processing_ms=900_000, # optional: give up (and retry) after 15 min +) +async def handle_request(message): + await call_llm(message) # may take many minutes +``` + +| Option | Default | Meaning | +|---|---|---| +| `renew_lease` | `False` | Renew in-flight entries. Requires `retry_on_idle_ms`. | +| `renew_lease_interval_ms` | `retry_on_idle_ms // 3` | How often to renew; must be `< retry_on_idle_ms`. The default tolerates two failed renewals in a row. Requires `renew_lease=True`. | +| `cancel_on_lease_lost` | `True` | Cancel a handler whose entry was reclaimed anyway (`LeaseLostError`); `False` logs and lets it finish. Requires `renew_lease=True`. | +| `max_processing_ms` | `None` | Handler deadline: cancel, raise `ProcessingTimeoutError`, NACK → normal retry. Requires `retry_on_idle_ms`. | + +**How it works.** Every interval, one Lua script checks per entry that it is still pending *for this consumer* and still in the stream, then runs `XCLAIM 0 JUSTID`: idle time reset, owner and delivery count unchanged. The ownership check means a renewal never takes an entry back from the reclaimer, whose own `XCLAIM` (`min-idle-time = retry_on_idle_ms`) skips renewed entries. The main and the `.retry` stream are renewed under their own group and consumer. Each call is bounded by half the interval; failures are logged (stream, group, consumer, ids) and retried, never stopping the consumer. The server must allow Lua scripting (`SCRIPT LOAD` / `EVALSHA`); if it refuses, an error is logged at `connect()` and handlers run unleased. + +**Lost leases.** An in-flight entry no longer pending for this consumer was normally reclaimed after renewals failed for longer than `retry_on_idle_ms` (a Redis outage, a blocked event loop), and is redelivered. By default the handler is cancelled and `eggai.transport.LeaseLostError` is raised instead of its result, so it can't overlap the redelivery; the entry is not acked. A `batch=True` batch is one unit: if any entry is lost, none is acked (the ones still owned are retried); `lost_ids` names the lost entries, `message_ids` the batch. An entry still pending for this consumer but deleted from the stream (`MAXLEN`, `XTRIM`, `XDEL`) is not lost: nothing can redeliver it, so the handler finishes (with a warning). An entry that left the PEL is lost even if also deleted, since with `delete_on_ack` the reclaimer deletes the original when moving it to `.retry`. + +**Prefetch and `max_records`.** FastStream reads up to `max_records` entries per `XREADGROUP` (all available ones when `None`), so entries queued behind a slow handler sit idle in this consumer's PEL; they are renewed too, also with `max_workers > 1`. Entries older than anything this process read (left by an earlier process with the same consumer name, e.g. a restarted container) are not renewed, so the reclaimer recovers them. As a leased entry stays with the worker that read it, `renew_lease=True` defaults `max_records` to **1** unless you set a number or use `batch=True` (an explicit `None` also becomes 1). + +**Handler deadline.** A hung handler would keep its lease forever. `max_processing_ms` cancels a handler still running at the deadline with `eggai.transport.ProcessingTimeoutError` (a `TimeoutError`); the entry is NACKed and retried after `retry_on_idle_ms`, counting towards `max_retries`. It also works without `renew_lease`. + +**Constraints.** + +- Handlers should be `async`. A sync handler's worker thread can't be interrupted, so it is never cancelled: a lost lease or missed deadline takes effect when the thread returns (result discarded, error raised). Past the deadline the entry stays leased until then, so its retry can't overlap the thread, but a sync handler that never returns keeps its lease; after a lost lease the thread may overlap the redelivery. +- A handler that blocks the event loop blocks renewal too, including a sync handler with `data_type` / `filter_by_message` (the filter wrapper calls it on the event loop). +- Requires `retry_on_idle_ms` and a consumer group; rejected with `no_ack=True` and `ack_policy=AckPolicy.MANUAL` (an entry acked mid-run would look lost). +- Delivery stays at-least-once: the lease removes the *parallel* duplicate, not the redelivery after a crash. Handlers must still be idempotent. + ### Automatic Recovery from Redis Stream Loss (NOGROUP) If Redis loses streams mid-life — due to a restart without persistence, a failover, or memory eviction — FastStream's consume loop retries `XREADGROUP` but never re-runs `XGROUP CREATE`, resulting in an infinite `NOGROUP` error loop. @@ -372,8 +411,14 @@ async def recovery_handler(message): ## Backward Compatibility -`retry_on_idle_ms` is fully opt-in. Existing subscriptions without it behave exactly as before — no extra streams, no background tasks, no changed semantics. +`retry_on_idle_ms` is fully opt-in. Existing subscriptions without it behave exactly as before — no extra streams, no background tasks, no changed semantics. The same holds for `renew_lease` and `max_processing_ms`: without them no lease client or renewal task is created and handlers are not wrapped. ## API Reference ::: eggai.transport.RedisTransport + +### Lease exceptions + +::: eggai.transport.LeaseLostError + +::: eggai.transport.ProcessingTimeoutError diff --git a/sdk/CHANGELOG.md b/sdk/CHANGELOG.md index afe3aa72..35154087 100644 --- a/sdk/CHANGELOG.md +++ b/sdk/CHANGELOG.md @@ -7,6 +7,49 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Redis `subscribe(..., renew_lease=True)`: in-flight lease renewal. Redis + resets a PEL entry's idle time only on delivery, so a handler running longer + than `retry_on_idle_ms` was reclaimed and redelivered while still running. + - Every `renew_lease_interval_ms` (default `retry_on_idle_ms // 3`) a Lua + script renews this consumer's in-flight entries with `XCLAIM ... 0 + JUSTID` (owner and delivery count unchanged) on the main and `.retry` + stream, including prefetched ones; a crashed worker's entries are still + reclaimed. `renew_lease` defaults `max_records` to 1. + - An entry reclaimed anyway cancels its handler with + `eggai.transport.LeaseLostError` (`cancel_on_lease_lost=False` lets it + finish); sync handlers are never interrupted, the error is raised when + their thread returns. + - Requires `retry_on_idle_ms` and Lua scripting; rejected with `no_ack=True` + and `AckPolicy.MANUAL`. Opt-in, default behaviour unchanged. +- Redis `subscribe(..., max_processing_ms=...)`: handler deadline. A handler + still running at it is cancelled with `eggai.transport.ProcessingTimeoutError` + and its entry NACKed for the normal retry. Requires `retry_on_idle_ms`. + +### Changed + +- Minimum FastStream is now 0.6.4 (`faststream>=0.6.4,<0.8`, was `>=0.6`). + The SDK never worked on the older 0.6 releases: on 0.6.0–0.6.2 every Redis + stream subscription failed with a `TypeError` (`StreamSub` has no + `min_idle_time` before 0.6.3), and on 0.6.3 a `min_idle_time` subscription + never claimed pending entries. Nothing that works today stops working; the + failure moves from runtime to install time. CI now also runs the test suite + against FastStream 0.6.4. +- `anyio` (`>=4.0,<5`, the range FastStream 0.6.4 already requires) is now a + declared dependency; the SDK imports it directly. + +### Fixed + +- Redis: consuming resumes after a lost consumer group (stream deleted, + Redis flushed or failed over) with FastStream >= 0.7.6. Those versions stop + a subscriber for good on `NOGROUP` instead of retrying the read, so the + group monitor recreated the group but nothing read from it any more. The + monitor now also restarts the stream subscribers `connect()` started, once + their group exists again (a partial loss still redelivers from id `0`). No + change on FastStream <= 0.7.5, where the read loop never stopped. The lock + file moves to FastStream 0.7.7 so CI covers this path. + ## [0.6.0] - 2026-09-25 ### Added diff --git a/sdk/README.md b/sdk/README.md index e1edd70d..d91e3667 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -253,6 +253,19 @@ if one joins later, the group monitor turns deletion off for that stream and log Deleted entries can't be replayed, and a generous `max_len` is still worth keeping as a backstop for the few paths that leave acked entries behind. +**Long-running handlers with `renew_lease`:** Redis resets an entry's idle time only on +delivery, so a handler running longer than `retry_on_idle_ms` is reclaimed and redelivered +while still running. `renew_lease=True` keeps renewing in-flight entries (a crashed worker's +are still reclaimed), and `max_processing_ms` adds an optional handler deadline. See +[Lease Renewal](https://docs.egg-ai.com/sdk/redis-transport/#long-running-handlers-lease-renewal). + +```python +@agent.subscribe(channel=requests, retry_on_idle_ms=60_000, renew_lease=True, + max_processing_ms=900_000) +async def handle_request(message): + await call_llm(message) # minutes; not redelivered while it runs +``` + **Retry backoff:** By default retries fire at a constant cadence equal to `retry_on_idle_ms`. Set `retry_backoff_multiplier > 1.0` to back off exponentially: a message is treated as due for reclaim once it has been idle for diff --git a/sdk/eggai/agent.py b/sdk/eggai/agent.py index 3cc8c807..ecbc85e3 100644 --- a/sdk/eggai/agent.py +++ b/sdk/eggai/agent.py @@ -8,6 +8,7 @@ from .hooks import eggai_register_stop from .transport import get_default_transport from .transport.base import Transport +from .transport.lease import LEASE_OPTION_KEYS, resolve_lease_options HANDLERS_IDS: defaultdict[str, int] = defaultdict(int) @@ -127,6 +128,11 @@ def decorator(handler: Callable[[dict[str, Any]], "asyncio.Future"]): "idle threshold); a cap below the base would disable backoff " "entirely." ) + # renew_lease / max_processing_ms: same checks as the transport, + # but at decoration time. Validation only: kwargs is read, not + # changed, and the transport resolves the options again. + if any(k in kwargs for k in LEASE_OPTION_KEYS): + resolve_lease_options(kwargs, kwargs.get("retry_on_idle_ms")) # Plugins see the kwargs as the caller wrote them (before the DLQ # key below is resolved to its namespaced form). original_kwargs = kwargs.copy() diff --git a/sdk/eggai/transport/__init__.py b/sdk/eggai/transport/__init__.py index 11515a0e..1ce47776 100644 --- a/sdk/eggai/transport/__init__.py +++ b/sdk/eggai/transport/__init__.py @@ -7,4 +7,6 @@ ) from .inmemory import InMemoryTransport as InMemoryTransport from .kafka import KafkaTransport as KafkaTransport +from .lease import LeaseLostError as LeaseLostError +from .lease import ProcessingTimeoutError as ProcessingTimeoutError from .redis import RedisTransport as RedisTransport diff --git a/sdk/eggai/transport/lease.py b/sdk/eggai/transport/lease.py new file mode 100644 index 00000000..410a12f1 --- /dev/null +++ b/sdk/eggai/transport/lease.py @@ -0,0 +1,781 @@ +"""In-flight lease renewal for Redis Streams consumer groups. + +With SDK-managed retries (``retry_on_idle_ms``) the reclaimer treats every PEL +entry idle for longer than ``retry_on_idle_ms`` as abandoned. Redis only resets +an entry's idle time when it is (re)delivered, not while a handler works on it, +so a handler that runs longer than ``retry_on_idle_ms`` has its entry reclaimed +and redelivered while it is still running: the same message is processed twice, +in parallel. + +``renew_lease=True`` keeps the entries a subscription is working on "fresh": +every ``renew_lease_interval_ms`` it runs, per entry and atomically (one Lua +script), ``XCLAIM 0 JUSTID`` on entries this +consumer still owns. That resets the idle time without changing the owner or the +delivery count, so the reclaimer skips them. A consumer that crashed stops +renewing, and its entries are reclaimed ``retry_on_idle_ms`` after the last +renewal, as before. + +Leased entries are the ones being handled plus the ones FastStream has already +read into this consumer's PEL but not handed to the handler yet (``max_records`` +> 1, ``max_workers`` > 1): they are idle while queued in the process too. + +``max_processing_ms`` is the matching deadline: a handler still running after it +is cancelled and its entry NACKed, so the normal retry path picks it up. Without +it, a hung handler would hold its lease forever. +""" + +import asyncio +import functools +import logging +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +import redis.asyncio as aioredis +from redis.exceptions import ResponseError + +from eggai.transport.middleware_utils import call_handler, is_async_callable + +logger = logging.getLogger(__name__) + +# Subscribe options owned by this module (popped by RedisTransport.subscribe). +LEASE_OPTION_KEYS = ( + "renew_lease", + "renew_lease_interval_ms", + "cancel_on_lease_lost", + "max_processing_ms", +) + +# Ids per renewal script call: keeps each Lua run (3 cheap commands per id) +# short, so Redis is never blocked noticeably. +_CHUNK = 100 +# XPENDING page size for the prefetch scan. The scan reads only this consumer's +# PEL from the oldest run in flight onwards, so one page covers it in practice; +# larger PELs are paged. +_SCAN_PAGE = 1000 +# Shortest pause between two renewal rounds, as a share of the interval. +_MIN_PAUSE_RATIO = 0.1 +# Bound on the up-front SCRIPT LOAD in LeaseManager.start() (runs in connect()). +_SCRIPT_CHECK_TIMEOUT_S = 5.0 + +# Per id: 1 = renewed, 0 = lost (not in this consumer's PEL: reclaimed, acked by +# someone else, or claimed by another consumer), 2 = trimmed (still ours in the +# PEL, but the entry was deleted from the stream, so nothing can redeliver it). +# Checked and claimed in one script so a reclaimer can't take the entry between +# the ownership check and the XCLAIM (a plain XCLAIM with min-idle-time 0 would +# steal it back from whoever holds it). +_RENEW_SCRIPT = """ +local out = {} +for i = 3, #ARGV do + local id = ARGV[i] + local p = redis.call('XPENDING', KEYS[1], ARGV[1], id, id, 1) + if #p == 0 or p[1][2] ~= ARGV[2] then + out[#out + 1] = 0 + elseif #redis.call('XRANGE', KEYS[1], id, id) == 0 then + out[#out + 1] = 2 + else + redis.call('XCLAIM', KEYS[1], ARGV[1], ARGV[2], 0, id, 'JUSTID') + out[#out + 1] = 1 + end +end +return out +""" +_LOST, _RENEWED, _TRIMMED = 0, 1, 2 + + +class LeaseLostError(Exception): + """A handler's stream entry left its consumer's PEL while it was running. + + Raised by a ``renew_lease`` subscription (with the default + ``cancel_on_lease_lost=True``) in place of the handler's result, after the + handler was cancelled. The entry was reclaimed (usually after renewals kept + failing for longer than ``retry_on_idle_ms``, e.g. during a Redis outage) + and has been, or will be, redelivered; cancelling this run keeps it from + running in parallel with that redelivery. Also raised, without running the + handler, for an entry that was lost while it was still queued in the + process. The entry is not acked: it is not this consumer's any more. + + A batch is one unit: if any of its entries is lost, the batch handler is + cancelled and none is acked; the entries still owned stay pending and are + retried. ``message_ids`` are all the run's entries, ``lost_ids`` the ones + that left the PEL. + """ + + def __init__( + self, + stream: str, + group: str, + message_ids: Iterable[str], + lost_ids: Iterable[str] | None = None, + ): + self.stream = stream + self.group = group + self.message_ids = tuple(message_ids) + self.lost_ids = tuple(lost_ids) if lost_ids is not None else self.message_ids + super().__init__( + f"lease lost on stream {stream!r} group {group!r} for " + f"{', '.join(self.lost_ids)}: the entry left this consumer's PEL " + "(reclaimed for redelivery); the handler was cancelled" + + ( + f" (batch of {', '.join(self.message_ids)})" + if self.lost_ids != self.message_ids + else "" + ) + ) + + +class ProcessingTimeoutError(TimeoutError): + """A handler ran longer than its subscription's ``max_processing_ms``. + + The handler was cancelled and the entry is NACKed (left in the PEL), so the + reclaimer retries it after ``retry_on_idle_ms``, counting towards + ``max_retries`` like any other failure. + """ + + def __init__( + self, + stream: str, + group: str, + message_ids: Iterable[str], + max_processing_ms: int, + ): + self.stream = stream + self.group = group + self.message_ids = tuple(message_ids) + self.max_processing_ms = max_processing_ms + super().__init__( + f"handler exceeded max_processing_ms={max_processing_ms} on stream " + f"{stream!r} group {group!r} for {', '.join(self.message_ids)}; " + "cancelled, the entry is left for retry" + ) + + +@dataclass(frozen=True) +class LeaseOptions: + """Validated lease-related subscribe options.""" + + renew_lease: bool = False + interval_ms: int | None = None # resolved: retry_on_idle_ms // 3 by default + cancel_on_lease_lost: bool = True + max_processing_ms: int | None = None + + +def _is_positive_int(value: Any) -> bool: + # bool is an int subclass: True must not pass as 1 ms. + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def resolve_lease_options( + options: Mapping[str, Any], retry_on_idle_ms: int | None +) -> LeaseOptions: + """Validate the lease options in ``options`` (a subscribe kwargs mapping). + + Shared by ``Agent.subscribe`` (decoration time) and + ``RedisTransport.subscribe``, so a bad combination fails the same way on + both paths. ``retry_on_idle_ms`` is the subscription's reclaim threshold. + ``options`` is only read, never modified. + """ + lease = options.get("renew_lease", False) + interval_ms = options.get("renew_lease_interval_ms") + cancel_on_lost = options.get("cancel_on_lease_lost", True) + max_processing_ms = options.get("max_processing_ms") + if not isinstance(lease, bool): + raise ValueError("renew_lease must be True or False") + if not isinstance(cancel_on_lost, bool): + raise ValueError("cancel_on_lease_lost must be True or False") + if not lease: + if interval_ms is not None or "cancel_on_lease_lost" in options: + raise ValueError( + "renew_lease_interval_ms / cancel_on_lease_lost require " + "renew_lease=True." + ) + else: + if retry_on_idle_ms is None: + raise ValueError( + "renew_lease requires retry_on_idle_ms: the lease keeps the " + "SDK-managed reclaimer from redelivering an entry that is still " + "being processed, and without retry_on_idle_ms there is none." + ) + if interval_ms is None: + interval_ms = max(1, retry_on_idle_ms // 3) + if interval_ms >= retry_on_idle_ms: + # Only for retry_on_idle_ms < 2: no renewal interval fits. + raise ValueError( + "renew_lease requires retry_on_idle_ms >= 2: leases are " + "renewed every renew_lease_interval_ms (a positive int " + "< retry_on_idle_ms)." + ) + elif not _is_positive_int(interval_ms): + raise ValueError("renew_lease_interval_ms must be a positive int") + elif interval_ms >= retry_on_idle_ms: + raise ValueError( + "renew_lease_interval_ms must be < retry_on_idle_ms: an entry " + "renewed less often than the reclaim threshold is reclaimed " + "between two renewals (the default is retry_on_idle_ms // 3)." + ) + if max_processing_ms is not None: + if not _is_positive_int(max_processing_ms): + raise ValueError("max_processing_ms must be a positive int") + if retry_on_idle_ms is None: + raise ValueError( + "max_processing_ms requires retry_on_idle_ms: a handler past its " + "deadline is NACKed, and only the SDK-managed reclaimer retries a " + "NACKed entry." + ) + return LeaseOptions( + renew_lease=lease, + interval_ms=interval_ms if lease else None, + cancel_on_lease_lost=cancel_on_lost, + max_processing_ms=max_processing_ms, + ) + + +def _id_key(msg_id: str) -> tuple[int, int]: + ms, _, seq = msg_id.partition("-") + return int(ms), int(seq or 0) + + +def _text(value: Any) -> str: + return value.decode() if isinstance(value, bytes) else str(value) + + +@dataclass(frozen=True) +class LeaseConfig: + stream: str # full stream key (main or .retry stream) + group: str + consumer: str # the FastStream consumer name that owns the PEL entries + interval_s: float + cancel_on_lost: bool = True + # Also renew entries read into this consumer's PEL but not yet handed to the + # handler (max_records > 1 without batch, or max_workers > 1). + scan_prefetched: bool = False + + @property + def timeout_s(self) -> float: + """Bound on each renewal round trip: half the interval, so a hung call + never delays the next renewal by a whole interval.""" + return self.interval_s / 2 + + +class _Invocation: + """One handler run and the stream entries it covers (several in batch mode).""" + + __slots__ = ("ids", "task", "reason", "lost", "lost_ids", "interruptible") + + def __init__(self, ids: Iterable[str], interruptible: bool = True): + self.ids: tuple[str, ...] = tuple(ids) + self.task: asyncio.Task | None = None + # Why the SDK stopped the run: "lease_lost" or "timeout". + self.reason: str | None = None + self.lost = False + # The ids that left the PEL (a subset of ids in batch mode). + self.lost_ids: list[str] = [] + # False for a sync handler: its worker thread can't be stopped, and + # cancelling the task that awaits it would only abandon the thread + # (the entry would be released for retry while the thread still runs). + self.interruptible = interruptible + + def cancel(self, reason: str) -> None: + if self.reason is None: + self.reason = reason + if self.interruptible and self.task is not None and not self.task.done(): + self.task.cancel() + + +class LeaseKeeper: + """Renews the in-flight entries of one (stream, group, consumer).""" + + def __init__(self, config: LeaseConfig): + self.config = config + self._in_flight: dict[str, _Invocation] = {} + # Entries whose handler already finished in this process but may still be + # in the PEL (a failure, NACKed for retry). Never renewed: the reclaimer + # must see them go idle. Only needed while other runs are in flight; + # the prefetch scan drops the ones no longer pending. + self._finished: set[str] = set() + # Prefetched entries found lost before their handler started. + self._lost_queued: set[str] = set() + # Entries trimmed from the stream while in flight: no longer renewed. + self._trimmed: set[str] = set() + # Ids whose run began while a renewal round is in progress (None + # between rounds): a queued entry that started, and was acked, during + # the round trip is not a lost one. + self._begun_this_round: set[str] | None = None + + @property + def in_flight(self) -> Mapping[str, _Invocation]: + """Read-only view of the runs in flight, by stream id.""" + return MappingProxyType(self._in_flight) + + def begin(self, ids: Iterable[str], interruptible: bool = True) -> _Invocation: + inv = _Invocation(ids, interruptible) + if self._begun_this_round is not None: + self._begun_this_round.update(inv.ids) + for msg_id in inv.ids: + self._in_flight[msg_id] = inv + self._finished.discard(msg_id) + if msg_id in self._lost_queued: + self._lost_queued.discard(msg_id) + inv.lost = True + inv.lost_ids.append(msg_id) + return inv + + def end(self, inv: _Invocation) -> None: + for msg_id in inv.ids: + if self._in_flight.get(msg_id) is inv: + del self._in_flight[msg_id] + self._trimmed.discard(msg_id) + if self._in_flight: + self._finished.update(inv.ids) + else: + # Nothing older than the next run can be renewed (see renew_once), + # so the finished set is not needed any more. + self._finished.clear() + + async def renew_once(self, client: Any, script: Any) -> None: + """Renew every entry this keeper holds; react to lost ones.""" + if not self._in_flight: + return + cfg = self.config + # Everything older than the oldest run in flight is either done or + # belongs to an earlier incarnation of this consumer name (same host and + # pid after a container restart): FastStream reads with ">", so entries + # delivered to this process are newer than anything left from before. + # Not renewing those lets the reclaimer recover them. + low = min(self._in_flight, key=_id_key) + in_flight_now = { + msg_id + for msg_id, inv in self._in_flight.items() + if not inv.lost and msg_id not in self._trimmed + } + self._begun_this_round = set() + try: + queued: set[str] = set() + if cfg.scan_prefetched: + finished_before = set(self._finished) + pending: set[str] = set() + cursor = low + while True: + page = await asyncio.wait_for( + client.xpending_range( + cfg.stream, + cfg.group, + min=cursor, + max="+", + count=_SCAN_PAGE, + consumername=cfg.consumer, + ), + cfg.timeout_s, + ) + for entry in page: + msg_id = _text(entry["message_id"]) + pending.add(msg_id) + if ( + msg_id not in self._in_flight + and msg_id not in self._finished + and msg_id not in self._lost_queued + ): + queued.add(msg_id) + if len(page) < _SCAN_PAGE: + break + cursor = "(" + _text(page[-1]["message_id"]) + # A finished entry the scan didn't find was acked (or reclaimed) + # and can't become pending for this consumer again: forget it, + # or a long run would keep every id finished beside it. Ids that + # finished while the scan was running may have been missed by + # it: kept. + self._finished -= finished_before - pending + candidates = sorted(in_flight_now | queued, key=_id_key) + for start in range(0, len(candidates), _CHUNK): + chunk = candidates[start : start + _CHUNK] + statuses = await asyncio.wait_for( + script(keys=[cfg.stream], args=[cfg.group, cfg.consumer, *chunk]), + cfg.timeout_s, + ) + for msg_id, status in zip(chunk, statuses, strict=True): + self._apply(msg_id, int(status), msg_id in queued) + self._prune() + finally: + self._begun_this_round = None + + def _prune(self) -> None: + """Forget finished / lost-queued ids older than the oldest run in flight + now (the scan never reaches them). Uses the state after renew_once's + awaits: runs may have ended, all of them even, while those were pending.""" + if not self._finished and not self._lost_queued: + return + if not self._in_flight: + self._finished.clear() + self._lost_queued.clear() + return + low = _id_key(min(self._in_flight, key=_id_key)) + self._finished = {i for i in self._finished if _id_key(i) >= low} + self._lost_queued = {i for i in self._lost_queued if _id_key(i) >= low} + + def _apply(self, msg_id: str, status: int, was_queued: bool) -> None: + if status == _RENEWED: + return + cfg = self.config + inv = self._in_flight.get(msg_id) + if status == _TRIMMED: + if inv is not None: + # Nothing can redeliver an entry that is gone from the stream, + # so the run is not duplicated: let it finish (its XACK clears + # the PEL entry). Stop renewing: once the reclaimer's XCLAIM + # drops the dangling PEL entry it would look lost. + self._trimmed.add(msg_id) + logger.warning( + "Lease renewal: entry %s was deleted from stream %s " + "(group %s) while being processed (MAXLEN/XTRIM/XDEL); " + "letting the handler finish, the entry cannot be redelivered.", + msg_id, + cfg.stream, + cfg.group, + ) + return + # Lost: the entry is not in this consumer's PEL any more. + if inv is not None: + if msg_id not in inv.lost_ids: + inv.lost_ids.append(msg_id) + if inv.lost: + return + inv.lost = True + if cfg.cancel_on_lost: + logger.error( + "Lease lost for entry %s on stream %s (group %s, consumer %s): " + "it was reclaimed for redelivery; %s.", + msg_id, + cfg.stream, + cfg.group, + cfg.consumer, + "cancelling the handler" + if inv.interruptible + else "the sync handler's thread can't be interrupted and may " + "overlap the redelivery; its result will be discarded", + ) + inv.cancel("lease_lost") + else: + logger.error( + "Lease lost for entry %s on stream %s (group %s, consumer %s): " + "it was reclaimed for redelivery; the handler keeps running " + "(cancel_on_lease_lost=False) and may overlap the redelivery.", + msg_id, + cfg.stream, + cfg.group, + cfg.consumer, + ) + elif was_queued: + if self._begun_this_round and msg_id in self._begun_this_round: + # Queued at scan time, then started (and acked) during the + # round trip: its own ack is why it left the PEL. + return + # Prefetched, handler not started yet: fail it when it starts. + self._lost_queued.add(msg_id) + logger.error( + "Lease lost for queued entry %s on stream %s (group %s) before " + "its handler started.", + msg_id, + cfg.stream, + cfg.group, + ) + # Otherwise the run finished while the renewal was in flight: its ack is + # why the entry left the PEL. + + async def run(self, manager: "LeaseManager") -> None: + cfg = self.config + loop = asyncio.get_running_loop() + delay = cfg.interval_s + overrunning = False + # The flag, not just cancellation, ends the loop: on CPython < 3.12 + # asyncio.wait_for inside redis-py can swallow a CancelledError. + while manager.running: + await asyncio.sleep(delay) + started = loop.time() + client, script = manager.client, manager.script + if client is not None and script is not None: + try: + await self.renew_once(client, script) + except asyncio.CancelledError: + raise + except Exception as e: + logger.warning( + "Lease renewal failed — stream=%s group=%s consumer=%s " + "ids=%s: %s: %s (retrying in %.3fs)", + cfg.stream, + cfg.group, + cfg.consumer, + sorted(self._in_flight, key=_id_key), + type(e).__name__, + e, + cfg.interval_s, + ) + # Fixed rate: a slow round (several chunks, a timed-out call) + # shortens the next wait instead of pushing every later renewal + # back. A round slower than the interval still gets a short pause, + # so a keeper that can't keep up doesn't run back-to-back rounds. + elapsed = loop.time() - started + if elapsed >= cfg.interval_s and not overrunning: + logger.warning( + "Lease renewal round took %.3fs, longer than the %.3fs " + "interval — stream=%s group=%s consumer=%s (%d in flight); " + "entries may become reclaimable if this persists.", + elapsed, + cfg.interval_s, + cfg.stream, + cfg.group, + cfg.consumer, + len(self._in_flight), + ) + overrunning = elapsed >= cfg.interval_s + delay = max(_MIN_PAUSE_RATIO * cfg.interval_s, cfg.interval_s - elapsed) + + +LeaseKey = tuple[str, str, str] + + +class LeaseManager: + """Owns the Redis client and one renewal task per lease keeper.""" + + def __init__(self, redis_url: str, connection_kwargs: dict[str, Any] | None = None): + self._redis_url = redis_url + self._connection_kwargs: dict[str, Any] = connection_kwargs or {} + self._keepers: dict[LeaseKey, LeaseKeeper] = {} + self._tasks: dict[LeaseKey, asyncio.Task] = {} + self.client: aioredis.Redis | None = None + self.script: Any = None + self.running = False + + def add(self, config: LeaseConfig) -> tuple[LeaseKey, LeaseKeeper, bool]: + """Register (or reuse) the keeper for ``config``'s (stream, group, + consumer). Returns ``(key, keeper, created)``; only a created keeper is + the caller's to roll back with :meth:`discard`. + + Subscriptions sharing a key share one PEL, so they share one keeper. A + second registration with different settings is refused rather than + replacing a keeper that handlers already wrapped may still be using. + """ + key = (config.stream, config.group, config.consumer) + keeper = self._keepers.get(key) + if keeper is not None: + if keeper.config != config: + raise ValueError( + f"renew_lease is already configured for stream " + f"{config.stream!r} group {config.group!r} consumer " + f"{config.consumer!r} with different settings " + f"({keeper.config} vs {config}); subscriptions sharing a " + "consumer must use the same lease options." + ) + return key, keeper, False + keeper = LeaseKeeper(config) + self._keepers[key] = keeper + return key, keeper, True + + def discard(self, key: LeaseKey) -> None: + """Remove a keeper (rolling back a partial subscribe) and its task.""" + self._keepers.pop(key, None) + task = self._tasks.pop(key, None) + if task is not None: + task.cancel() + + async def start(self) -> None: + """Open the client and start missing renewal tasks. Idempotent: also + called for keepers registered after the first start.""" + if self.client is None: + # decode_responses: ids come back as str. The forwarded connection + # kwargs are resilience settings only, never decode_responses. + self.client = aioredis.from_url( + self._redis_url, + **{**self._connection_kwargs, "decode_responses": True}, + ) + self.script = self.client.register_script(_RENEW_SCRIPT) + await self._check_scripting() + self.running = True + for key, keeper in self._keepers.items(): + task = self._tasks.get(key) + if task is not None and not task.done(): + continue + self._tasks[key] = asyncio.create_task( + keeper.run(self), + name=f"lease-renewal:{keeper.config.stream}:{keeper.config.group}", + ) + + async def _check_scripting(self) -> None: + """Load the renewal script once up front, so a server that refuses Lua + scripting (an ACL without ``@scripting``, a managed-Redis policy) is + reported clearly at startup instead of only as a renewal warning every + interval. Never raises: the transport keeps running, unleased.""" + client = self.client + if client is None: + return + try: + await asyncio.wait_for( + client.script_load(_RENEW_SCRIPT), _SCRIPT_CHECK_TIMEOUT_S + ) + except asyncio.CancelledError: + raise + except ResponseError as e: + logger.error( + "Lease renewal will not work: the Redis server refused to load " + "the renewal Lua script (%s). renew_lease needs Lua scripting " + "(SCRIPT LOAD / EVALSHA; ACL category @scripting). Handlers keep " + "running, but unleased: entries of handlers running longer than " + "retry_on_idle_ms will be redelivered while still running.", + e, + ) + except Exception as e: + logger.warning( + "Lease renewal: could not check Lua scripting on the Redis " + "server (%s: %s); renewals will retry it.", + type(e).__name__, + e, + ) + + async def stop(self) -> None: + """Stop all renewal tasks and close the client.""" + self.running = False + tasks = list(self._tasks.values()) + self._tasks.clear() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + client, self.client, self.script = self.client, None, None + if client is not None: + try: + await client.aclose() + except Exception: + logger.warning( + "Lease renewal: closing the Redis client failed", exc_info=True + ) + + +async def _drain(task: asyncio.Task) -> None: + """Wait for a cancelled task to finish, ignoring its outcome.""" + try: + await task + except (asyncio.CancelledError, Exception): + pass + + +def wrap_handler_with_lease( + handler: Callable, + *, + stream: str, + group: str, + keeper: LeaseKeeper | None, + message_ids: Callable[[], list[str] | None], + max_processing_ms: int | None, +) -> Callable: + """Run ``handler`` under a lease (``keeper``) and/or a deadline. + + The handler runs in a child task so the lease keeper (lost lease) or the + deadline can cancel it without cancelling FastStream's consume loop. A + cancellation of the wrapper itself (agent stop) is forwarded to the handler. + ``message_ids`` returns the stream ids of the message being handled (from + FastStream's context); without them the handler runs unleased (logged). + + Sync handlers run in a worker thread, which cannot be interrupted, so they + are never cancelled: a lost lease or a missed deadline takes effect when the + thread returns (its result is discarded and the error raised). Past the + deadline the entry stays leased until then, so its retry cannot overlap the + thread; a sync handler that never returns therefore keeps its lease. + """ + is_async = is_async_callable(handler) + deadline_s = max_processing_ms / 1000 if max_processing_ms is not None else None + + def _error(inv: _Invocation) -> Exception: + if inv.reason == "timeout" and max_processing_ms is not None: + return ProcessingTimeoutError(stream, group, inv.ids, max_processing_ms) + return LeaseLostError(stream, group, inv.ids, inv.lost_ids or None) + + def _read_ids() -> list[str]: + try: + ids = message_ids() + except Exception: + if keeper is not None: + logger.error( + "Lease renewal: could not read the message ids on %s (group " + "%s); handling the message without a lease", + stream, + group, + exc_info=True, + ) + return [] + if not ids and keeper is not None: + logger.warning( + "Lease renewal: no stream ids for the message on %s (group %s); " + "handling it without a lease", + stream, + group, + ) + return list(ids or ()) + + @functools.wraps(handler) + async def leased_handler(*args, **kwargs): + ids = _read_ids() + lease = keeper if ids else None + inv = ( + lease.begin(ids, interruptible=is_async) + if lease is not None + else _Invocation(ids, interruptible=is_async) + ) + try: + if inv.lost and lease is not None and lease.config.cancel_on_lost: + raise LeaseLostError(stream, group, inv.ids, inv.lost_ids or None) + # Nothing awaits between begin() and here, so the keeper can't have + # stopped the run yet; honour it anyway should that ever change. + if inv.reason is not None: + raise _error(inv) + task = asyncio.create_task(call_handler(handler, is_async, *args, **kwargs)) + inv.task = task + try: + done, _ = await asyncio.wait({task}, timeout=deadline_s) + if not done: + logger.error( + "Handler on %s (group %s) exceeded max_processing_ms=%s " + "for %s; %s, the entry is left for retry.", + stream, + group, + max_processing_ms, + ", ".join(inv.ids) or "?", + "cancelling it" + if is_async + else "it runs in a worker thread that can't be " + "interrupted: waiting for it to return (the lease is " + "kept until then)", + ) + inv.cancel("timeout") + # A sync handler's thread keeps running: wait for it, still + # leased, so its retry can't overlap it. + await asyncio.wait({task}) + except asyncio.CancelledError: + # asyncio.wait never raises the child's cancellation, so this is + # the wrapper's own (agent stop): forward it to the handler. No + # Task.cancelling() needed, which Python 3.10 doesn't have. + task.cancel() + await _drain(task) + raise + # The task is done: result() doesn't suspend, so from here on no + # outer cancellation can be mistaken for the handler's outcome. + try: + result = task.result() + except asyncio.CancelledError: + if inv.reason is None: + raise # the handler cancelled itself + raise _error(inv) from None + except Exception as exc: + if inv.reason is not None: + raise _error(inv) from exc + raise + if inv.reason is not None: + # The handler swallowed the cancellation, or a sync handler's + # thread returned after the lease was lost / the deadline + # passed: the entry is not ours (lost) or overdue (timeout), so + # discard the result and don't let it be acked. + raise _error(inv) + return result + finally: + if lease is not None: + lease.end(inv) + + return leased_handler diff --git a/sdk/eggai/transport/middleware_utils.py b/sdk/eggai/transport/middleware_utils.py index 7740f1b7..0fcc7879 100644 --- a/sdk/eggai/transport/middleware_utils.py +++ b/sdk/eggai/transport/middleware_utils.py @@ -10,10 +10,13 @@ the behaviour identical across the Kafka, Redis, and in-memory transports. """ +import asyncio +import functools import inspect from collections.abc import Callable from typing import Any +import anyio from pydantic import BaseModel, ValidationError from pydantic_core import PydanticUndefined @@ -37,6 +40,41 @@ def _carry_identity(wrapper: Callable, handler: Callable) -> Callable: return wrapper +def _is_coroutine_function(obj: Any) -> bool: + if inspect.iscoroutinefunction(obj): + return True + # Marked the pre-3.12 way (e.g. AsyncMock on 3.10/3.11), which only + # asyncio.iscoroutinefunction recognised; that one is deprecated in 3.14. + marker = getattr(asyncio.coroutines, "_is_coroutine", None) + return marker is not None and getattr(obj, "_is_coroutine", None) is marker + + +def is_async_callable(handler: Callable) -> bool: + """True for a coroutine function, an object with an async ``__call__``, or a + ``functools.partial`` of either.""" + while isinstance(handler, functools.partial): + handler = handler.func + return _is_coroutine_function(handler) or _is_coroutine_function( + type(handler).__call__ + ) + + +async def call_handler(handler: Callable, is_async: bool, *args, **kwargs) -> Any: + """Run ``handler`` the way FastStream would have without an SDK wrapper. + + An async handler is awaited on the event loop; a sync one is sent to a + worker thread (FastStream only does that itself when it sees the sync + function, not an async wrapper around it). A sync callable that still + returned an awaitable has it awaited rather than dropped. + """ + if is_async: + return await handler(*args, **kwargs) + result = await anyio.to_thread.run_sync(functools.partial(handler, *args, **kwargs)) + if inspect.isawaitable(result): + result = await result + return result + + async def _invoke(handler: Callable, arg: Any) -> Any: """Call ``handler`` with ``arg``, awaiting the result only if it is awaitable. diff --git a/sdk/eggai/transport/redis.py b/sdk/eggai/transport/redis.py index 53d64216..ed4a0dd2 100644 --- a/sdk/eggai/transport/redis.py +++ b/sdk/eggai/transport/redis.py @@ -1,6 +1,5 @@ import asyncio import functools -import inspect import logging import os import socket @@ -8,7 +7,6 @@ from dataclasses import dataclass from typing import Any -import anyio import redis.asyncio as aioredis from faststream import AckPolicy from faststream.redis import RedisBroker, StreamSub @@ -16,7 +14,20 @@ from redis.exceptions import ResponseError from eggai.transport.base import Transport -from eggai.transport.middleware_utils import wrap_handler_with_filters +from eggai.transport.lease import ( + LEASE_OPTION_KEYS, + LeaseConfig, + LeaseKeeper, + LeaseKey, + LeaseManager, + resolve_lease_options, + wrap_handler_with_lease, +) +from eggai.transport.middleware_utils import ( + call_handler, + is_async_callable, + wrap_handler_with_filters, +) from eggai.transport.pending_reclaimer import ( PendingReclaimerManager, ReclaimerConfig, @@ -65,6 +76,16 @@ class _StreamGroupInfo: group_create_id: str +def _stream_group(sub: Any) -> tuple[str, str] | None: + """(stream, group) of a consumer-group stream subscriber, else None.""" + stream_sub: Any = getattr(sub, "stream_sub", None) + name = getattr(stream_sub, "name", None) + group = getattr(stream_sub, "group", None) + if not name or not group: + return None + return str(name), str(group) + + class RedisTransport(Transport): """ Redis-based transport layer adapted to use FastStream's RedisBroker for message publishing and consumption. @@ -202,6 +223,13 @@ def __init__( # Stream keys where the monitor turned delete_on_ack off because another # group appeared; shared with the reclaimer. Plain XACK from then on. self._delete_on_ack_disabled: set[str] = set() + # renew_lease: one keeper per (stream, group, consumer), renewed by the + # manager's own client. Created on the first renew_lease subscribe. + self._lease_manager: LeaseManager | None = None + # Stream subscribers connect() started, keyed by id(). FastStream + # >= 0.7.6 stops a subscriber for good when its group is gone (NOGROUP); + # the group monitor restarts these once it has recreated the group. + self._started_subscribers: dict[int, Any] = {} async def connect(self): """ @@ -241,13 +269,28 @@ async def connect(self): await self._delete_client.aclose() self._delete_client = None raise - if not getattr(self.broker, "running", False): - await self.broker.start() - else: - await self.broker.connect() - for sub in self.broker.subscribers: - if not getattr(sub, "running", False): - await sub.start() + # Lease renewal runs before any consumer starts, so the first entries + # read (and any prefetched behind them) are renewed from the start. + started_lease_manager = False + if self._lease_manager is not None: + started_lease_manager = not self._lease_manager.running + await self._lease_manager.start() + try: + if not getattr(self.broker, "running", False): + await self.broker.start() + else: + await self.broker.connect() + for sub in self.broker.subscribers: + if not getattr(sub, "running", False): + await sub.start() + except BaseException: + # Don't leave renewal tasks and their client behind a failed start. + if started_lease_manager and self._lease_manager is not None: + await self._lease_manager.stop() + raise + for sub in self.broker.subscribers: + if getattr(sub, "running", False) and _stream_group(sub): + self._started_subscribers[id(sub)] = sub if self._reclaimer_manager is not None: # start() is idempotent — it skips reclaimer tasks already running. await self._reclaimer_manager.start() @@ -281,6 +324,10 @@ async def disconnect(self): if self._reclaimer_manager is not None: await self._reclaimer_manager.stop() await self.broker.stop() + # After broker.stop(): handlers still finishing during the graceful + # shutdown keep their leases until they are done. + if self._lease_manager is not None: + await self._lease_manager.stop() # After broker.stop(): handlers still finishing may need it to delete. if self._delete_client is not None: await self._delete_client.aclose() @@ -419,6 +466,40 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: A shared DLQ is written WITHOUT ``MAXLEN`` unless ``RedisTransport(dlq_max_len=...)`` is set (``retry_max_len`` does not apply): trimming is stream-wide, so any writer's cap would delete other writers' unconsumed dead letters. By default retention of a shared DLQ is the sink's job. + renew_lease (bool, optional): Renew the lease of every entry this subscription is processing, so a + handler that runs longer than ``retry_on_idle_ms`` is not reclaimed and redelivered while it is still + running (default False). Redis resets a PEL entry's idle time only when it is (re)delivered; with + this option each in-flight entry is renewed every ``renew_lease_interval_ms`` with + ``XCLAIM 0 JUSTID`` (owner and delivery count unchanged, idle + time reset), guarded by an ownership check in the same Lua script. A consumer that crashes stops + renewing, so its entries are still reclaimed ``retry_on_idle_ms`` after the last renewal. Covers the + main stream and the SDK retry stream (each with its own group/consumer), and entries FastStream has + read into this consumer's PEL but not handed to the handler yet. Defaults ``max_records`` to 1 + (unless you set a number, or ``batch=True``) so a busy worker doesn't hold prefetched entries an + idle worker could take. Requires ``retry_on_idle_ms`` and a consumer group; incompatible with + ``no_ack=True`` and ``ack_policy=AckPolicy.MANUAL``. Needs Lua scripting (``EVALSHA``) on the + server. + renew_lease_interval_ms (int, optional): How often leases are renewed (default + ``retry_on_idle_ms // 3``, so two renewals can fail before an entry becomes reclaimable). Must be + less than ``retry_on_idle_ms``. Each renewal round trip is bounded by half the interval; a failed + renewal is logged (stream, group, ids) and retried at the next interval, it never stops the + consumer. Requires ``renew_lease=True``. + cancel_on_lease_lost (bool, optional): What to do when a renewal finds that an in-flight entry left + this consumer's PEL, i.e. it was reclaimed for redelivery (typically after renewals failed for + longer than ``retry_on_idle_ms``, e.g. a Redis outage). True (default): cancel the handler and + raise ``eggai.transport.LeaseLostError`` so it cannot run in parallel with the redelivery (the + entry is not acked; it is no longer this consumer's). False: log an error and let it finish. An + entry that is still owned but was deleted from the stream (``MAXLEN``/``XTRIM``/``XDEL``) is not a + lost lease: nothing can redeliver it, so the handler always finishes (a warning is logged). + Requires ``renew_lease=True``. + max_processing_ms (int, optional): Handler deadline. A handler still running after it is cancelled and + ``eggai.transport.ProcessingTimeoutError`` is raised, so the entry is NACKed and retried by the + reclaimer like any other failure (counts towards ``max_retries``). Default None (no deadline). + Pairs with ``renew_lease``, where a hung handler would otherwise hold its lease forever. Requires + ``retry_on_idle_ms``. Sync handlers run in a worker thread and cannot be interrupted, so they are + never cancelled: a lost lease or a missed deadline takes effect when the thread returns (result + discarded, error raised). Past the deadline the entry stays leased until then, so its retry can't + overlap the thread; a sync handler that never returns keeps its lease. retry_on_error (bool, optional): Whether to retry handler on error (default is True). # Durability parameters @@ -453,6 +534,12 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: retry_backoff_max_ms = kwargs.pop("retry_backoff_max_ms", None) retry_backoff_jitter = kwargs.pop("retry_backoff_jitter", 0.0) _internal_retry = kwargs.pop("_internal_retry", False) + # The recursive retry-stream subscribe gets the main subscription's + # already-validated lease options (it has no retry_on_idle_ms of its own). + _lease_options = kwargs.pop("_lease_options", None) + # The parent subscribe's list of lease keys to roll back if it fails + # after this (retry-stream) subscribe succeeded. + _parent_lease_keys: list | None = kwargs.pop("_lease_keys_added", None) # EggAI applies content filtering (filter_by_message) and typed-subscription # support (data_type / filter_by_data) by wrapping the handler — see @@ -623,6 +710,45 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: "retries/DLQ never trigger. Use the default NACK_ON_ERROR." ) + # Lease renewal / processing deadline (see eggai.transport.lease). + if _lease_options is None: + _lease_options = resolve_lease_options(kwargs, retry_on_idle_ms) + lease_opts = _lease_options + for option in LEASE_OPTION_KEYS: + kwargs.pop(option, None) + if lease_opts.renew_lease: + # retry_on_idle_ms (required above) already implies a group; checked + # explicitly anyway rather than trusting that further down. + if not group: + raise ValueError( + "renew_lease requires a consumer group: pass handler_id= or " + "group= (Agent.subscribe / Channel.subscribe set one " + "automatically)." + ) + if not consumer: + raise ValueError( + "renew_lease requires a consumer name: pass handler_id= or " + "consumer= (Agent.subscribe / Channel.subscribe set one " + "automatically)." + ) + if no_ack: + raise ValueError( + "renew_lease is incompatible with no_ack=True: with no ack " + "there is no PEL entry to renew." + ) + if kwargs.get("ack_policy") == AckPolicy.MANUAL: + raise ValueError( + "renew_lease is incompatible with ack_policy=AckPolicy.MANUAL: " + "an entry acked by the handler leaves the PEL mid-run and would " + "look like a lost lease." + ) + # FastStream reads max_records entries per XREADGROUP (all available + # ones when None) and hands them to the handler one by one; the + # queued ones sit in this consumer's PEL, leased to a busy worker + # instead of going to an idle one. Long handlers want one at a time. + if not batch and max_records is None: + max_records = 1 + # Default max_retries only applies when retry_on_idle_ms is set. if retry_on_idle_ms is None: max_retries = None @@ -690,6 +816,50 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: # allowing them to be redelivered for retry ack_policy = kwargs.pop("ack_policy", AckPolicy.NACK_ON_ERROR) + # Lease / deadline wrapper: around tracing( filters( handler ) ), inside + # delete_on_ack, so renewal stops before the entry is acked (and deleted). + # Applied per stream: the retry subscribe below wraps its own handler + # with the retry stream's group and consumer. + # Lease keys this call (and its retry subscribe) registered, rolled back + # if anything below fails. A keeper that already existed is left alone. + added_lease_keys: list[LeaseKey] = [] + if lease_opts.renew_lease or lease_opts.max_processing_ms is not None: + keeper: LeaseKeeper | None = None + interval_ms = lease_opts.interval_ms + if lease_opts.renew_lease and group and consumer and interval_ms: + # (group, consumer and interval were validated above.) + if self._lease_manager is None: + self._lease_manager = LeaseManager( + self._redis_url, connection_kwargs=self._connection_kwargs + ) + max_workers = kwargs.get("max_workers") or 1 + lease_key, keeper, created = self._lease_manager.add( + LeaseConfig( + stream=channel, + group=group, + consumer=consumer, + interval_s=interval_ms / 1000, + cancel_on_lost=lease_opts.cancel_on_lease_lost, + scan_prefetched=max_workers > 1 + or (not batch and max_records != 1), + ) + ) + if created: + added_lease_keys.append(lease_key) + if self._lease_manager.running: + # Registered after connect() (e.g. a second Agent on a + # shared transport): start its renewal task now rather + # than relying on the next connect(). + await self._lease_manager.start() + handler = wrap_handler_with_lease( + handler, + stream=channel, + group=str(group), + keeper=keeper, + message_ids=self._current_message_ids, + max_processing_ms=lease_opts.max_processing_ms, + ) + # Outermost wrapper, so a message a filter skips is deleted too (it is # acked either way). Deletes exactly where ack_policy acks. Applied per stream: the recursive retry subscribe # below wraps its own handler with the retry group. @@ -698,9 +868,13 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: self._delete_on_ack_groups.add((channel, group)) # stream must be passed as keyword-only argument - registered_handler = self.broker.subscriber( - stream=stream_sub, ack_policy=ack_policy, **kwargs - )(handler) + try: + registered_handler = self.broker.subscriber( + stream=stream_sub, ack_policy=ack_policy, **kwargs + )(handler) + except BaseException: + self._discard_lease_keys(added_lease_keys) + raise if retry_on_idle_ms is not None and not _internal_retry: # Per-handler retry/dlq stream keys (handler_suffix / retry_stream @@ -805,6 +979,8 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: no_ack=no_ack, ack_policy=ack_policy, delete_on_ack=delete_on_ack, + _lease_options=lease_opts, + _lease_keys_added=added_lease_keys, **kwargs, ) @@ -843,10 +1019,21 @@ async def subscribe(self, channel: str, handler, **kwargs) -> Callable: if self._reclaimer_manager is not None: for key in added_reclaimer_keys: self._reclaimer_manager.discard(key) + # Main and retry-stream keepers (the retry subscribe records + # its keys in added_lease_keys once it has succeeded). + self._discard_lease_keys(added_lease_keys) raise + if _parent_lease_keys is not None: + _parent_lease_keys.extend(added_lease_keys) return registered_handler + def _discard_lease_keys(self, keys: list[LeaseKey]) -> None: + if self._lease_manager is not None: + for key in keys: + self._lease_manager.discard(key) + keys.clear() + async def _create_groups(self) -> None: """Create every registered consumer group at its group_start before the subscribers start. FastStream creates groups at "$" for last_id=">", so @@ -948,12 +1135,58 @@ async def _monitor_stream_groups(self) -> None: info.stream_key, e, ) + await self._restart_stopped_subscribers() await self._check_delete_on_ack_groups(client) except asyncio.CancelledError: pass finally: await client.aclose() + async def _restart_stopped_subscribers(self) -> None: + """Restart stream subscribers FastStream stopped (e.g. after a NOGROUP). + + Up to 0.7.5 FastStream's consume loop kept retrying XREADGROUP, so once + this monitor recreated the group, consuming resumed by itself. From + 0.7.6 a NOGROUP stops the subscriber for good ("restart the + application"). Runs right after the groups were ensured above, so + FastStream's own XGROUP CREATE in start() is a BUSYGROUP no-op and the + group keeps the start id chosen here (id="0" on a partial loss). + Subscribers are only started once their stop() has finished (no task + left running), so a restart can't be cancelled by a stop still in + progress. No-op on FastStream versions that never stop them. The cause + of a stop is not known here, so the log says what was restarted, not + why it stopped. + """ + for sub in list(self._started_subscribers.values()): + if not self._running: + return + if getattr(sub, "running", True): + continue + if any(not task.done() for task in getattr(sub, "tasks", ())): + continue # still stopping (graceful wait for in-flight handlers) + stream, group = _stream_group(sub) or ("?", "?") + try: + await sub.start() + except Exception as e: + logger.warning( + "Failed to restart the consumer for stream %s group %s " + "(retrying in %.1fs): %s: %s", + stream, + group, + self._group_monitor_interval_s, + type(e).__name__, + e, + ) + else: + logger.warning( + "Restarted the consumer for stream %s group %s, which " + "FastStream had stopped (it stops a consumer for good on " + "some errors, e.g. NOGROUP from 0.7.6 on); its consumer " + "group exists again.", + stream, + group, + ) + async def _check_delete_on_ack_groups(self, client: Any) -> None: """Handle a consumer group that joined a delete_on_ack stream after startup: turn deletion off for that stream and log an error (once per @@ -986,6 +1219,16 @@ async def _check_delete_on_ack_groups(self, client: Any) -> None: foreign, ) + def _current_message_ids(self) -> list[str] | None: + """Stream ids of the message FastStream is handling in this context (a + list even for one entry, several in batch mode).""" + message = self.broker.context.get_local("message") + raw = getattr(message, "raw_message", None) + ids = raw.get("message_ids") if isinstance(raw, dict) else None + if not ids: + return None + return [i.decode() if isinstance(i, bytes) else str(i) for i in ids] + def _wrap_delete_on_ack( self, handler: Callable, group: str, ack_policy: AckPolicy ) -> Callable: @@ -1009,9 +1252,7 @@ def _wrap_delete_on_ack( """ delete_on_error = ack_policy in (AckPolicy.ACK, AckPolicy.ACK_FIRST) # An object with an async __call__ is async too, not a sync function. - is_async = asyncio.iscoroutinefunction(handler) or asyncio.iscoroutinefunction( - type(handler).__call__ - ) + is_async = is_async_callable(handler) async def delete_current() -> None: # Everything, lookups included, inside the try: raising here would @@ -1043,16 +1284,7 @@ async def delete_current() -> None: @functools.wraps(handler) async def delete_on_ack_handler(*args, **kwargs): try: - if is_async: - result = await handler(*args, **kwargs) - else: - result = await anyio.to_thread.run_sync( - functools.partial(handler, *args, **kwargs) - ) - # A callable that still returned a coroutine: await it here - # rather than drop it (it would never run). - if inspect.isawaitable(result): - result = await result + result = await call_handler(handler, is_async, *args, **kwargs) except Exception: if delete_on_error: await delete_current() diff --git a/sdk/poetry.lock b/sdk/poetry.lock index a0ecb282..cb019707 100644 --- a/sdk/poetry.lock +++ b/sdk/poetry.lock @@ -1063,14 +1063,14 @@ test = ["pytest (>=6)"] [[package]] name = "fast-depends" -version = "3.0.8" +version = "3.0.9" description = "FastDepends - extracted and cleared from HTTP domain logic FastAPI Dependency Injection System. Async and sync are both supported." optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "fast_depends-3.0.8-py3-none-any.whl", hash = "sha256:4c52c8a3907bca46d43e70e4364d6d016872d9a3aae4bc0c1c85e72e0a6a21c7"}, - {file = "fast_depends-3.0.8.tar.gz", hash = "sha256:896b16f79a512b6ea1df721b0aa1708a192a06f964be6597e01fcf5412559101"}, + {file = "fast_depends-3.0.9-py3-none-any.whl", hash = "sha256:a9fda223e4f9f92ade11dabfdc31d20e62fedc25340f1a5b1f01db2ae25ab1d9"}, + {file = "fast_depends-3.0.9.tar.gz", hash = "sha256:c31c3985d0196ea4998da9f5913a76ede75cc4484d89eeafcea1a6bd64fbca82"}, ] [package.dependencies] @@ -1163,20 +1163,20 @@ server = ["authlib (>=1.6.11)", "cyclopts (>=4.0.0)", "exceptiongroup (>=1.2.2)" [[package]] name = "faststream" -version = "0.7.5" -description = "FastStream: the simplest way to work with a messaging queues" +version = "0.7.7" +description = "A thin, typed Python client for Kafka, RabbitMQ, NATS, Redis and MQTT — with AsyncAPI docs and in-memory testing." optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "faststream-0.7.5-py3-none-any.whl", hash = "sha256:27d1bb3a2945cbd914b326d9b78d3e51d6678293835562be962a8a2189abe20b"}, - {file = "faststream-0.7.5.tar.gz", hash = "sha256:a8185385e4fe64f2f0420c8cd38a55a183f4599e5343b77bf16abc9f62ccaadf"}, + {file = "faststream-0.7.7-py3-none-any.whl", hash = "sha256:f15f1c2f389d218e5e9b2d2e22286155faffed20f7e16be8d26ccf5e9f1dd0ce"}, + {file = "faststream-0.7.7.tar.gz", hash = "sha256:dfba491c665f1b962f7c1262496201340bec8711f7d19a786f13cb1e9e2a5bdd"}, ] [package.dependencies] aiokafka = {version = ">=0.9,<0.15", optional = true, markers = "extra == \"kafka\""} -anyio = ">=4.0,<5" -fast-depends = {version = ">=3.0.0", extras = ["pydantic"]} +anyio = ">=4.14.2,<5" +fast-depends = {version = ">=3.0.9", extras = ["pydantic"]} redis = {version = ">=5.0.0,<9.0.0", optional = true, markers = "extra == \"redis\""} typing-extensions = ">=4.12.0" @@ -1185,7 +1185,7 @@ cli = ["typer (>=0.9,!=0.12,<=0.27.1)", "watchfiles (>=0.15.0,<1.2.0)"] confluent = ["confluent-kafka (>=2,!=2.8.1,<3) ; python_full_version < \"3.13.0\"", "confluent-kafka (>=2.6,!=2.8.1,<3) ; python_full_version >= \"3.13.0\""] kafka = ["aiokafka (>=0.9,<0.15)"] mqtt = ["zmqtt (>=0.2.0,<0.3.0)"] -nats = ["nats-py (>=2.12.0,<=3.0.0)"] +nats = ["nats-py (>=2.14.0,<=3.0.0)"] otel = ["opentelemetry-sdk (>=1.24.0,<2.0.0)"] prometheus = ["prometheus-client (>=0.20.0,<0.30.0)"] rabbit = ["aio-pika (>=9,<11)"] @@ -4190,4 +4190,4 @@ otel = ["opentelemetry-api", "opentelemetry-exporter-otlp-proto-grpc", "opentele [metadata] lock-version = "2.1" python-versions = "^3.10" -content-hash = "59e37b5cd75c9dc95bc18b0a99e481f5d52edbc58d6cbfac6bec4c2d9f696c9a" +content-hash = "64c7403f70d2f5fbb8c98aa6040111e1a7d5e7053ed06fc0aed5593b0897ce22" diff --git a/sdk/pyproject.toml b/sdk/pyproject.toml index b4595900..0fa91664 100644 --- a/sdk/pyproject.toml +++ b/sdk/pyproject.toml @@ -34,7 +34,10 @@ classifiers = [ [tool.poetry.dependencies] python = "^3.10" pydantic = "^2.11.4" -faststream = {extras = ["kafka", "redis"], version = ">=0.6,<0.8"} +# Imported directly (sync handlers run in a worker thread); same range as +# the lowest supported FastStream, so it adds no constraint of its own. +anyio = ">=4.0,<5" +faststream = {extras = ["kafka", "redis"], version = ">=0.6.4,<0.8"} click = {version = "^8.0.0", optional = true} jinja2 = {version = "^3.1.6", optional = true} a2a-sdk = {version = "^0.3.0", optional = true} diff --git a/sdk/requirements.txt b/sdk/requirements.txt index 96b64642..c4d2c628 100644 --- a/sdk/requirements.txt +++ b/sdk/requirements.txt @@ -1,4 +1,5 @@ -faststream[kafka]==0.7.5 +faststream[kafka]==0.7.7 pydantic==2.13.5 click==8.5.0 -jinja2==3.1.6 \ No newline at end of file +jinja2==3.1.6 +anyio==4.15.1 diff --git a/sdk/tests/test_redis.py b/sdk/tests/test_redis.py index ad939ed0..a470ce55 100644 --- a/sdk/tests/test_redis.py +++ b/sdk/tests/test_redis.py @@ -9,6 +9,7 @@ """ import asyncio +import inspect import json import logging import uuid @@ -1859,12 +1860,16 @@ async def _make_envelope(body: dict) -> bytes: """Build a FastStream BinaryMessageFormatV1 envelope around a JSON body. Used by the backoff tests to plant a stream entry with a known _retry_count. + ``encode`` is sync in FastStream 0.6 and async from 0.7. """ from faststream.redis.parser.binary import BinaryMessageFormatV1 - return await BinaryMessageFormatV1.encode( + envelope = BinaryMessageFormatV1.encode( message=body, reply_to=None, headers=None, correlation_id="test-cid" ) + if inspect.isawaitable(envelope): + envelope = await envelope + return envelope def _make_backoff_config(stream: str, retry_stream: str, **overrides): @@ -2248,3 +2253,56 @@ async def handler(message): await redis_client.delete(f"eggai.{channel_name}") await redis_client.aclose() await producer.disconnect() + + +@pytest.mark.asyncio +async def test_group_monitor_restarts_subscribers_faststream_stopped(caplog): + """FastStream >= 0.7.6 stops a subscriber for good on NOGROUP. The monitor + restarts subscribers connect() started, once their stop() has finished; + running ones and ones still stopping are left alone, and a failed restart + is logged and retried on the next pass.""" + from types import SimpleNamespace + + transport = RedisTransport() + transport._running = True + + class FakeSubscriber: + def __init__(self, running, tasks=(), fail=False): + self.stream_sub = SimpleNamespace(name="s", group="g") + self.running = running + self.tasks = list(tasks) + self.fail = fail + self.starts = 0 + + async def start(self): + self.starts += 1 + if self.fail: + raise ConnectionError("redis down") + self.running = True + + pending = asyncio.get_running_loop().create_future() + stopped = FakeSubscriber(running=False) + stopping = FakeSubscriber(running=False, tasks=[pending]) + running = FakeSubscriber(running=True) + failing = FakeSubscriber(running=False, fail=True) + for sub in (stopped, stopping, running, failing): + transport._started_subscribers[id(sub)] = sub + + caplog.set_level(logging.WARNING) + await transport._restart_stopped_subscribers() + + assert (stopped.starts, stopping.starts, running.starts) == (1, 0, 0) + assert stopped.running + assert failing.starts == 1 and "Failed to restart the consumer" in caplog.text + assert "Restarted the consumer for stream s group g" in caplog.text + + failing.fail = False + pending.cancel() + await transport._restart_stopped_subscribers() + assert failing.starts == 2 and failing.running + assert stopping.starts == 1 and stopped.starts == 1 + + transport._running = False + stopped.running = False + await transport._restart_stopped_subscribers() + assert stopped.starts == 1 # never restarted once disconnect() began diff --git a/sdk/tests/test_redis_lease_renewal.py b/sdk/tests/test_redis_lease_renewal.py new file mode 100644 index 00000000..c05e2300 --- /dev/null +++ b/sdk/tests/test_redis_lease_renewal.py @@ -0,0 +1,1702 @@ +""" +renew_lease: entries a handler is still working on are renewed (XCLAIM ... +JUSTID), so the SDK reclaimer does not redeliver them while they run. Plus the +max_processing_ms deadline. Integration tests against a real Redis at +localhost:6379, and unit tests of the keeper and the handler wrapper. +""" + +import asyncio +import functools +import logging +import threading +import time +import uuid +from unittest.mock import AsyncMock + +import pytest +import pytest_asyncio +import redis.asyncio as redis +from faststream import AckPolicy + +from eggai import Agent, Channel +from eggai.transport import LeaseLostError, ProcessingTimeoutError, RedisTransport +from eggai.transport.lease import ( + _RENEW_SCRIPT, + _SCAN_PAGE, + LeaseConfig, + LeaseKeeper, + LeaseManager, + resolve_lease_options, + wrap_handler_with_lease, +) +from eggai.transport.middleware_utils import is_async_callable +from eggai.transport.redis import _CONSUMER_INSTANCE + +# Reclaim threshold used by the integration tests: handlers below run for +# several multiples of it, so without a lease they are always redelivered. +IDLE_MS = 500 +RECLAIM_S = 0.2 + + +async def _wait_for(predicate, timeout=10.0, interval=0.05): + deadline = asyncio.get_running_loop().time() + timeout + while True: + if await predicate(): + return + if asyncio.get_running_loop().time() > deadline: + raise AssertionError("condition not met before timeout") + await asyncio.sleep(interval) + + +async def _true(value): + return value + + +def _retries(message) -> int: + """The SDK's _retry_count of a delivery (0 for a first one), whether it is + written as a string or as a number.""" + return int(message.get("_retry_count", 0)) + + +@pytest_asyncio.fixture +async def redis_client(): + client = redis.Redis(host="localhost", port=6379, decode_responses=True) + yield client + await client.aclose() + + +def _names(prefix): + test_id = uuid.uuid4().hex[:8] + agent_name = f"{prefix}-{test_id}" + channel_name = f"{prefix}-ch-{test_id}" + group = f"{agent_name}-handler-1" + return agent_name, channel_name, f"eggai.{channel_name}", group + + +async def _read_as(stream, group, consumer): + """XREADGROUP one entry as ``consumer`` (binary client: FastStream bodies).""" + client = redis.Redis(host="localhost", port=6379) + try: + await client.xreadgroup(group, consumer, {stream: ">"}, count=1) + finally: + await client.aclose() + + +async def _pending_entries(client, stream, group): + return await client.xpending_range(stream, group, min="-", max="+", count=100) + + +class _Tracker: + """Counts handler runs and how many overlap.""" + + def __init__(self): + self.calls: list[dict] = [] + self.active = 0 + self.max_active = 0 + self.done = 0 + + async def run(self, message, seconds): + self.calls.append(message) + self.active += 1 + self.max_active = max(self.max_active, self.active) + try: + await asyncio.sleep(seconds) + finally: + self.active -= 1 + self.done += 1 + + +# -------------------------------------------------------------------------- +# Exactly once with a lease, twice without +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_slow_handler_processed_once_with_lease_renewal(redis_client): + agent_name, channel_name, stream, group = _names("lease-once") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + samples = [] + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(message): + await tracker.run(message, 2.0) # 4x retry_on_idle_ms + + await agent.start() + try: + await channel.publish({"type": "t", "n": 1}) + await _wait_for(lambda: _true(tracker.active == 1)) + # While it runs: idle stays below the threshold, owner and delivery + # count unchanged (JUSTID does not bump it). + for _ in range(8): + await asyncio.sleep(0.2) + (entry,) = await _pending_entries(redis_client, stream, group) + samples.append(entry) + await _wait_for(lambda: _true(tracker.done == 1)) + await asyncio.sleep(1.5) # several more reclaim cycles + finally: + await agent.stop() + + assert len(tracker.calls) == 1 + assert tracker.max_active == 1 + assert all(s["time_since_delivered"] < IDLE_MS for s in samples), samples + assert all(s["times_delivered"] == 1 for s in samples), samples + assert {s["consumer"] for s in samples} == {f"{group}-{_CONSUMER_INSTANCE}"} + assert await _pending_entries(redis_client, stream, group) == [] + assert await redis_client.xlen(f"{stream}.{group}.retry") == 0 + + +@pytest.mark.asyncio +async def test_slow_handler_processed_twice_without_lease_renewal(): + """Contrast (the bug): the same slow handler is redelivered while it runs.""" + agent_name, channel_name, _stream, _group = _names("lease-off") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + ) + async def handler(message): + await tracker.run(message, 2.0) + + await agent.start() + try: + await channel.publish({"type": "t", "n": 1}) + await _wait_for(lambda: _true(len(tracker.calls) >= 2), timeout=5.0) + finally: + await agent.stop() + + assert tracker.max_active == 2 # in parallel + assert _retries(tracker.calls[1]) == 1 + assert transport._lease_manager is None # default: nothing new runs + + +# -------------------------------------------------------------------------- +# Crash recovery still works +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_crashed_consumers_entry_is_still_reclaimed(redis_client): + """An entry read by a consumer that died (no renewals, no ack) is reclaimed + after retry_on_idle_ms and handled via the retry stream, even though this + subscription renews its own leases.""" + agent_name, channel_name, stream, group = _names("lease-crash") + await redis_client.xgroup_create(stream, group, id="$", mkstream=True) + producer = RedisTransport() + await producer.connect() + try: + await producer.publish(stream, {"type": "t", "n": 1}) + finally: + await producer.disconnect() + # The "crashed" worker: read it, then never ack or renew it. + await _read_as(stream, group, "crashed-worker") + + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + seen = [] + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(message): + seen.append(message) + + await agent.start() + try: + await _wait_for(lambda: _true(len(seen) == 1)) + finally: + await agent.stop() + assert seen[0]["n"] == 1 + assert _retries(seen[0]) == 1 + + +@pytest.mark.asyncio +async def test_renewal_stopping_mid_run_lets_the_entry_be_reclaimed(redis_client): + """If renewal stops (what a crash looks like from Redis) the entry is + reclaimed after retry_on_idle_ms, as without the option.""" + agent_name, channel_name, stream, group = _names("lease-stop") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + cancel_on_lease_lost=False, + ) + async def handler(message): + await tracker.run(message, 0.0 if _retries(message) else 3.0) + + await agent.start() + try: + await channel.publish({"type": "t"}) + await _wait_for(lambda: _true(tracker.active == 1)) + assert transport._lease_manager is not None + for task in transport._lease_manager._tasks.values(): + task.cancel() + await _wait_for(lambda: _true(len(tracker.calls) == 2), timeout=5.0) + finally: + await agent.stop() + assert _retries(tracker.calls[1]) == 1 + + +@pytest.mark.asyncio +async def test_stale_entry_of_same_consumer_name_is_not_renewed(redis_client): + """An entry left in this consumer name's PEL by an earlier incarnation (same + host and pid after a container restart) is older than anything this process + reads, so the prefetch scan skips it and the reclaimer recovers it while a + leased handler is running.""" + agent_name, channel_name, stream, group = _names("lease-stale") + consumer = f"{group}-{_CONSUMER_INSTANCE}" + await redis_client.xgroup_create(stream, group, id="$", mkstream=True) + producer = RedisTransport() + await producer.connect() + try: + await producer.publish(stream, {"type": "t", "n": "stale"}) + finally: + await producer.disconnect() + await _read_as(stream, group, consumer) + + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + seen = [] + slow_done = asyncio.Event() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + max_records=10, # prefetch scan on + ) + async def handler(message): + seen.append((message["n"], _retries(message))) + if message["n"] == "slow": + await asyncio.sleep(2.0) + slow_done.set() + + await agent.start() + try: + await channel.publish({"type": "t", "n": "slow"}) + await asyncio.wait_for(slow_done.wait(), timeout=10.0) + await asyncio.sleep(0.5) + finally: + await agent.stop() + assert ("stale", 1) in seen # reclaimed while "slow" held a lease + assert [s for s in seen if s[0] == "slow"] == [("slow", 0)] + + +# -------------------------------------------------------------------------- +# Lost leases +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_lost_lease_cancels_the_handler(redis_client, caplog): + agent_name, channel_name, stream, group = _names("lease-lost") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + started = asyncio.Event() + cancelled = asyncio.Event() + calls = [] + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + max_retries=None, + ) + async def handler(message): + calls.append(_retries(message)) + if len(calls) > 1: + return + started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.set() + raise + + caplog.set_level(logging.ERROR) + await agent.start() + try: + await channel.publish({"type": "t"}) + await asyncio.wait_for(started.wait(), timeout=10.0) + (entry,) = await _pending_entries(redis_client, stream, group) + # Someone else takes the entry (what the reclaimer does after renewals + # failed for longer than retry_on_idle_ms). + await redis_client.xclaim( + stream, group, "someone-else", 0, [entry["message_id"]], justid=True + ) + await asyncio.wait_for(cancelled.wait(), timeout=3.0) + await _wait_for(lambda: _true("LeaseLostError" in caplog.text)) + # Not acked by the cancelled run: still in the PEL, under its new owner. + pending = await _pending_entries(redis_client, stream, group) + assert f"{group}-{_CONSUMER_INSTANCE}" not in {p["consumer"] for p in pending} + finally: + await agent.stop() + assert "Lease lost for entry" in caplog.text + + +@pytest.mark.asyncio +async def test_lost_lease_without_cancel_lets_the_handler_finish(redis_client, caplog): + agent_name, channel_name, stream, group = _names("lease-keep") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + started = asyncio.Event() + finished = asyncio.Event() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + cancel_on_lease_lost=False, + ) + async def handler(message): + if _retries(message): + return + started.set() + await asyncio.sleep(1.0) + finished.set() + + caplog.set_level(logging.ERROR) + await agent.start() + try: + await channel.publish({"type": "t"}) + await asyncio.wait_for(started.wait(), timeout=10.0) + (entry,) = await _pending_entries(redis_client, stream, group) + await redis_client.xack(stream, group, entry["message_id"]) + await asyncio.wait_for(finished.wait(), timeout=5.0) + finally: + await agent.stop() + assert "keeps running" in caplog.text + assert "LeaseLostError" not in caplog.text + + +@pytest.mark.asyncio +async def test_trimmed_entry_is_not_a_lost_lease(redis_client, caplog): + """An entry deleted from the stream can't be redelivered: let it finish.""" + agent_name, channel_name, stream, group = _names("lease-trim") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + started = asyncio.Event() + finished = asyncio.Event() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(message): + started.set() + await asyncio.sleep(1.5) + finished.set() + + caplog.set_level(logging.WARNING) + await agent.start() + try: + await channel.publish({"type": "t"}) + await asyncio.wait_for(started.wait(), timeout=10.0) + (entry,) = await _pending_entries(redis_client, stream, group) + await redis_client.xdel(stream, entry["message_id"]) + await asyncio.wait_for(finished.wait(), timeout=5.0) + await asyncio.sleep(0.3) + finally: + await agent.stop() + assert "was deleted from stream" in caplog.text + assert "LeaseLostError" not in caplog.text + assert await _pending_entries(redis_client, stream, group) == [] + + +# -------------------------------------------------------------------------- +# Retry stream, prefetch, batch +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_retry_stream_delivery_is_renewed(redis_client): + """A redelivery is a new entry on the .retry stream with its own group and + consumer: that one is renewed too.""" + agent_name, channel_name, stream, group = _names("lease-retry") + retry_stream = f"{stream}.{group}.retry" + retry_group = f"{group}-retry" + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + samples = [] + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(message): + if not _retries(message): + raise RuntimeError("transient") # fast failure, NACK -> retry + await tracker.run(message, 2.0) + + await agent.start() + try: + await channel.publish({"type": "t"}) + await _wait_for(lambda: _true(tracker.active == 1)) + for _ in range(8): + await asyncio.sleep(0.2) + (entry,) = await _pending_entries(redis_client, retry_stream, retry_group) + samples.append(entry) + await _wait_for(lambda: _true(tracker.done == 1)) + await asyncio.sleep(1.5) + finally: + await agent.stop() + + assert [_retries(c) for c in tracker.calls] == [1] + assert all(s["time_since_delivered"] < IDLE_MS for s in samples), samples + assert {s["consumer"] for s in samples} == {f"{retry_group}-{_CONSUMER_INSTANCE}"} + assert await _pending_entries(redis_client, retry_stream, retry_group) == [] + + +async def _prefill(stream, group, n): + """Publish n entries before the group's consumer starts (read in one go).""" + producer = RedisTransport() + await producer.connect() + try: + for i in range(n): + await producer.publish(stream, {"type": "t", "n": i}) + finally: + await producer.disconnect() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lease", [True, False]) +async def test_prefetched_entries_are_renewed(redis_client, lease): + """max_records=5 reads the whole backlog into this consumer's PEL; the ones + queued behind a slow handler are renewed too (and redelivered without).""" + agent_name, channel_name, stream, group = _names("lease-prefetch") + await _prefill(stream, group, 3) + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + options = {"renew_lease": True} if lease else {} + + @agent.subscribe( + channel=channel, + group_start="0", + max_records=5, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + **options, + ) + async def handler(message): + await tracker.run(message, 0.8) + + await agent.start() + try: + await _wait_for(lambda: _true(tracker.done >= 3), timeout=10.0) + await asyncio.sleep(1.0) + finally: + await agent.stop() + retried = [c for c in tracker.calls if _retries(c)] + if lease: + assert sorted(c["n"] for c in tracker.calls) == [0, 1, 2] + assert retried == [] + else: + assert retried # the queued ones went idle and were redelivered + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lease", [True, False]) +async def test_concurrent_workers_entries_are_renewed(redis_client, lease): + """max_workers=2: FastStream keeps reading while two handlers run, and the + entries waiting for a worker slot are renewed too (redelivered without).""" + agent_name, channel_name, stream, group = _names("lease-workers") + await _prefill(stream, group, 5) + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + + @agent.subscribe( + channel=channel, + group_start="0", + max_workers=2, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + **({"renew_lease": True} if lease else {}), + ) + async def handler(message): + await tracker.run(message, 0.8) + + await agent.start() + try: + await _wait_for(lambda: _true(tracker.done >= 5), timeout=15.0) + await asyncio.sleep(1.0) + finally: + await agent.stop() + retried = [c for c in tracker.calls if _retries(c)] + if lease: + assert sorted(c["n"] for c in tracker.calls) == [0, 1, 2, 3, 4] + assert retried == [] + else: + assert retried + + +@pytest.mark.asyncio +async def test_lease_renewal_defaults_max_records_to_one(redis_client): + agent_name, channel_name, stream, group = _names("lease-one") + await _prefill(stream, group, 3) + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + tracker = _Tracker() + pending_counts = [] + + @agent.subscribe( + channel=channel, + group_start="0", + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(message): + pending_counts.append(len(await _pending_entries(redis_client, stream, group))) + await tracker.run(message, 0.3) + + await agent.start() + try: + await _wait_for(lambda: _true(tracker.done == 3)) + finally: + await agent.stop() + assert pending_counts == [1, 1, 1] # one entry read at a time + + +@pytest.mark.asyncio +async def test_batch_entries_are_all_renewed(redis_client): + agent_name, channel_name, stream, group = _names("lease-batch") + await _prefill(stream, group, 3) + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + batches = [] + samples = [] + + @agent.subscribe( + channel=channel, + group_start="0", + batch=True, + max_records=10, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(messages: list): + batches.append(messages) + for _ in range(8): + await asyncio.sleep(0.2) + samples.extend(await _pending_entries(redis_client, stream, group)) + + await agent.start() + try: + await _wait_for(lambda: _true(len(batches) >= 1 and len(samples) >= 24)) + await asyncio.sleep(1.0) + finally: + await agent.stop() + assert len(batches) == 1 and len(batches[0]) == 3 + assert all(s["time_since_delivered"] < IDLE_MS for s in samples), samples + + +# -------------------------------------------------------------------------- +# max_processing_ms +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_max_processing_ms_cancels_and_retries(caplog): + agent_name, channel_name, _stream, _group = _names("lease-deadline") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + cancelled = asyncio.Event() + retried = asyncio.Event() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + max_processing_ms=300, + ) + async def handler(message): + if _retries(message) == 1: + retried.set() + return + try: + await asyncio.sleep(30) # hung call + except asyncio.CancelledError: + cancelled.set() + raise + + caplog.set_level(logging.ERROR) + await agent.start() + try: + await channel.publish({"type": "t"}) + await asyncio.wait_for(cancelled.wait(), timeout=10.0) + await asyncio.wait_for(retried.wait(), timeout=10.0) + finally: + await agent.stop() + assert "ProcessingTimeoutError" in caplog.text + + +@pytest.mark.asyncio +async def test_sync_handler_past_its_deadline_is_not_retried_while_running(caplog): + """A sync handler's thread can't be stopped: past max_processing_ms its + entry stays leased until the thread returns, so the retry never overlaps + it, and the timeout is raised only then.""" + agent_name, channel_name, _stream, _group = _names("lease-sync-deadline") + transport = RedisTransport() + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + lock = threading.Lock() + runs: list[tuple[int, float, float]] = [] # (retry_count, start, end) + active = 0 + max_active = 0 + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + max_processing_ms=300, + ) + def handler(message): + nonlocal active, max_active + retry_count = _retries(message) + start = time.monotonic() + with lock: + active += 1 + max_active = max(max_active, active) + try: + if retry_count == 0: + time.sleep(2.0) # 4x retry_on_idle_ms, well past the deadline + finally: + with lock: + active -= 1 + runs.append((retry_count, start, time.monotonic())) + + caplog.set_level(logging.ERROR) + await agent.start() + try: + await channel.publish({"type": "t"}) + await _wait_for(lambda: _true(len(runs) >= 2), timeout=10.0) + finally: + await agent.stop() + + first, retry = runs[0], runs[1] + assert (first[0], retry[0]) == (0, 1) + assert max_active == 1 + assert retry[1] >= first[2] # the retry started after the thread returned + assert "ProcessingTimeoutError" in caplog.text + assert "can't be interrupted" in caplog.text + + +# -------------------------------------------------------------------------- +# Lifecycle +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_agent_stop_cancels_handler_and_stops_renewal(): + agent_name, channel_name, _stream, _group = _names("lease-shutdown") + transport = RedisTransport(graceful_timeout=0.5) + agent = Agent(agent_name, transport=transport) + channel = Channel(channel_name, transport=transport) + started = asyncio.Event() + cancelled = asyncio.Event() + + @agent.subscribe( + channel=channel, + retry_on_idle_ms=IDLE_MS, + retry_reclaim_interval_s=RECLAIM_S, + renew_lease=True, + ) + async def handler(message): + started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.set() + raise + + await agent.start() + manager = transport._lease_manager + assert manager is not None and manager.client is not None + tasks = list(manager._tasks.values()) + assert len(tasks) == 2 # main stream + retry stream + await channel.publish({"type": "t"}) + await asyncio.wait_for(started.wait(), timeout=10.0) + await asyncio.wait_for(agent.stop(), timeout=10.0) + assert cancelled.is_set() + assert all(t.done() for t in tasks) + assert manager.client is None and manager._tasks == {} + assert all(not k.in_flight for k in manager._keepers.values()) + + +# -------------------------------------------------------------------------- +# The renewal script against a real Redis +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_renew_script_statuses(redis_client): + stream = f"lease-script-{uuid.uuid4().hex[:8]}" + group, me = "g", "me" + await redis_client.xgroup_create(stream, group, id="$", mkstream=True) + ids = [await redis_client.xadd(stream, {"k": str(i)}) for i in range(3)] + await redis_client.xreadgroup(group, me, {stream: ">"}) + await asyncio.sleep(0.3) + await redis_client.xclaim(stream, group, "other", 0, [ids[1]], justid=True) + await redis_client.xdel(stream, ids[2]) + script = redis_client.register_script(_RENEW_SCRIPT) + missing = "1-1" + + statuses = await script(keys=[stream], args=[group, me, *ids, missing]) + + assert statuses == [1, 0, 2, 0] # renewed, other owner, trimmed, not pending + pending = { + p["message_id"]: p + for p in await redis_client.xpending_range(stream, group, "-", "+", 10) + } + assert pending[ids[0]]["consumer"] == me + assert pending[ids[0]]["time_since_delivered"] < 100 # idle reset + assert pending[ids[0]]["times_delivered"] == 1 # JUSTID: count unchanged + assert pending[ids[1]]["consumer"] == "other" # not stolen back + await redis_client.delete(stream) + + +# -------------------------------------------------------------------------- +# Keeper / wrapper unit tests (no Redis) +# -------------------------------------------------------------------------- + + +def _keeper(**overrides): + config = { + "stream": "s", + "group": "g", + "consumer": "c", + "interval_s": 0.05, + **overrides, + } + return LeaseKeeper(LeaseConfig(**config)) + + +def _manager(script): + manager = LeaseManager("redis://unused") + manager.client = object() # type: ignore[assignment] + manager.script = script + manager.running = True + return manager + + +@pytest.mark.asyncio +async def test_renewal_failure_is_logged_and_retried(caplog): + keeper = _keeper() + keeper.begin(["5-0"]) + script = AsyncMock(side_effect=[ConnectionError("redis down"), [1], [1], [1]]) + manager = _manager(script) + caplog.set_level(logging.WARNING) + task = asyncio.create_task(keeper.run(manager)) + try: + await _wait_for(lambda: _true(script.await_count >= 2), timeout=2.0) + finally: + manager.running = False + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert "Lease renewal failed" in caplog.text + assert "stream=s group=g consumer=c ids=['5-0']" in caplog.text + assert "ConnectionError: redis down" in caplog.text + + +@pytest.mark.asyncio +async def test_hung_renewal_call_is_bounded(caplog): + keeper = _keeper(interval_s=0.1) + keeper.begin(["5-0"]) + + async def hang(**_kwargs): + await asyncio.sleep(10) + + manager = _manager(hang) + caplog.set_level(logging.WARNING) + task = asyncio.create_task(keeper.run(manager)) + try: + await _wait_for(lambda: _true("TimeoutError" in caplog.text), timeout=2.0) + finally: + manager.running = False + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_renewal_skips_when_nothing_in_flight(): + keeper = _keeper() + script = AsyncMock(return_value=[]) + await keeper.renew_once(object(), script) + script.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lost_status_cancels_only_its_invocation(): + keeper = _keeper() + blocker = asyncio.Event() + lost_task = asyncio.create_task(blocker.wait()) + kept_task = asyncio.create_task(blocker.wait()) + lost, kept = keeper.begin(["1-0"]), keeper.begin(["2-0"]) + lost.task, kept.task = lost_task, kept_task + + await keeper.renew_once(object(), AsyncMock(return_value=[0, 1])) + await asyncio.sleep(0) + + assert lost_task.cancelled() and lost.reason == "lease_lost" + assert not kept_task.done() and kept.reason is None + kept_task.cancel() + + +@pytest.mark.asyncio +async def test_finished_failures_are_not_renewed_by_the_prefetch_scan(): + """A failed entry stays in the PEL for retry: never renew it.""" + keeper = _keeper(scan_prefetched=True) + older = keeper.begin(["1-0"]) + failed = keeper.begin(["2-0"]) + keeper.end(failed) # handler raised: NACKed, still pending + client = AsyncMock() + client.xpending_range.return_value = [ + {"message_id": "1-0"}, + {"message_id": "2-0"}, + {"message_id": "3-0"}, # prefetched, queued + ] + script = AsyncMock(return_value=[1, 1]) + + await keeper.renew_once(client, script) + + assert script.await_args.kwargs["args"] == ["g", "c", "1-0", "3-0"] + keeper.end(older) + + +@pytest.mark.asyncio +async def test_finished_entries_no_longer_pending_are_forgotten(): + """Beside a long run, every finished id would be kept until it ends: the + scan drops the ones that are no longer pending (acked), keeps the NACKed + ones, and keeps ids that finished while the scan was running.""" + keeper = _keeper(scan_prefetched=True) + long_run = keeper.begin(["1-0"]) + for i in range(2, 502): # 500 short runs finish beside it + keeper.end(keeper.begin([f"{i}-0"])) + assert len(keeper._finished) == 500 + + async def xpending_range(*_args, **_kwargs): + # Finishes while this scan is in flight; the scan may miss it. + keeper.end(keeper.begin(["600-0"])) + return [{"message_id": "1-0"}, {"message_id": "7-0"}] # 7-0 NACKed + + client = AsyncMock() + client.xpending_range.side_effect = xpending_range + + await keeper.renew_once(client, AsyncMock(return_value=[1])) + + assert keeper._finished == {"7-0", "600-0"} + keeper.end(long_run) + + +@pytest.mark.asyncio +async def test_queued_entry_acked_during_the_round_is_not_reported_lost(caplog): + """An entry queued at scan time that starts, finishes and is acked while the + renewal script is in flight comes back 'not pending': not a lost lease.""" + keeper = _keeper(scan_prefetched=True) + long_run = keeper.begin(["1-0"]) + client = AsyncMock() + client.xpending_range.return_value = [{"message_id": "1-0"}, {"message_id": "2-0"}] + + async def script(keys, args): + keeper.end(keeper.begin(["2-0"])) # starts and is acked mid-round + return [1 if msg_id == "1-0" else 0 for msg_id in args[2:]] + + caplog.set_level(logging.ERROR) + await keeper.renew_once(client, script) + + assert "Lease lost" not in caplog.text + assert keeper._lost_queued == set() + keeper.end(long_run) + + +@pytest.mark.asyncio +async def test_queued_entry_lost_before_start_raises_without_running(): + keeper = _keeper(scan_prefetched=True) + running = keeper.begin(["1-0"]) + client = AsyncMock() + client.xpending_range.return_value = [{"message_id": "1-0"}, {"message_id": "2-0"}] + await keeper.renew_once(client, AsyncMock(return_value=[1, 0])) + keeper.end(running) + called = [] + + async def handler(message): + called.append(message) + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["2-0"], + max_processing_ms=None, + ) + with pytest.raises(LeaseLostError) as exc: + await wrapped({"n": 2}) + assert called == [] + assert exc.value.message_ids == ("2-0",) + assert not keeper.in_flight + + +@pytest.mark.asyncio +async def test_wrapper_passes_results_and_errors_through_and_releases(): + keeper = _keeper() + + async def ok(message): + return message["n"] * 2 + + async def boom(message): + raise ValueError("handler failed") + + def sync_ok(message): + return "sync" + + def wrap(handler): + return wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["7-0"], + max_processing_ms=None, + ) + + assert await wrap(ok)({"n": 21}) == 42 + assert await wrap(sync_ok)({"n": 1}) == "sync" + with pytest.raises(ValueError, match="handler failed"): + await wrap(boom)({"n": 1}) + assert not keeper.in_flight + + +@pytest.mark.asyncio +async def test_wrapper_forwards_outer_cancellation(): + keeper = _keeper() + started = asyncio.Event() + inner_cancelled = asyncio.Event() + + async def handler(message): + started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + inner_cancelled.set() + raise + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["7-0"], + max_processing_ms=None, + ) + outer = asyncio.create_task(wrapped({})) + await started.wait() + outer.cancel() + with pytest.raises(asyncio.CancelledError): + await outer + assert inner_cancelled.is_set() + assert not keeper.in_flight + + +@pytest.mark.asyncio +async def test_deadline_raises_processing_timeout_error(): + async def handler(message): + await asyncio.sleep(30) + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=None, + message_ids=lambda: None, + max_processing_ms=50, + ) + with pytest.raises(ProcessingTimeoutError) as exc: + await wrapped({}) + assert exc.value.max_processing_ms == 50 + assert isinstance(exc.value, TimeoutError) + + +@pytest.mark.parametrize("reason", ["timeout", "lease_lost"]) +@pytest.mark.asyncio +async def test_sync_handler_is_never_abandoned(reason): + """A sync handler past its deadline (or with a lost lease) is waited for, + not abandoned: past the deadline it is still renewed until its thread + returns, and the error is raised only then, with the result discarded.""" + keeper = _keeper() + returned = threading.Event() + renewed_after_return: list[bool] = [] + + def handler(message): + time.sleep(0.5) + returned.set() + return "late result" + + async def script(keys, args): + renewed_after_return.append(returned.is_set()) + return [0 if reason == "lease_lost" else 1] * (len(args) - 2) + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["7-0"], + max_processing_ms=100 if reason == "timeout" else None, + ) + run = asyncio.create_task(wrapped({})) + await asyncio.sleep(0.2) # past the deadline; the thread is still running + await keeper.renew_once(object(), script) + assert not run.done() # not abandoned + expected = ProcessingTimeoutError if reason == "timeout" else LeaseLostError + with pytest.raises(expected): + await run + assert returned.is_set() # raised only after the thread returned + if reason == "timeout": + assert renewed_after_return == [False] # still renewed past the deadline + assert not keeper.in_flight + + +# -------------------------------------------------------------------------- +# Option validation and defaults +# -------------------------------------------------------------------------- + + +def test_default_interval_is_a_third_of_retry_on_idle_ms(): + opts = resolve_lease_options({"renew_lease": True}, 900) + assert opts.interval_ms == 300 + assert opts.cancel_on_lease_lost is True + assert resolve_lease_options({}, 900).renew_lease is False + + +def test_default_interval_must_fit_below_a_tiny_retry_on_idle_ms(): + assert resolve_lease_options({"renew_lease": True}, 2).interval_ms == 1 + with pytest.raises(ValueError, match="retry_on_idle_ms >= 2"): + resolve_lease_options({"renew_lease": True}, 1) + + +def test_in_flight_is_read_only(): + keeper = _keeper() + inv = keeper.begin(["1-0"]) + with pytest.raises(TypeError): + keeper.in_flight["2-0"] = inv # type: ignore[index] + assert list(keeper.in_flight) == ["1-0"] + keeper.end(inv) + + +def test_stream_group_of_a_subscriber_without_a_stream_name(): + from types import SimpleNamespace + + from eggai.transport.redis import _stream_group + + def sub(**stream_sub): + return SimpleNamespace(stream_sub=SimpleNamespace(**stream_sub)) + + assert _stream_group(sub(name="s", group="g")) == ("s", "g") + assert _stream_group(sub(group="g")) is None + assert _stream_group(sub(name="s", group=None)) is None + assert _stream_group(SimpleNamespace()) is None + + +@pytest.mark.parametrize( + ("options", "match"), + [ + ({"renew_lease": True}, "renew_lease requires retry_on_idle_ms"), + ( + { + "renew_lease": True, + "retry_on_idle_ms": 500, + "renew_lease_interval_ms": 500, + }, + "must be < retry_on_idle_ms", + ), + ( + { + "renew_lease": True, + "retry_on_idle_ms": 500, + "renew_lease_interval_ms": 0, + }, + "must be a positive int", + ), + ( + {"retry_on_idle_ms": 500, "renew_lease_interval_ms": 100}, + "require renew_lease=True", + ), + ( + {"retry_on_idle_ms": 500, "cancel_on_lease_lost": False}, + "require renew_lease=True", + ), + ({"max_processing_ms": 1000}, "max_processing_ms requires retry_on_idle_ms"), + ( + {"retry_on_idle_ms": 500, "max_processing_ms": 0}, + "max_processing_ms must be a positive int", + ), + ({"renew_lease": "yes", "retry_on_idle_ms": 500}, "True or False"), + ( + { + "renew_lease": True, + "retry_on_idle_ms": 500, + "renew_lease_interval_ms": True, + }, + "must be a positive int", + ), + ( + {"retry_on_idle_ms": 500, "max_processing_ms": True}, + "max_processing_ms must be a positive int", + ), + ( + { + "renew_lease": True, + "retry_on_idle_ms": 500, + "cancel_on_lease_lost": "no", + }, + "cancel_on_lease_lost must be True or False", + ), + ], +) +def test_agent_subscribe_validates_lease_options(options, match): + agent = Agent("lease-validation", transport=RedisTransport()) + channel = Channel("lease-validation", transport=agent._get_transport()) + + with pytest.raises(ValueError, match=match): + + @agent.subscribe(channel=channel, **options) + async def handler(message): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("options", "match"), + [ + ({"renew_lease": True}, "renew_lease requires retry_on_idle_ms"), + ( + {"renew_lease": True, "retry_on_idle_ms": 500, "no_ack": True}, + "incompatible with no_ack", + ), + ( + { + "renew_lease": True, + "retry_on_idle_ms": 500, + "ack_policy": AckPolicy.MANUAL, + }, + "incompatible with ack_policy=AckPolicy.MANUAL", + ), + ( + {"renew_lease": True, "retry_on_idle_ms": 500, "group": "g"}, + "requires a consumer name", + ), + ], +) +async def test_transport_subscribe_validates_lease_options(options, match): + transport = RedisTransport() + if "group" not in options: + options = {"handler_id": "lease-val-h", **options} + + async def handler(message): + pass + + with pytest.raises(ValueError, match=match): + await transport.subscribe("lease-val", handler, **options) + assert transport._lease_manager is None or not transport._lease_manager._keepers + + +# -------------------------------------------------------------------------- +# Cancellation, registration and rollback edge cases +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_outer_cancel_while_draining_a_timed_out_handler_is_not_swallowed(): + """Agent stop while a timed-out handler is still unwinding: the wrapper's + own cancellation propagates instead of turning into ProcessingTimeoutError + (independent of Task.cancelling(), which Python 3.10 lacks).""" + unwinding = asyncio.Event() + + async def handler(message): + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + unwinding.set() + await asyncio.sleep(0.3) # slow cleanup after the deadline + raise + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=None, + message_ids=lambda: ["7-0"], + max_processing_ms=50, + ) + outer = asyncio.create_task(wrapped({})) + await asyncio.wait_for(unwinding.wait(), timeout=2.0) + outer.cancel() + with pytest.raises(asyncio.CancelledError): + await outer + + +@pytest.mark.asyncio +async def test_handler_cancelling_itself_propagates_cancelled_error(): + keeper = _keeper() + + async def handler(message): + raise asyncio.CancelledError + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["7-0"], + max_processing_ms=None, + ) + with pytest.raises(asyncio.CancelledError): + await wrapped({}) + assert not keeper.in_flight + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ids", [None, [], "raise"]) +async def test_missing_message_ids_run_unleased_and_are_logged(ids, caplog): + keeper = _keeper() + seen = [] + + def message_ids(): + if ids == "raise": + raise KeyError("message_ids") + return ids + + async def handler(message): + seen.append(dict(keeper.in_flight)) + return "ok" + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=message_ids, + max_processing_ms=None, + ) + caplog.set_level(logging.WARNING) + assert await wrapped({}) == "ok" + assert seen == [{}] # never registered with the keeper + assert "without a lease" in caplog.text + + +def test_is_async_callable(): + async def coro(message): + pass + + def sync(message): + pass + + class AsyncCallable: + async def __call__(self, message): + pass + + assert is_async_callable(coro) + assert is_async_callable(functools.partial(coro)) + assert is_async_callable(functools.partial(functools.partial(coro))) + assert is_async_callable(AsyncCallable()) + assert is_async_callable(functools.partial(AsyncCallable())) + assert is_async_callable(AsyncMock()) + assert not is_async_callable(sync) + assert not is_async_callable(functools.partial(sync)) + + +@pytest.mark.asyncio +async def test_async_partial_handler_runs_on_the_event_loop(): + loop = asyncio.get_running_loop() + + async def handler(message, factor): + assert asyncio.get_running_loop() is loop + return message["n"] * factor + + wrapped = wrap_handler_with_lease( + functools.partial(handler, factor=3), + stream="s", + group="g", + keeper=_keeper(), + message_ids=lambda: ["7-0"], + max_processing_ms=None, + ) + assert await wrapped({"n": 2}) == 6 + + +@pytest.mark.asyncio +async def test_prune_uses_state_after_the_renewal_round(): + """Runs that end while a renewal call is pending are pruned afterwards.""" + keeper = _keeper() + running = keeper.begin(["1-0"]) + failed = keeper.begin(["2-0"]) + keeper.end(failed) # NACKed, remembered while "1-0" runs + assert keeper._finished == {"2-0"} + + async def script(**_kwargs): + keeper.end(running) # the last run ends mid-round + return [1] + + await keeper.renew_once(object(), script) + assert keeper._finished == set() and keeper._lost_queued == set() + + +@pytest.mark.asyncio +async def test_prefetch_scan_pages_through_a_large_pel(): + keeper = _keeper(scan_prefetched=True) + run = keeper.begin(["1-0"]) + first = [{"message_id": f"{i}-0"} for i in range(1, _SCAN_PAGE + 1)] + second = [{"message_id": f"{_SCAN_PAGE + 1}-0"}] + client = AsyncMock() + client.xpending_range.side_effect = [first, second] + + async def script(keys, args): + return [1] * (len(args) - 2) + + await keeper.renew_once(client, script) + + calls = client.xpending_range.await_args_list + assert [c.kwargs["min"] for c in calls] == ["1-0", f"({_SCAN_PAGE}-0"] + assert all(c.kwargs["count"] == _SCAN_PAGE for c in calls) + keeper.end(run) + + +@pytest.mark.asyncio +async def test_renewal_runs_at_a_fixed_rate_despite_slow_rounds(): + keeper = _keeper(interval_s=0.4) + keeper.begin(["5-0"]) + loop = asyncio.get_running_loop() + starts: list[float] = [] + + async def slow_script(**_kwargs): + starts.append(loop.time()) + await asyncio.sleep(0.15) # below the 0.2 s call bound + return [1] + + manager = _manager(slow_script) + task = asyncio.create_task(keeper.run(manager)) + try: + await _wait_for(lambda: _true(len(starts) >= 5), timeout=5.0) + finally: + manager.running = False + task.cancel() + await asyncio.gather(task, return_exceptions=True) + gaps = sorted(b - a for a, b in zip(starts, starts[1:], strict=False)) + # interval + round time would be ~0.55 s; fixed rate keeps it ~0.4 s. + assert gaps[len(gaps) // 2] < 0.48, gaps + + +@pytest.mark.asyncio +async def test_manager_add_reuses_same_config_refuses_a_different_one(): + manager = LeaseManager("redis://unused") + config = LeaseConfig(stream="s", group="g", consumer="c", interval_s=0.1) + key, keeper, created = manager.add(config) + assert created + assert manager.add(config) == (key, keeper, False) + with pytest.raises(ValueError, match="different settings"): + manager.add(LeaseConfig(stream="s", group="g", consumer="c", interval_s=0.2)) + assert manager._keepers[key] is keeper + + +@pytest.mark.asyncio +async def test_manager_start_loads_the_renewal_script(redis_client): + manager = LeaseManager("redis://localhost:6379") + await redis_client.script_flush() + await manager.start() + try: + assert manager.script is not None + assert await redis_client.script_exists(manager.script.sha) == [True] + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_manager_start_reports_refused_scripting_without_failing( + monkeypatch, caplog +): + """A server that refuses Lua scripting is reported once, up front; the + transport still starts (renewals then keep failing, logged).""" + from redis.exceptions import NoPermissionError + + import eggai.transport.lease as lease_mod + + client = AsyncMock() + client.register_script = lambda _src: AsyncMock() + client.script_load.side_effect = NoPermissionError( + "this user has no permissions to run the 'script|load' command" + ) + monkeypatch.setattr(lease_mod.aioredis, "from_url", lambda *a, **k: client) + manager = LeaseManager("redis://unused") + key, _, _ = manager.add( + LeaseConfig(stream="s", group="g", consumer="c", interval_s=0.05) + ) + caplog.set_level(logging.ERROR) + await manager.start() + try: + assert "Lease renewal will not work" in caplog.text + assert "@scripting" in caplog.text + assert manager.running and key in manager._tasks + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_manager_discard_cancels_the_running_task(): + # An empty keeper never calls Redis (start()'s script check only warns). + manager = LeaseManager("redis://unused:1") + key, _, _ = manager.add( + LeaseConfig(stream="s", group="g", consumer="c", interval_s=0.05) + ) + await manager.start() + try: + task = manager._tasks[key] + manager.discard(key) + await asyncio.gather(task, return_exceptions=True) + assert task.cancelled() + assert key not in manager._keepers and key not in manager._tasks + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_lease_subscription_added_after_connect_is_renewed(redis_client): + """A second subscription on an already connected transport gets its + renewal task straight away, without another connect().""" + transport = RedisTransport() + channel = f"lease-late-{uuid.uuid4().hex[:8]}" + + async def handler(message): + pass + + await transport.subscribe( + channel, + handler, + handler_id=f"{channel}-a", + retry_on_idle_ms=IDLE_MS, + renew_lease=True, + ) + await transport.connect() + try: + manager = transport._lease_manager + assert manager is not None and len(manager._tasks) == 2 # main + retry + await transport.subscribe( + channel, + handler, + handler_id=f"{channel}-b", + retry_on_idle_ms=IDLE_MS, + renew_lease=True, + ) + assert len(manager._keepers) == 4 + assert set(manager._tasks) == set(manager._keepers) + assert all(not t.done() for t in manager._tasks.values()) + finally: + await transport.disconnect() + + +@pytest.mark.asyncio +async def test_failed_subscribe_rolls_back_main_and_retry_keepers(monkeypatch): + transport = RedisTransport() + calls = {"n": 0} + real_setup = transport._setup_reclaimer + + def failing_setup(**kwargs): + calls["n"] += 1 + if calls["n"] == 2: # the retry-stream reclaimer, after both keepers + raise RuntimeError("boom") + return real_setup(**kwargs) + + monkeypatch.setattr(transport, "_setup_reclaimer", failing_setup) + + async def handler(message): + pass + + with pytest.raises(RuntimeError, match="boom"): + await transport.subscribe( + "lease-rollback", + handler, + handler_id="lease-rollback-h", + consumer="custom-consumer", + retry_on_idle_ms=IDLE_MS, + renew_lease=True, + ) + assert transport._lease_manager is not None + assert transport._lease_manager._keepers == {} + + +@pytest.mark.asyncio +async def test_rollback_keeps_a_keeper_registered_by_an_earlier_subscribe( + monkeypatch, +): + transport = RedisTransport() + + async def handler(message): + pass + + options = { + "handler_id": "lease-shared-h", + "retry_on_idle_ms": IDLE_MS, + "renew_lease": True, + } + await transport.subscribe("lease-shared", handler, **options) + keepers = dict(transport._lease_manager._keepers) + assert len(keepers) == 2 + + def failing_setup(**kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(transport, "_setup_reclaimer", failing_setup) + with pytest.raises(RuntimeError, match="boom"): + await transport.subscribe("lease-shared", handler, **options) + assert transport._lease_manager._keepers == keepers + + +@pytest.mark.asyncio +async def test_lost_entry_in_a_batch_names_the_lost_ids(): + """A batch is one unit: one lost entry cancels it, and the error says which + entry left the PEL.""" + keeper = _keeper() + started = asyncio.Event() + + async def handler(messages): + started.set() + await asyncio.sleep(30) + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["1-0", "2-0"], + max_processing_ms=None, + ) + run = asyncio.create_task(wrapped([{}, {}])) + await asyncio.wait_for(started.wait(), timeout=2.0) + await keeper.renew_once(object(), AsyncMock(return_value=[0, 1])) + + with pytest.raises(LeaseLostError) as exc: + await asyncio.wait_for(run, timeout=2.0) + assert exc.value.message_ids == ("1-0", "2-0") + assert exc.value.lost_ids == ("1-0",) + assert "1-0" in str(exc.value) and "batch of 1-0, 2-0" in str(exc.value) + assert not keeper.in_flight + + +@pytest.mark.asyncio +async def test_run_cancelled_before_its_task_exists_never_starts(monkeypatch): + """If the keeper cancels an invocation before the wrapper has attached the + handler task, the handler still doesn't run.""" + keeper = _keeper() + real_begin = keeper.begin + + def begin_then_lose(ids, **kwargs): + inv = real_begin(ids, **kwargs) + inv.lost_ids.append("7-0") + inv.cancel("lease_lost") # task not attached yet: only the reason is set + return inv + + monkeypatch.setattr(keeper, "begin", begin_then_lose) + called = [] + + async def handler(message): + called.append(message) + + wrapped = wrap_handler_with_lease( + handler, + stream="s", + group="g", + keeper=keeper, + message_ids=lambda: ["7-0"], + max_processing_ms=None, + ) + with pytest.raises(LeaseLostError): + await wrapped({}) + assert called == [] + assert not keeper.in_flight + + +@pytest.mark.asyncio +async def test_rounds_slower_than_the_interval_still_pause_and_warn_once( + monkeypatch, caplog +): + import eggai.transport.lease as lease_mod + + keeper = _keeper(interval_s=0.1) + keeper.begin(["5-0"]) + real_sleep = asyncio.sleep + delays: list[float] = [] + rounds = 0 + runner: list[asyncio.Task] = [] + + async def slow_round(client, script): + nonlocal rounds + rounds += 1 + await real_sleep(0.12) # longer than the 0.1 s interval + + async def recording_sleep(delay, *args, **kwargs): + if asyncio.current_task() in runner: # only the renewal loop's waits + delays.append(delay) + await real_sleep(delay, *args, **kwargs) + + monkeypatch.setattr(keeper, "renew_once", slow_round) + monkeypatch.setattr(lease_mod.asyncio, "sleep", recording_sleep) + manager = _manager(AsyncMock()) + caplog.set_level(logging.WARNING) + task = asyncio.create_task(keeper.run(manager)) + runner.append(task) + try: + await _wait_for(lambda: _true(rounds >= 4), timeout=3.0) + finally: + manager.running = False + task.cancel() + await asyncio.gather(task, return_exceptions=True) + # delays[0] is the initial interval; every later wait follows a slow + # round and is the minimum pause, never 0 (back-to-back rounds). + assert len(delays) >= 4, delays + assert delays[1:] == [pytest.approx(0.01)] * len(delays[1:]), delays + assert caplog.text.count("longer than the") == 1