chore: implement task heartbeat (#909)

* chore: renew Redis stream message ownership during worker execution

Add a worker heartbeat so long-running tasks are not reclaimed by
XAUTOCLAIM while they are still running. ReceivedMessage now carries
both ack() and renew() callbacks, with renew implemented via XCLAIM
to reset the pending idle timer for the current consumer. The worker
starts this heartbeat for broker-delivered messages and stops it when
execution finishes. Interactive execution is unchanged.

* test: run_local.sh has parallel workers

Author: chaen

diracx-tasks/src/diracx/tasks/plumbing/broker/models.py

@@ -109,6 +109,7 @@ class ReceivedMessage(BaseModel):
 
     data: bytes
     ack: Callable[[], Awaitable[None]]
+    renew: Callable[[], Awaitable[None]]
 
 
 def _prepare_arg(arg: Any) -> Any:

diracx-tasks/src/diracx/tasks/plumbing/broker/redis_streams.py

@@ -138,6 +138,28 @@ async def _ack() -> None:
 
         return _ack
 
+    def _renew_generator(
+        self, msg_id: str | bytes, queue_name: str | bytes
+    ) -> Callable[[], Awaitable[None]]:
+        """Return a coroutine that resets the PEL idle timer for a message.
+
+        Calls XCLAIM with min-idle-time=0, which always succeeds and resets
+        the idle clock — preventing the autoclaim loop from reclaiming a
+        message that is still being actively processed.
+        """
+
+        async def _renew() -> None:
+            async with Redis(connection_pool=self.connection_pool) as redis:
+                await redis.xclaim(
+                    queue_name,
+                    self.consumer_group_name,
+                    self.consumer_name,
+                    min_idle_time=0,
+                    message_ids=[msg_id],
+                )
+
+        return _renew
+
     async def listen(self) -> AsyncGenerator[ReceivedMessage, None]:
         """Yield messages from streams in strict priority order.
 
@@ -163,6 +185,9 @@ async def listen(self) -> AsyncGenerator[ReceivedMessage, None]:
                         yield ReceivedMessage(
                             data=msg[b"data"],
                             ack=self._ack_generator(msg_id=msg_id, queue_name=stream),
+                            renew=self._renew_generator(
+                                msg_id=msg_id, queue_name=stream
+                            ),
                         )
 
                 # Reclaim unacknowledged messages (throttled to idle_timeout interval)
@@ -187,10 +212,11 @@ async def listen(self) -> AsyncGenerator[ReceivedMessage, None]:
                         )
 
                         if pending[1]:
-                            logger.debug(
-                                "Reclaimed %d unacked messages from %s",
+                            logger.info(
+                                "Reclaimed %d unacked messages from %s (message ids: %s)",
                                 len(pending[1]),
                                 sname,
+                                [msg_id for msg_id, msg in pending[1]],
                             )
 
                         for msg_id, msg in pending[1]:
@@ -199,4 +225,7 @@ async def listen(self) -> AsyncGenerator[ReceivedMessage, None]:
                                 ack=self._ack_generator(
                                     msg_id=msg_id, queue_name=sname
                                 ),
+                                renew=self._renew_generator(
+                                    msg_id=msg_id, queue_name=sname
+                                ),
                             )

diracx-tasks/src/diracx/tasks/plumbing/worker/worker.py

@@ -6,7 +6,7 @@
 import logging
 from datetime import UTC, datetime, timedelta
 from time import time
-from typing import Any, Callable
+from typing import Any, Awaitable, Callable
 
 import msgpack
 from opentelemetry import metrics, trace
@@ -47,6 +47,28 @@
 _LOCK_RETRY_DELAY_SECONDS = 5
 
 
+async def _message_heartbeat(
+    renew: Callable[[], Awaitable[None]],
+    stop_event: asyncio.Event,
+    interval: float,
+) -> None:
+    """Periodically reset the PEL idle timer while a message is being processed.
+
+    Calls ``renew()`` at ``interval`` seconds so the autoclaim loop never
+    considers the message stale as long as execution is in progress.
+    """
+    while not stop_event.is_set():
+        try:
+            await asyncio.wait_for(stop_event.wait(), timeout=interval)
+            return
+        except asyncio.TimeoutError:
+            pass
+        try:
+            await renew()
+        except Exception:
+            logger.warning("Failed to renew message ownership", exc_info=True)
+
+
 class Worker:
     """Execute tasks consumed from a broker.
 
@@ -261,22 +283,42 @@ async def process_message(self, message: bytes | ReceivedMessage) -> None:
             "Executing task %s (ID: %s)", task_message.task_name, task_message.task_id
         )
 
-        result = await self.run_task(task_func, task_message)
-
-        # Handle failure: retry or dead letter queue
-        if result.is_err:
-            await self._handle_failure(task_message, result)
-        else:
-            await self._handle_success(task_message, result)
+        # Start heartbeat to keep PEL ownership while the task runs.
+        # Renew at half the idle_timeout so there's always a safety margin.
+        heartbeat_stop = asyncio.Event()
+        heartbeat_interval = self.broker.idle_timeout / 1000 / 2
+        heartbeat_task: asyncio.Task[None] | None = None
+        if isinstance(message, ReceivedMessage):
+            heartbeat_task = asyncio.create_task(
+                _message_heartbeat(message.renew, heartbeat_stop, heartbeat_interval)
+            )
 
-        # Always persist the result to the backend
         try:
-            if self.broker.result_backend:
-                await self.broker.result_backend.set_result(
-                    task_message.task_id, result
-                )
-        except Exception:
-            logger.exception("Failed to save result")
+            result = await self.run_task(task_func, task_message)
+
+            # Handle failure: retry or dead letter queue
+            if result.is_err:
+                await self._handle_failure(task_message, result)
+            else:
+                await self._handle_success(task_message, result)
+
+            # Always persist the result to the backend
+            try:
+                if self.broker.result_backend:
+                    await self.broker.result_backend.set_result(
+                        task_message.task_id, result
+                    )
+            except Exception:
+                logger.exception("Failed to save result")
+
+        finally:
+            if heartbeat_task is not None:
+                heartbeat_stop.set()
+                heartbeat_task.cancel()
+                try:
+                    await heartbeat_task
+                except asyncio.CancelledError:
+                    pass
 
         if isinstance(message, ReceivedMessage):
             await message.ack()

diracx-tasks/tests/test_redis_streams.py

@@ -2,8 +2,11 @@
 
 from __future__ import annotations
 
+from unittest.mock import AsyncMock, patch
+
 from diracx.tasks.plumbing.broker.redis_streams import (
     ALL_STREAM_NAMES,
+    RedisStreamBroker,
     stream_name_for,
 )
 from diracx.tasks.plumbing.enums import Priority, Size
@@ -25,3 +28,29 @@ def test_all_stream_names():
     assert "diracx:tasks:realtime:small" in ALL_STREAM_NAMES
     assert "diracx:tasks:normal:medium" in ALL_STREAM_NAMES
     assert "diracx:tasks:background:large" in ALL_STREAM_NAMES
+
+
+async def test_renew_generator_calls_xclaim() -> None:
+    broker = RedisStreamBroker("redis://example.invalid")
+
+    redis_cm = AsyncMock()
+    redis = AsyncMock()
+    redis_cm.__aenter__.return_value = redis
+    redis_cm.__aexit__.return_value = False
+
+    with patch(
+        "diracx.tasks.plumbing.broker.redis_streams.Redis", return_value=redis_cm
+    ):
+        renew = broker._renew_generator(
+            msg_id="1234-0",
+            queue_name="diracx:tasks:normal:medium",
+        )
+        await renew()
+
+    redis.xclaim.assert_awaited_once_with(
+        "diracx:tasks:normal:medium",
+        broker.consumer_group_name,
+        broker.consumer_name,
+        min_idle_time=0,
+        message_ids=["1234-0"],
+    )

diracx-tasks/tests/test_worker_integration.py

@@ -2,10 +2,11 @@
 
 from __future__ import annotations
 
+import asyncio
 from unittest.mock import AsyncMock, patch
 
 from diracx.tasks.plumbing.base_task import BaseTask
-from diracx.tasks.plumbing.broker.models import TaskMessage, TaskResult
+from diracx.tasks.plumbing.broker.models import ReceivedMessage, TaskMessage, TaskResult
 from diracx.tasks.plumbing.depends import CallbackSpawner
 from diracx.tasks.plumbing.enums import Priority, Size
 from diracx.tasks.plumbing.factory import wrap_task
@@ -322,6 +323,40 @@ async def test_process_message_acks_on_parse_error(
     await worker.process_message(b"not valid msgpack at all!!")
 
 
+async def test_process_message_renews_ownership_while_running(
+    broker, task_class_registry, wrapped_registry
+):
+    """Worker should heartbeat message ownership and ack once done."""
+    worker = Worker(
+        broker=broker,
+        task_registry=wrapped_registry,
+        task_class_registry=task_class_registry,
+    )
+    worker.broker.idle_timeout = 20  # ms
+
+    task_msg = TaskMessage(
+        task_id="t-renew",
+        task_name="test:SuccessTask",
+        labels={"priority": "normal", "size": "small"},
+        task_args=[],
+        task_kwargs={},
+    )
+
+    ack = AsyncMock()
+    renew = AsyncMock()
+    message = ReceivedMessage(data=task_msg.dumpb(), ack=ack, renew=renew)
+
+    async def _slow_run_task(*_args, **_kwargs):
+        await asyncio.sleep(0.05)
+        return TaskResult.from_value("ok", execution_time=0.05)
+
+    with patch.object(worker, "run_task", side_effect=_slow_run_task):
+        await worker.process_message(message)
+
+    assert renew.await_count >= 1
+    ack.assert_awaited_once()
+
+
 # ---------------------------------------------------------------------------
 # TaskBroker dependency injection resolution
 # ---------------------------------------------------------------------------

docs/dev/explanations/tasks/class-details.md

@@ -148,7 +148,7 @@ classDiagram
     class RedisStreamBroker {
         9 streams: 3 priorities × 3 sizes
         +enqueue(TaskMessage)
-        +consume(size) ReceivedMessage
+        +listen() AsyncGenerator~ReceivedMessage~
         +startup()
     }
 
@@ -175,6 +175,7 @@ classDiagram
     class ReceivedMessage {
         +data: bytes
         +ack() awaitable
+        +renew() awaitable
     }
 
     class TaskBinding {
@@ -185,8 +186,8 @@ classDiagram
     }
 
     class RedisResultBackend {
-        +store(task_id, TaskResult)
-        +get(task_id) TaskResult
+        +set_result(task_id, TaskResult)
+        +get_result(task_id) TaskResult
     }
 
     RedisStreamBroker ..> TaskMessage : enqueues
@@ -195,7 +196,7 @@ classDiagram
     TaskBinding --> RedisStreamBroker : references
 ```
 
-`TaskMessage` is the wire-protocol message serialized to msgpack. `ReceivedMessage` wraps the raw bytes with an `ack()` callback so the worker can acknowledge the message after processing. `TaskBinding` maps a task class to its broker, providing the `submit()` method used by `BaseTask.schedule()`.
+`TaskMessage` is the wire-protocol message serialized to msgpack. `ReceivedMessage` wraps the raw bytes with `ack()` and `renew()` callbacks: `ack()` acknowledges completion, while `renew()` refreshes ownership of in-flight messages during long executions. `TaskBinding` maps a task class to its broker, providing the `submit()` method used by `BaseTask.schedule()`.
 
 ## Callback subsystem
 

docs/dev/explanations/tasks/index.md

@@ -28,7 +28,7 @@ Occasionally tasks must be scheduled to run at some point in the future, most co
 
 ## What is a task?
 
-Tasks are async Python functions which have extremely low overhead, allowing for many tasks to be spawned for even cheap operations.
+Tasks are Python classes derived from `BaseTask`. Each task class implements an async `execute(...)` method, and task instances can be scheduled with `task.schedule(...)`.
 Tasks can be executed in four different ways:
 
 - **Standalone tasks:** The task performs its work and then returns. For example synchronising the IAM to DiracX configuration queries IAM and then updates the DiracX CS.

docs/dev/reference/tasks.md

@@ -93,6 +93,20 @@ ConcurrencyLimiter(obj, key, *extra_keys, ttl_ms=30000)
 
 `PeriodicBaseTask` overrides this with a `MutexLock` keyed by the task class name, preventing concurrent execution. `PeriodicVoAwareBaseTask` adds the VO name to the lock key, so each VO gets its own mutex.
 
+### Lock watchdog vs message heartbeat
+
+Two independent watchdog-style mechanisms exist while a worker executes tasks:
+
+- **Lock watchdog** (`_lock_watchdog` in `factory.py`) periodically calls `extend()` on acquired locks, so lock TTLs do not expire mid-execution.
+- **Message heartbeat** (`_message_heartbeat` in `worker.py`) periodically renews stream-message ownership while a task is running.
+
+The message heartbeat uses Redis `XCLAIM` with `min_idle_time=0` to reset the pending-entry idle timer for the in-flight message. This prevents `XAUTOCLAIM` from reclaiming a long-running task simply because it exceeded the idle timeout.
+
+These mechanisms protect different things:
+
+- lock watchdog protects lock ownership
+- message heartbeat protects consumer-group message ownership
+
 ______________________________________________________________________
 
 ## Schedules

run_local.sh

@@ -174,9 +174,9 @@ uvicorn --factory diracx.testing.routers:create_app --reload > "${tmp_dir}/logs/
 diracx_pid=$!
 diracx-task-run scheduler > "${tmp_dir}/logs/scheduler.log" 2>&1 &
 scheduler_pid=$!
-diracx-task-run worker --worker-size small --max-concurrent-tasks 1 > "${tmp_dir}/logs/worker-sm.log" 2>&1 &
+diracx-task-run worker --worker-size small --max-concurrent-tasks 3 > "${tmp_dir}/logs/worker-sm.log" 2>&1 &
 worker_small_pid=$!
-diracx-task-run worker --worker-size medium --max-concurrent-tasks 1 > "${tmp_dir}/logs/worker-md.log" 2>&1 &
+diracx-task-run worker --worker-size medium --max-concurrent-tasks 2 > "${tmp_dir}/logs/worker-md.log" 2>&1 &
 worker_medium_pid=$!
 diracx-task-run worker --worker-size large --max-concurrent-tasks 1 > "${tmp_dir}/logs/worker-lg.log" 2>&1 &
 worker_large_pid=$!