diff --git a/docs/extending-taskiq/broker.md b/docs/extending-taskiq/broker.md index 537227f9..419b1138 100644 --- a/docs/extending-taskiq/broker.md +++ b/docs/extending-taskiq/broker.md @@ -40,9 +40,16 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]: # So you either set here method of a message, # or you can make a closure. ack=message.ack, + # Optionally expose a callback that reports the message + # is still being processed without finally acknowledging it. + ack_progress=message.ack_progress, ) ``` +Tasks can call the optional progress callback through `Context.ack_progress()`. +The callback may be invoked multiple times and should extend or reset the broker's +acknowledgement deadline without marking the message as finally acknowledged. + ## Conventions For brokers, we have several conventions. It's good if your broker implements them. diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 8b8a7ab4..8f471a58 100644 --- a/docs/guide/cli.md +++ b/docs/guide/cli.md @@ -74,12 +74,19 @@ async def critical_task() -> None: @broker.task(ack_type="manual") async def manually_acked_task(context: Context = TaskiqDepends()) -> None: - ... + while still_working: + await do_some_work() + if context.is_ack_progressable: + await context.ack_progress() await context.ack() ``` If a task has no `ack_type` label, the worker-level `--ack-type` value is used. Manual acknowledgement requires a broker that yields `AckableMessage`. +Some brokers can also report that a task is still being processed through +`Context.ack_progress()` to prevent acknowledgement timeouts during long-running tasks. +Currently, this is supported only by the `PushBasedJetStreamBroker` and +`PullBasedJetStreamBroker` implementations from `taskiq-nats`. ### Type casts diff --git a/taskiq/acks.py b/taskiq/acks.py index 894a9160..9721328c 100644 --- a/taskiq/acks.py +++ b/taskiq/acks.py @@ -55,13 +55,19 @@ class AckableMessage(BaseModel): data: bytes ack: Callable[[], None | Awaitable[None]] + ack_progress: Callable[[], None | Awaitable[None]] | None = None class AckController: """Controls acknowledgement state for a received message.""" - def __init__(self, ack: Callable[[], None | Awaitable[None]] | None) -> None: + def __init__( + self, + ack: Callable[[], None | Awaitable[None]] | None, + ack_progress: Callable[[], None | Awaitable[None]] | None = None, + ) -> None: self._ack = ack + self._ack_progress = ack_progress self.is_acked = False @property @@ -69,6 +75,11 @@ def is_ackable(self) -> bool: """Whether the current message supports acknowledgement.""" return self._ack is not None + @property + def is_ack_progressable(self) -> bool: + """Whether the current message supports acknowledgement progress.""" + return self._ack_progress is not None + async def ack(self) -> None: """Acknowledge the current message once.""" if self._ack is None: @@ -77,3 +88,11 @@ async def ack(self) -> None: return await maybe_awaitable(self._ack()) self.is_acked = True + + async def ack_progress(self) -> None: + """Report that the current message is still being processed.""" + if self._ack_progress is None: + raise RuntimeError("Current message does not support ack progress.") + if self.is_acked: + return + await maybe_awaitable(self._ack_progress()) diff --git a/taskiq/context.py b/taskiq/context.py index 48058ebb..3491bfd1 100644 --- a/taskiq/context.py +++ b/taskiq/context.py @@ -34,6 +34,14 @@ def is_acked(self) -> bool: """Whether the current message has already been acknowledged.""" return self._ack_controller is not None and self._ack_controller.is_acked + @property + def is_ack_progressable(self) -> bool: + """Whether the current message supports acknowledgement progress.""" + return ( + self._ack_controller is not None + and self._ack_controller.is_ack_progressable + ) + async def ack(self) -> None: """ Acknowledge current message. @@ -44,6 +52,16 @@ async def ack(self) -> None: raise RuntimeError("Current message is not ackable.") await self._ack_controller.ack() + async def ack_progress(self) -> None: + """ + Report that the current message is still being processed. + + :raises RuntimeError: if acknowledgement progress is not supported. + """ + if self._ack_controller is None: + raise RuntimeError("Current message does not support ack progress.") + await self._ack_controller.ack_progress() + async def requeue(self) -> None: """ Requeue task. diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index 871d13d6..58c5d1b6 100644 --- a/taskiq/receiver/receiver.py +++ b/taskiq/receiver/receiver.py @@ -118,6 +118,7 @@ async def callback( # noqa: C901, PLR0912 """ ack_controller = AckController( message.ack if isinstance(message, AckableMessage) else None, + message.ack_progress if isinstance(message, AckableMessage) else None, ) message_data = message.data if isinstance(message, AckableMessage) else message try: diff --git a/tests/receiver/test_receiver.py b/tests/receiver/test_receiver.py index d724b326..abadf0bd 100644 --- a/tests/receiver/test_receiver.py +++ b/tests/receiver/test_receiver.py @@ -539,6 +539,7 @@ async def test_manual_task_ack_from_context() -> None: async def my_task(context: Context = Depends()) -> int: events.append("task") assert context.is_ackable + assert not context.is_ack_progressable assert not context.is_acked await context.ack() assert context.is_acked @@ -560,6 +561,55 @@ def ack_callback() -> None: assert events == ["task", "ack", "post_execute", "save", "post_save"] +async def test_task_reports_ack_progress_from_context() -> None: + """A task can repeatedly report acknowledgement progress through context.""" + events: list[str] = [] + broker = ( + InMemoryBroker() + .with_result_backend( + _EventResultBackend(events), + ) + .with_middlewares(_EventMiddleware(events)) + ) + + @broker.task(ack_type="manual") + async def my_task(context: Context = Depends()) -> int: + events.append("task") + assert context.is_ack_progressable + await context.ack_progress() + await context.ack_progress() + await context.ack() + await context.ack_progress() + return 1 + + def ack_callback() -> None: + events.append("ack") + + def ack_progress_callback() -> None: + events.append("progress") + + receiver = get_receiver(broker) + broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message()) + + await receiver.callback( + AckableMessage( + data=broker_message.message, + ack=ack_callback, + ack_progress=ack_progress_callback, + ), + ) + + assert events == [ + "task", + "progress", + "progress", + "ack", + "post_execute", + "save", + "post_save", + ] + + async def test_manual_task_ack_is_idempotent() -> None: """Calling Context.ack twice acknowledges the message once.""" events: list[str] = [] diff --git a/tests/test_acks.py b/tests/test_acks.py index fefd3f53..94ea751b 100644 --- a/tests/test_acks.py +++ b/tests/test_acks.py @@ -1,6 +1,37 @@ import pytest -from taskiq.acks import AcknowledgeType, parse_acknowledge_type +from taskiq.acks import AckController, AcknowledgeType, parse_acknowledge_type + + +async def test_ack_progress_can_be_reported_repeatedly() -> None: + events: list[str] = [] + + def ack_callback() -> None: + events.append("ack") + + async def ack_progress_callback() -> None: + events.append("progress") + + controller = AckController(ack_callback, ack_progress_callback) + + assert controller.is_ackable + assert controller.is_ack_progressable + await controller.ack_progress() + await controller.ack_progress() + assert not controller.is_acked + await controller.ack() + await controller.ack_progress() + + assert controller.is_acked + assert events == ["progress", "progress", "ack"] + + +async def test_ack_progress_raises_when_unsupported() -> None: + controller = AckController(lambda: None) + + assert not controller.is_ack_progressable + with pytest.raises(RuntimeError, match="does not support ack progress"): + await controller.ack_progress() def test_parse_acknowledge_type_from_enum() -> None: