diff --git a/taskbadger.yaml b/taskbadger.yaml index fde0ad1..e369e50 100644 --- a/taskbadger.yaml +++ b/taskbadger.yaml @@ -712,6 +712,15 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of: + type: string + minLength: 1 + nullable: true + description: 'ID of the first attempt of the task this task is a retry of. + The first 3 retries of a task don''t count towards the monthly task allocation. + A retry is nested like a child of the original: under it for a top-level + task, otherwise alongside it under the same parent. Retries of a retry + point at the original. Can only be set when the task is created.' name: type: string minLength: 1 @@ -823,6 +832,14 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of: + type: string + nullable: true + description: 'ID of the first attempt of the task this task is a retry of. + The first 3 retries of a task don''t count towards the monthly task allocation. + A retry is nested like a child of the original: under it for a top-level + task, otherwise alongside it under the same parent. Retries of a retry + point at the original. Can only be set when the task is created.' name: type: string description: Name of the task @@ -932,6 +949,15 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of: + type: string + minLength: 1 + nullable: true + description: 'ID of the first attempt of the task this task is a retry of. + The first 3 retries of a task don''t count towards the monthly task allocation. + A retry is nested like a child of the original: under it for a top-level + task, otherwise alongside it under the same parent. Retries of a retry + point at the original. Can only be set when the task is created.' name: type: string minLength: 1 diff --git a/taskbadger/_current_task.py b/taskbadger/_current_task.py index 3613bc9..ac5fb1c 100644 --- a/taskbadger/_current_task.py +++ b/taskbadger/_current_task.py @@ -14,21 +14,23 @@ from contextvars import ContextVar -# (id of the running task, id of its parent or None if it is a root task) -_current: ContextVar[tuple[str, str | None] | None] = ContextVar("taskbadger_current_task", default=None) +# (id of the running task, id of its parent or None if it is a root task, +# id of the first attempt if the running task is a retry, else None) +_current: ContextVar[tuple[str, str | None, str | None] | None] = ContextVar("taskbadger_current_task", default=None) -def enter_task(task_id: str, parent: str = None): +def enter_task(task_id: str, parent: str = None, retry_of: str = None): """Mark `task_id` as the task running in this context. Arguments: task_id: The running task. parent: The running task's own parent, if it has one. + retry_of: The first attempt, if the running task is itself a retry. Returns: A token to pass to `exit_task`. """ - return _current.set((task_id, parent)) + return _current.set((task_id, parent, retry_of)) def exit_task(token) -> None: @@ -49,5 +51,20 @@ def parent_id() -> str | None: current = _current.get() if current is None: return None - task_id, parent = current + task_id, parent, _ = current return parent or task_id + + +def retry_of_id() -> str | None: + """The id to use as `retry_of` for a retry of the task running right now. + + The first attempt: the running task itself, or the task it is a retry of. Retries all point at + the first attempt, so the API can count them per task. `None` outside a tracked task. The + retry's `parent` is still `parent_id()`: under the first attempt for a top-level task, + alongside it for a child. + """ + current = _current.get() + if current is None: + return None + task_id, _, retry_of = current + return retry_of or task_id diff --git a/taskbadger/_integrations.py b/taskbadger/_integrations.py index c9d52ff..c309324 100644 --- a/taskbadger/_integrations.py +++ b/taskbadger/_integrations.py @@ -88,6 +88,12 @@ def parent_of(task) -> str | None: return parent or None +def retry_of(task) -> str | None: + """The id of the first attempt ``task`` is a retry of, or ``None`` (also for ``None``/unset).""" + value = getattr(task, "retry_of", None) if task is not None else None + return value or None + + def match_task_name(task_name: str, includes, excludes) -> bool: """Return True if ``task_name`` should be tracked under the given rules. diff --git a/taskbadger/celery.py b/taskbadger/celery.py index 26c4e67..f397e75 100644 --- a/taskbadger/celery.py +++ b/taskbadger/celery.py @@ -14,10 +14,17 @@ from kombu import serialization from . import sdk -from ._current_task import enter_task, exit_task, parent_id +from ._current_task import enter_task, exit_task, parent_id, retry_of_id from ._error_context import capture_error_data, reset_error_context, start_error_context from ._heartbeat import heartbeat -from ._integrations import TERMINAL_STATES, parent_of, resolve_heartbeat_options, safe_get_task, task_cache +from ._integrations import ( + TERMINAL_STATES, + parent_of, + resolve_heartbeat_options, + retry_of, + safe_get_task, + task_cache, +) from .internal.models import StatusEnum from .mug import Badger from .safe_sdk import create_task_safe, update_task_safe @@ -199,6 +206,16 @@ def _is_workflow_successor(task_name): return any(name == task_name for sig in successors for name in _signature_names(sig)) +def _is_retry_of_running_task(headers): + """Whether the message being published is a retry of the task that is running. + + ``self.retry()`` republishes with the running task's own id, which separates it from an + unrelated task that happens to carry a ``retries`` header. + """ + request = getattr(celery.current_task, "request", None) + return request is not None and request.id == headers.get("id") + + @before_task_publish.connect def task_publish_handler(sender=None, headers=None, body=None, **kwargs): routing_key = kwargs.get("routing_key") @@ -236,6 +253,10 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs): # is itself a tracked task this nests the new task under it. enclosing_task = parent_id() if enclosing_task and not _is_workflow_successor(sender): + if "parent" not in kwargs and headers.get("retries") and _is_retry_of_running_task(headers): + # A Celery retry republishes the same task. It names the first attempt so the API + # can leave the first few retries of it out of the monthly allocation. + kwargs["retry_of"] = retry_of_id() kwargs.setdefault("parent", enclosing_task) name = kwargs.pop("name", headers["task"]) @@ -372,6 +393,9 @@ def _maybe_create_task(signal_sender): # explicit `None`, which asks for a root task. "parent": header_kwargs["parent"] if "parent" in header_kwargs else parent_id(), } + if "parent" not in header_kwargs and getattr(signal_sender.request, "retries", 0): + # An eager retry runs nested inside the attempt that retried it. + create_kwargs["retry_of"] = retry_of_id() heartbeat_interval, stale_timeout = resolve_heartbeat_options( header_kwargs.get("heartbeat_interval", getattr(signal_sender, TB_HEARTBEAT_INTERVAL, None)), header_kwargs.get("stale_timeout", getattr(signal_sender, TB_STALE_TIMEOUT, None)), @@ -407,7 +431,8 @@ def task_prerun_handler(sender=None, **kwargs): # Mark this as the running task so anything it publishes is nested under # it. `_update_task` above has already put it in the cache, so this is a # lookup rather than a fetch. - current_token = enter_task(task_id, parent_of(safe_get_task(task_id))) + running = safe_get_task(task_id) + current_token = enter_task(task_id, parent_of(running), retry_of(running)) sender.request.update({TB_CURRENT_TASK_TOKEN: current_token}) diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index 519cf02..1319d2c 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/taskbadger/internal/models/patched_task_request.py @@ -24,6 +24,10 @@ class PatchedTaskRequest: Attributes: parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of (None | str | Unset): ID of the first attempt of the task this task is a retry of. The first 3 retries + of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under + it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. + Can only be set when the task is created. name (str | Unset): Name of the task queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with @@ -53,6 +57,7 @@ class PatchedTaskRequest: """ parent: None | str | Unset = UNSET + retry_of: None | str | Unset = UNSET name: str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET @@ -77,6 +82,12 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + retry_of: None | str | Unset + if isinstance(self.retry_of, Unset): + retry_of = UNSET + else: + retry_of = self.retry_of + name = self.name queue = self.queue @@ -140,6 +151,8 @@ def to_dict(self) -> dict[str, Any]: field_dict.update({}) if parent is not UNSET: field_dict["parent"] = parent + if retry_of is not UNSET: + field_dict["retry_of"] = retry_of if name is not UNSET: field_dict["name"] = name if queue is not UNSET: @@ -184,6 +197,15 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + def _parse_retry_of(data: object) -> None | str | Unset: + if data is None: + return data + if isinstance(data, Unset): + return data + return cast(None | str | Unset, data) + + retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) + name = d.pop("name", UNSET) queue = d.pop("queue", UNSET) @@ -280,6 +302,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: patched_task_request = cls( parent=parent, + retry_of=retry_of, name=name, queue=queue, external_id=external_id, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index 67a4812..01250a2 100644 --- a/taskbadger/internal/models/task.py +++ b/taskbadger/internal/models/task.py @@ -33,6 +33,10 @@ class Task: public_url (str): parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of (None | str | Unset): ID of the first attempt of the task this task is a retry of. The first 3 retries + of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under + it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. + Can only be set when the task is created. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -70,6 +74,7 @@ class Task: url: str public_url: str parent: None | str | Unset = UNSET + retry_of: None | str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -112,6 +117,12 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + retry_of: None | str | Unset + if isinstance(self.retry_of, Unset): + retry_of = UNSET + else: + retry_of = self.retry_of + queue = self.queue external_id = self.external_id @@ -185,6 +196,8 @@ def to_dict(self) -> dict[str, Any]: ) if parent is not UNSET: field_dict["parent"] = parent + if retry_of is not UNSET: + field_dict["retry_of"] = retry_of if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -249,6 +262,15 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + def _parse_retry_of(data: object) -> None | str | Unset: + if data is None: + return data + if isinstance(data, Unset): + return data + return cast(None | str | Unset, data) + + retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) + queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -352,6 +374,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: url=url, public_url=public_url, parent=parent, + retry_of=retry_of, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index fe22822..481f6c0 100644 --- a/taskbadger/internal/models/task_request.py +++ b/taskbadger/internal/models/task_request.py @@ -25,6 +25,10 @@ class TaskRequest: name (str): Name of the task parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of (None | str | Unset): ID of the first attempt of the task this task is a retry of. The first 3 retries + of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under + it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. + Can only be set when the task is created. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -54,6 +58,7 @@ class TaskRequest: name: str parent: None | str | Unset = UNSET + retry_of: None | str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -79,6 +84,12 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + retry_of: None | str | Unset + if isinstance(self.retry_of, Unset): + retry_of = UNSET + else: + retry_of = self.retry_of + queue = self.queue external_id = self.external_id @@ -144,6 +155,8 @@ def to_dict(self) -> dict[str, Any]: ) if parent is not UNSET: field_dict["parent"] = parent + if retry_of is not UNSET: + field_dict["retry_of"] = retry_of if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -187,6 +200,15 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + def _parse_retry_of(data: object) -> None | str | Unset: + if data is None: + return data + if isinstance(data, Unset): + return data + return cast(None | str | Unset, data) + + retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) + queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -282,6 +304,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: task_request = cls( name=name, parent=parent, + retry_of=retry_of, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index 4564580..a9b651b 100644 --- a/taskbadger/sdk.py +++ b/taskbadger/sdk.py @@ -184,6 +184,7 @@ def create_task( queue: str = None, external_id: str = None, parent: str = None, + retry_of: str = None, ) -> "Task": """Create a Task. @@ -203,6 +204,9 @@ def create_task( parent: ID of the parent task. Tasks nest a single level deep, so this must be the ID of a task that is not itself a child. The Celery and Procrastinate integrations set this automatically for tasks enqueued from within a tracked task. + retry_of: ID of the first attempt, to mark the task as a retry of it. The first 3 retries + of a task don't count towards your monthly task allocation. The Celery integration + sets this automatically for retries. Returns: Task: The created Task object. @@ -213,6 +217,8 @@ def create_task( } if parent is not None: task_dict["parent"] = parent + if retry_of is not None: + task_dict["retry_of"] = retry_of if queue is not None: task_dict["queue"] = queue if external_id is not None: @@ -392,6 +398,7 @@ def create( queue: str = None, external_id: str = None, parent: str = None, + retry_of: str = None, ) -> "Task": """Create a new task @@ -411,6 +418,7 @@ def create( queue=queue, external_id=external_id, parent=parent, + retry_of=retry_of, ) def __init__(self, task): diff --git a/tests/test_celery.py b/tests/test_celery.py index 56ca7bc..23b09f7 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -466,12 +466,18 @@ def add_retry(self, a, b, is_retry=False): mock.patch("taskbadger.sdk.get_task") as get_task, ): create.return_value = task_for_test() + update.return_value = task_for_test() get_task.return_value = task_for_test() result = add_retry.delay(2, 2) assert result.get(timeout=10, propagate=True) == 4 _wait_for_call_count(update, 4) assert create.call_count == 2 + first, retry = create.call_args_list + # the original attempt is an ordinary task; the retry is marked and nested under it + assert "retry_of" not in first.kwargs + assert retry.kwargs["retry_of"] is not None + assert retry.kwargs["parent"] == retry.kwargs["retry_of"] # a top-level task's retry nests under it assert update.call_args_list == [ mock.call(mock.ANY, status=StatusEnum.PROCESSING, data=mock.ANY), mock.call(mock.ANY, status=StatusEnum.ERROR, data=mock.ANY), diff --git a/tests/test_parents.py b/tests/test_parents.py index f648278..8d67a30 100644 --- a/tests/test_parents.py +++ b/tests/test_parents.py @@ -430,6 +430,68 @@ def test_celery_eager_explicit_none_parent_makes_a_root_task(): assert create.call_args.kwargs["parent"] is None +@pytest.mark.usefixtures("_bind_settings") +def test_celery_publish_of_a_retry_of_a_fan_out_child_sits_beside_it(): + """A retry is published with the running task's own id and a `retries` header. For a + child of another task it points at the child and stays under the same parent.""" + add = _celery_app() + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.sdk.get_task"), + mock.patch("taskbadger.celery._is_retry_of_running_task", return_value=True), + ): + create.return_value = task_for_test() + token = enter_task("child", "batch") + try: + add.apply_async((2, 2), retries=1) + finally: + exit_task(token) + + assert create.call_args.kwargs["retry_of"] == "child" + assert create.call_args.kwargs["parent"] == "batch" + + +@pytest.mark.usefixtures("_bind_settings") +def test_celery_publish_without_retries_header_is_not_a_retry(): + add = _celery_app() + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.sdk.get_task"), + mock.patch("taskbadger.celery._is_retry_of_running_task", return_value=True), + ): + create.return_value = task_for_test() + token = enter_task("root_id") + try: + add.apply_async((2, 2)) + finally: + exit_task(token) + + assert "retry_of" not in create.call_args.kwargs + assert create.call_args.kwargs["parent"] == "root_id" + + +@pytest.mark.usefixtures("_bind_settings") +def test_celery_eager_retry_points_at_the_first_attempt(): + add = _celery_app(task_always_eager=True, task_eager_propagates=True) + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.celery.update_task_safe"), + mock.patch("taskbadger.sdk.get_task"), + ): + create.return_value = task_for_test() + token = enter_task("first") + try: + add.apply_async((2, 2), retries=1) + finally: + exit_task(token) + + assert create.call_args.kwargs["retry_of"] == "first" + assert create.call_args.kwargs["parent"] == "first" + + # --- Procrastinate ------------------------------------------------------------ @@ -495,3 +557,25 @@ def records_parent(): assert seen["parent"] == "job_tb" assert current_task_id() is None + + +def test_retry_of_id(): + from taskbadger._current_task import enter_task, exit_task, parent_id, retry_of_id + + assert retry_of_id() is None # outside a tracked task + + token = enter_task("first") # a top-level task: its retries point at it and nest under it + assert (retry_of_id(), parent_id()) == ("first", "first") + exit_task(token) + + token = enter_task("retry", "first", retry_of="first") # retries of retries point at the first attempt + assert (retry_of_id(), parent_id()) == ("first", "first") + exit_task(token) + + token = enter_task("child", "batch") # a fan-out child's retries sit beside it under the batch + assert (retry_of_id(), parent_id()) == ("child", "batch") + exit_task(token) + + token = enter_task("child-retry", "batch", retry_of="child") + assert (retry_of_id(), parent_id()) == ("child", "batch") + exit_task(token)