feat(celery): mark retries with retry_of

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

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

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.
 

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})
 
 

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,

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,

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,