Merge pull request #58 from taskbadger/sk/keep-alive

Keep long-running tasks from going stale

Author: snopoke

README.md

@@ -98,3 +98,30 @@ taskbadger.init(
 - **`task.configure(...).defer(...)` is not tracked.** Procrastinate's `configure()` returns a separate `JobDeferrer` whose methods bypass our wrapper. Use `task.defer(...)` directly for tracked deferrals. Tasks deferred via `configure().defer()` will run normally but will not appear in TaskBadger.
 - **`task.batch_defer*` is not tracked.** Same reason as `configure().defer()`.
 - **Tasks added via `app.add_tasks_from(blueprint)` after `ProcrastinateSystemIntegration` is constructed are not auto-instrumented.** Construct the integration after all blueprints are registered, or apply `@track` to those tasks explicitly.
+
+### Keeping long-running tasks fresh
+
+A task with a `stale_timeout` is marked `stale` by Task Badger if it goes too long without an
+update. Set `heartbeat_interval` (seconds) to have the SDK ping the task for you while it runs,
+rather than updating it from the task body.
+
+For Procrastinate, on the task or on `ProcrastinateSystemIntegration(...)`:
+
+```python
+@track(heartbeat_interval=60)
+@app.task
+async def slow_job():
+    ...
+```
+
+For Celery, on `CelerySystemIntegration(...)`, on the task, or per call with
+`slow_job.apply_async(taskbadger_heartbeat_interval=60)`:
+
+```python
+@app.task(base=taskbadger.Task, taskbadger_heartbeat_interval=60)
+def slow_job():
+    ...
+```
+
+Unless `stale_timeout` is given explicitly it is set to twice the interval. All running tasks are
+pinged from a single background thread, started the first time a task with a heartbeat runs.

integration_tests/tasks.py

@@ -1,7 +1,13 @@
+import time
+
 from celery import shared_task
 
 import taskbadger.celery
 
+HEARTBEAT_INTERVAL = 1
+# long enough to sample the task's `updated` time while it is still running
+SLOW_ADD_DURATION = 6
+
 
 @shared_task(bind=True, base=taskbadger.celery.Task)
 def add(self, x, y):
@@ -14,3 +20,10 @@ def add(self, x, y):
 def add_auto_track(self, x, y):
     assert self.request.taskbadger_task_id is not None, "missing task ID on self.request"
     return x + y
+
+
+@shared_task(bind=True, base=taskbadger.celery.Task, taskbadger_heartbeat_interval=HEARTBEAT_INTERVAL)
+def slow_add(self, x, y):
+    """Runs long enough to go stale without a heartbeat, and never updates itself."""
+    time.sleep(SLOW_ADD_DURATION)
+    return x + y

integration_tests/test_celery.py

@@ -1,11 +1,13 @@
 import logging
 import random
+import time
 
 import pytest
 
+import taskbadger
 from taskbadger import StatusEnum
 
-from .tasks import add, add_auto_track
+from .tasks import HEARTBEAT_INTERVAL, add, add_auto_track, slow_add
 
 
 @pytest.fixture(autouse=True)
@@ -41,3 +43,30 @@ def test_celery_auto_track(celery_session_app, celery_session_worker):
     a, b = random.randint(1, 1000), random.randint(1, 1000)
     result = add_auto_track.delay(a, b)
     assert result.get(timeout=10, propagate=True) == a + b
+
+
+def test_celery_heartbeat(celery_session_app, celery_session_worker):
+    """The worker pings the task while it runs, so it doesn't go stale."""
+    a, b = random.randint(1, 1000), random.randint(1, 1000)
+    result = slow_add.delay(a, b)
+
+    running = _wait_for_status(result.taskbadger_task_id, StatusEnum.PROCESSING)
+    assert running.stale_timeout == HEARTBEAT_INTERVAL * 2
+
+    time.sleep(HEARTBEAT_INTERVAL * 2)
+    pinged = taskbadger.get_task(running.id)
+
+    assert result.get(timeout=30, propagate=True) == a + b
+    # still running, so the task can only have been touched by the heartbeat
+    assert pinged.status == StatusEnum.PROCESSING
+    assert pinged.updated > running.updated, "task was not pinged while it was running"
+
+
+def _wait_for_status(task_id, status, timeout=15):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        task = taskbadger.get_task(task_id)
+        if task.status == status:
+            return task
+        time.sleep(0.2)
+    pytest.fail(f"task '{task_id}' did not reach status '{status}'")

integration_tests/test_procrastinate.py

@@ -8,9 +8,11 @@
 pyproject.toml.
 """
 
+import datetime
 import logging
 import os
 import random
+import time
 
 import procrastinate
 import psycopg
@@ -26,6 +28,10 @@
     "postgresql://postgres:postgres@localhost:5432/procrastinate",
 )
 
+HEARTBEAT_INTERVAL = 1
+# long enough for the heartbeat to fire while the task is still running
+SLOW_TASK_DURATION = 3
+
 
 @pytest.fixture(autouse=True)
 def _check_log_errors(caplog):
@@ -102,6 +108,38 @@ def add_manual(a, b):
     assert fetched.data == {"result": a + b}
 
 
+def test_heartbeat(app):
+    """The worker pings the task while it runs, so it doesn't go stale."""
+
+    @track(heartbeat_interval=HEARTBEAT_INTERVAL)
+    @app.task(name="slow", queue="taskbadger_int_hb")
+    def slow():
+        # `run_worker` blocks the test, so sample the task's `updated` time from
+        # inside the body. Fetched directly to bypass the integration's cache.
+        tb_id = current_task().id
+        before = taskbadger.get_task(tb_id).updated
+        time.sleep(SLOW_TASK_DURATION)
+        after = taskbadger.get_task(tb_id).updated
+        current_task().update(data={"before": before.isoformat(), "after": after.isoformat()})
+
+    job_id = slow.defer()
+    app.run_worker(
+        queues=["taskbadger_int_hb"],
+        wait=False,
+        install_signal_handlers=False,
+        listen_notify=False,
+    )
+
+    args = _fetch_job_args(job_id)
+    fetched = taskbadger.get_task(args["__taskbadger_task_id__"])
+
+    assert fetched.status == StatusEnum.SUCCESS
+    assert fetched.stale_timeout == HEARTBEAT_INTERVAL * 2
+    before = datetime.datetime.fromisoformat(fetched.data["before"])
+    after = datetime.datetime.fromisoformat(fetched.data["after"])
+    assert after > before, "task was not pinged while it was running"
+
+
 def test_auto_track_via_system(app):
     ProcrastinateSystemIntegration(app=app, auto_track_tasks=True)
 

taskbadger/_heartbeat.py

@@ -0,0 +1,127 @@
+"""Background heartbeat support for the system integrations (Celery,
+Procrastinate). Not part of the public API.
+
+A task that sets ``stale_timeout`` is marked ``stale`` by the API if it goes
+too long without an update. Long-running queue tasks that don't report progress
+would trip that timeout while perfectly healthy, so the integrations can ping
+the task periodically for the duration of the run.
+
+All in-flight tasks are pinged from a single daemon thread rather than one
+thread per task. The thread is started lazily on first use and sleeps
+indefinitely while nothing is registered.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import threading
+import time
+
+from ._integrations import TERMINAL_STATES, is_valid_interval
+from .mug import Badger, Settings
+from .safe_sdk import update_task_safe
+
+log = logging.getLogger("taskbadger")
+
+
+@dataclasses.dataclass
+class _Entry:
+    interval: float
+    settings: Settings
+    due: float
+
+
+class Heartbeat:
+    """Periodically pings registered tasks to keep them from going stale."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._entries: dict[str, _Entry] = {}
+        self._wake = threading.Event()
+        self._thread = None
+
+    def start(self, task_id: str, interval: float | None) -> None:
+        """Begin pinging ``task_id`` every ``interval`` seconds.
+
+        No-op if there is no interval or Task Badger isn't configured in the
+        calling thread. Registering the same task again resets its schedule.
+        """
+        if not task_id or not is_valid_interval(interval):
+            return
+
+        settings = Badger.current.settings
+        if settings is None:
+            return
+
+        with self._lock:
+            self._entries[task_id] = _Entry(interval, settings, time.monotonic() + interval)
+            self._ensure_thread()
+        self._wake.set()
+
+    def stop(self, task_id: str) -> None:
+        """Stop pinging ``task_id``. No-op if it isn't registered."""
+        with self._lock:
+            self._entries.pop(task_id, None)
+
+    def stop_all(self) -> None:
+        with self._lock:
+            self._entries.clear()
+
+    def _ensure_thread(self) -> None:
+        """Start the beat thread if it isn't running. Called with the lock held."""
+        if self._thread is not None and self._thread.is_alive():
+            return
+
+        # `is_alive()` is also False for a thread inherited from a parent
+        # process, so a forked worker (e.g. Celery's prefork pool) starts its
+        # own thread the first time it runs a tracked task.
+        self._thread = threading.Thread(target=self._run, name="taskbadger-heartbeat", daemon=True)
+        self._thread.start()
+
+    def _run(self) -> None:
+        while True:
+            try:
+                timeout = self._beat()
+            except Exception:
+                # Never let the thread die: every registered task would then go
+                # stale with nothing to restart the pings.
+                log.exception("heartbeat beat failed")
+                timeout = 1.0
+            self._wake.wait(timeout)
+            self._wake.clear()
+
+    def _beat(self) -> float | None:
+        """Ping every task that is due and return the seconds to sleep for.
+
+        ``None`` means "sleep until a task is registered".
+        """
+        now = time.monotonic()
+        with self._lock:
+            due = [(task_id, entry) for task_id, entry in self._entries.items() if entry.due <= now]
+            for _, entry in due:
+                entry.due = now + entry.interval
+
+        for task_id, entry in due:
+            self._ping(task_id, entry)
+
+        with self._lock:
+            if not self._entries:
+                return None
+            next_due = min(entry.due for entry in self._entries.values())
+        return max(next_due - time.monotonic(), 0)
+
+    def _ping(self, task_id: str, entry: _Entry) -> None:
+        # The beat thread has its own context, so bind the settings captured
+        # when the task registered rather than relying on inheritance.
+        if Badger.current.settings is not entry.settings:
+            Badger.current.bind(entry.settings)
+
+        log.debug("heartbeat ping '%s'", task_id)
+        task = update_task_safe(task_id)
+        if task is not None and task.status in TERMINAL_STATES:
+            # The task finished (or was cancelled) elsewhere.
+            self.stop(task_id)
+
+
+heartbeat = Heartbeat()

taskbadger/_integrations.py

@@ -12,9 +12,11 @@
 
 import collections
 import logging
+import math
 import re
 
 from . import sdk
+from .exceptions import ConfigurationError
 from .internal.models import StatusEnum
 from .systems import System
 
@@ -27,6 +29,10 @@
     StatusEnum.STALE,
 }
 
+# When a heartbeat is configured but no `stale_timeout` is given, the timeout
+# is derived from the interval using this factor (as the CLI's `run` does).
+STALE_TIMEOUT_FACTOR = 2
+
 
 class TaskCache:
     """Bounded LRU-ish cache for TaskBadger Task objects.
@@ -91,18 +97,71 @@ def match_task_name(task_name: str, includes, excludes) -> bool:
     return True
 
 
+def is_valid_interval(interval) -> bool:
+    """Return True if ``interval`` is usable as a heartbeat interval."""
+    return isinstance(interval, int | float) and interval > 0
+
+
+def validate_interval(interval) -> None:
+    """Raise if ``interval`` is set but isn't usable as a heartbeat interval."""
+    if interval is not None and not is_valid_interval(interval):
+        raise ConfigurationError(f"heartbeat_interval must be a positive number of seconds: {interval!r}")
+
+
+def resolve_heartbeat_options(heartbeat_interval, stale_timeout, system):
+    """Resolve the heartbeat interval and stale timeout for a single task.
+
+    Values set on the task win over those set on the system integration. If a
+    heartbeat is configured without a stale timeout, one is derived from the
+    interval.
+
+    Returns:
+        A tuple of ``(heartbeat_interval, stale_timeout)``, either of which may
+        be ``None``.
+    """
+    if system is not None:
+        if heartbeat_interval is None:
+            heartbeat_interval = system.heartbeat_interval
+        if stale_timeout is None:
+            stale_timeout = system.stale_timeout
+
+    if heartbeat_interval is not None and not is_valid_interval(heartbeat_interval):
+        # Per-task values don't go through `validate_interval`, and a bad one
+        # shouldn't stop the task from being tracked.
+        log.warning("Ignoring invalid heartbeat_interval: %r", heartbeat_interval)
+        heartbeat_interval = None
+
+    if stale_timeout is None and heartbeat_interval:
+        # `stale_timeout` is whole seconds, so round up: a sub-second interval
+        # must not produce a timeout of 0.
+        stale_timeout = max(1, math.ceil(heartbeat_interval * STALE_TIMEOUT_FACTOR))
+
+    return heartbeat_interval, stale_timeout
+
+
 class BaseSystemIntegration(System):
     """Common ctor + ``track_task`` body for system integrations.
 
     Subclasses set ``identifier`` and may override ``track_task`` to add
     additional filtering (e.g. skipping built-in tasks).
     """
 
-    def __init__(self, auto_track_tasks=True, includes=None, excludes=None, record_task_args=False):
+    def __init__(
+        self,
+        auto_track_tasks=True,
+        includes=None,
+        excludes=None,
+        record_task_args=False,
+        heartbeat_interval=None,
+        stale_timeout=None,
+    ):
+        validate_interval(heartbeat_interval)
         self.auto_track_tasks = auto_track_tasks
         self.includes = includes
         self.excludes = excludes
         self.record_task_args = record_task_args
+        self.heartbeat_interval = heartbeat_interval
+        self.stale_timeout = stale_timeout
 
     def track_task(self, task_name: str) -> bool:
         if not self.auto_track_tasks: