diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 6b061b33..8b8a7ab4 100644 --- a/docs/guide/cli.md +++ b/docs/guide/cli.md @@ -47,10 +47,11 @@ We have two options for this: ### Acknowledgements -The taskiq supports three types of acknowledgements: +The taskiq supports four types of acknowledgements: * `when_received` - task is acknowledged when it is **received** by the worker. * `when_executed` - task is acknowledged right after it is **executed** by the worker. * `when_saved` - task is acknowledged when the result of execution is saved in the result backend. +* `manual` - task is acknowledged by calling `Context.ack()` inside the task. This can be configured using `--ack-type` parameter. For example: @@ -69,9 +70,16 @@ async def best_effort_task() -> None: @broker.task(ack_type="when_saved") async def critical_task() -> None: ... + + +@broker.task(ack_type="manual") +async def manually_acked_task(context: Context = TaskiqDepends()) -> None: + ... + 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`. ### Type casts @@ -159,7 +167,7 @@ The number of signals before a hard kill can be configured with the `--hardkill- * `--no-propagate-errors` - if this parameter is enabled, exceptions won't be thrown in generator dependencies. * `--receiver` - python path to custom receiver class. * `--receiver_arg` - custom args for receiver. -* `--ack-type` - Type of acknowledgement. This parameter is used to set when to acknowledge the task. Possible values are `when_received`, `when_executed`, `when_saved`. Default is `when_saved`. +* `--ack-type` - Type of acknowledgement. This parameter is used to set when to acknowledge the task. Possible values are `when_received`, `when_executed`, `when_saved`, `manual`. Default is `when_saved`. * `--max-tasks-per-child` - maximum number of tasks to be executed by a single worker process before restart. * `--max-fails` - Maximum number of child process exits. * `--shutdown-timeout` - maximum amount of time for graceful broker's shutdown in seconds (default 5). diff --git a/taskiq/acks.py b/taskiq/acks.py index f8179755..894a9160 100644 --- a/taskiq/acks.py +++ b/taskiq/acks.py @@ -4,6 +4,8 @@ from pydantic import BaseModel +from taskiq.utils import maybe_awaitable + @enum.unique class AcknowledgeType(str, enum.Enum): @@ -19,6 +21,9 @@ class AcknowledgeType(str, enum.Enum): # acknowledged when the task will be saved # only after it's saved in the result backend. WHEN_SAVED = "when_saved" + # This option means that the task is responsible + # for acknowledging the message through Context.ack. + MANUAL = "manual" def parse_acknowledge_type(value: Any) -> AcknowledgeType: @@ -50,3 +55,25 @@ class AckableMessage(BaseModel): data: bytes ack: Callable[[], None | Awaitable[None]] + + +class AckController: + """Controls acknowledgement state for a received message.""" + + def __init__(self, ack: Callable[[], None | Awaitable[None]] | None) -> None: + self._ack = ack + self.is_acked = False + + @property + def is_ackable(self) -> bool: + """Whether the current message supports acknowledgement.""" + return self._ack is not None + + async def ack(self) -> None: + """Acknowledge the current message once.""" + if self._ack is None: + raise RuntimeError("Current message is not ackable.") + if self.is_acked: + return + await maybe_awaitable(self._ack()) + self.is_acked = True diff --git a/taskiq/context.py b/taskiq/context.py index f9f1d0ee..48058ebb 100644 --- a/taskiq/context.py +++ b/taskiq/context.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING from taskiq.abc.broker import AsyncBroker +from taskiq.acks import AckController from taskiq.exceptions import NoResultError, TaskRejectedError from taskiq.message import TaskiqMessage @@ -11,12 +12,38 @@ class Context: """Context class.""" - def __init__(self, message: TaskiqMessage, broker: AsyncBroker) -> None: + def __init__( + self, + message: TaskiqMessage, + broker: AsyncBroker, + ack_controller: AckController | None = None, + ) -> None: self.message = message self.broker = broker + self._ack_controller = ack_controller self.state: TaskiqState = None # type: ignore self.state = broker.state + @property + def is_ackable(self) -> bool: + """Whether the current message supports acknowledgement.""" + return self._ack_controller is not None and self._ack_controller.is_ackable + + @property + 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 + + async def ack(self) -> None: + """ + Acknowledge current message. + + :raises RuntimeError: if current broker message is not ackable. + """ + if self._ack_controller is None: + raise RuntimeError("Current message is not ackable.") + await self._ack_controller.ack() + async def requeue(self) -> None: """ Requeue task. diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index 8773c68e..871d13d6 100644 --- a/taskiq/receiver/receiver.py +++ b/taskiq/receiver/receiver.py @@ -15,7 +15,7 @@ from taskiq.abc.broker import AckableMessage, AsyncBroker from taskiq.abc.middleware import TaskiqMiddleware -from taskiq.acks import AcknowledgeType, parse_acknowledge_type +from taskiq.acks import AckController, AcknowledgeType, parse_acknowledge_type from taskiq.context import Context from taskiq.exceptions import NoResultError from taskiq.message import TaskiqMessage @@ -116,6 +116,9 @@ async def callback( # noqa: C901, PLR0912 :param raise_err: raise an error if cannot save result in result_backend. """ + ack_controller = AckController( + message.ack if isinstance(message, AckableMessage) else None, + ) message_data = message.data if isinstance(message, AckableMessage) else message try: taskiq_msg = self.broker.formatter.loads(message=message_data) @@ -155,22 +158,17 @@ async def callback( # noqa: C901, PLR0912 taskiq_msg.task_id, ) - if ack_time == AcknowledgeType.WHEN_RECEIVED and isinstance( - message, - AckableMessage, - ): - await maybe_awaitable(message.ack()) + if ack_time == AcknowledgeType.WHEN_RECEIVED and ack_controller.is_ackable: + await ack_controller.ack() result = await self.run_task( target=task.original_func, message=taskiq_msg, + ack_controller=ack_controller, ) - if ack_time == AcknowledgeType.WHEN_EXECUTED and isinstance( - message, - AckableMessage, - ): - await maybe_awaitable(message.ack()) + if ack_time == AcknowledgeType.WHEN_EXECUTED and ack_controller.is_ackable: + await ack_controller.ack() for middleware in reversed(self.broker.middlewares): if middleware.__class__.post_execute != TaskiqMiddleware.post_execute: @@ -193,11 +191,8 @@ async def callback( # noqa: C901, PLR0912 if raise_err: raise exc - if ack_time == AcknowledgeType.WHEN_SAVED and isinstance( - message, - AckableMessage, - ): - await maybe_awaitable(message.ack()) + if ack_time == AcknowledgeType.WHEN_SAVED and ack_controller.is_ackable: + await ack_controller.ack() def _get_ack_time(self, message: TaskiqMessage) -> AcknowledgeType: """ @@ -219,6 +214,7 @@ async def run_task( # noqa: C901, PLR0912, PLR0915 self, target: Callable[..., Any], message: TaskiqMessage, + ack_controller: AckController | None = None, ) -> TaskiqResult[Any]: """ This function actually executes functions. @@ -259,7 +255,7 @@ async def run_task( # noqa: C901, PLR0912, PLR0915 broker_ctx = self.broker.custom_dependency_context broker_ctx.update( { - Context: Context(message, self.broker), + Context: Context(message, self.broker, ack_controller), TaskiqState: self.broker.state, }, ) diff --git a/tests/receiver/test_receiver.py b/tests/receiver/test_receiver.py index 695d5c1a..d724b326 100644 --- a/tests/receiver/test_receiver.py +++ b/tests/receiver/test_receiver.py @@ -16,6 +16,7 @@ from taskiq.abc.result_backend import AsyncResultBackend from taskiq.acks import AcknowledgeType from taskiq.brokers.inmemory_broker import InMemoryBroker +from taskiq.context import Context from taskiq.exceptions import NoResultError, TaskiqResultTimeoutError from taskiq.message import TaskiqMessage from taskiq.receiver import Receiver @@ -523,6 +524,112 @@ def ack_callback() -> None: assert events == ["ack", "task", "post_execute", "save", "post_save"] +async def test_manual_task_ack_from_context() -> None: + """Manual ack lets a task acknowledge the message 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_ackable + assert not context.is_acked + await context.ack() + assert context.is_acked + return 1 + + def ack_callback() -> None: + events.append("ack") + + receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_SAVED) + broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message()) + + await receiver.callback( + AckableMessage( + data=broker_message.message, + ack=ack_callback, + ), + ) + + assert events == ["task", "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] = [] + 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") + await context.ack() + await context.ack() + return 1 + + def ack_callback() -> None: + events.append("ack") + + receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_SAVED) + broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message()) + + await receiver.callback( + AckableMessage( + data=broker_message.message, + ack=ack_callback, + ), + ) + + assert events == ["task", "ack", "post_execute", "save", "post_save"] + + +async def test_manual_task_ack_without_ackable_message_saves_error() -> None: + """Context.ack fails clearly when broker messages do not support acking.""" + events: list[str] = [] + result_backend = _EventResultBackend(events) + broker = ( + InMemoryBroker() + .with_result_backend(result_backend) + .with_middlewares(_EventMiddleware(events)) + ) + + @broker.task(ack_type="manual") + async def my_task(context: Context = Depends()) -> int: + assert not context.is_ackable + await context.ack() + return 1 + + receiver = get_receiver(broker) + broker_message = broker.formatter.dumps( + TaskiqMessage( + task_id="task_id", + task_name=my_task.task_name, + labels={"ack_type": "manual"}, + args=[], + kwargs={}, + ), + ) + + await receiver.callback(broker_message.message) + + result = result_backend.results["task_id"] + assert result.is_err + assert isinstance(result.error, RuntimeError) + assert str(result.error) == "Current message is not ackable." + assert events == ["post_execute", "save", "post_save"] + + async def test_invalid_task_ack_type_raises_before_execution() -> None: """Invalid task ack_type fails before task execution and acknowledgement.""" events: list[str] = [] diff --git a/tests/test_acks.py b/tests/test_acks.py index 899e731a..fefd3f53 100644 --- a/tests/test_acks.py +++ b/tests/test_acks.py @@ -14,6 +14,10 @@ def test_parse_acknowledge_type_from_string() -> None: assert parse_acknowledge_type("when_executed") == AcknowledgeType.WHEN_EXECUTED +def test_parse_manual_acknowledge_type_from_string() -> None: + assert parse_acknowledge_type("manual") == AcknowledgeType.MANUAL + + def test_parse_acknowledge_type_unknown_string() -> None: with pytest.raises(ValueError, match="Unknown acknowledge type value"): parse_acknowledge_type("when_save")