feat(taskworker): Send heartbeats for running tasks

.github/workflows/ci.yml

@@ -419,6 +419,44 @@ jobs:
         run: |
           make test-upkeep-expiry
 
+  heartbeat-integration-test:
+    name: Heartbeat integration test
+    runs-on: ubuntu-24.04
+
+    steps:
+      - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2
+
+      - name: Install cmake
+        uses: lukka/get-cmake@28983e0d3955dba2bb0a6810caae0c6cf268ec0c # latest
+
+      - name: Install libcurl
+        run: sudo apt-get update && sudo apt-get install -y libcurl4-openssl-dev
+
+      - uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # pin@v1
+        with:
+          toolchain: stable
+          profile: minimal
+          override: true
+
+      - uses: swatinem/rust-cache@81d053bdb0871dcd3f10763c8cc60d0adc41762b # pin@v1
+        with:
+          key: ${{ github.job }}
+
+      - uses: astral-sh/setup-uv@5a7eac68fb9809dea845d802897dc5c723910fa3 # v7.1.3
+        with:
+          version: '0.13.7'
+          # we just cache the venv-dir directly in action-setup-venv
+          enable-cache: false
+
+      - uses: getsentry/action-setup-venv@5a80476d175edf56cb205b08bc58986fa99d1725 # v3.2.0
+        with:
+          cache-dependency-path: uv.lock
+          install-cmd: uv sync --frozen --only-dev --active
+
+      - name: Run heartbeat integration test
+        run: |
+          make test-heartbeat
+
   upkeep-delay-integration-test:
     name: Upkeep delay integration test
     runs-on: ubuntu-24.04

Makefile

@@ -89,12 +89,17 @@ test-failed-tasks: build reset-kafka ## Run the failed tasks integration test
 	rm -r integration_tests/.tests_output/test_failed_tasks
 .PHONY: test-failed-tasks
 
+test-heartbeat: build reset-kafka ## Run the task heartbeat integration test
+	python -m pytest integration_tests/integration_tests/test_heartbeat.py -s
+	rm -r integration_tests/.tests_output/test_heartbeat
+.PHONY: test-heartbeat
+
 test-push-postgres: build reset-kafka ## Run the postgres-backed push mode integration test
 	python -m pytest integration_tests/integration_tests/test_push_postgres.py -s
 	rm -r integration_tests/.tests_output/test_push_postgres
 .PHONY: test-push-postgres
 
-integration-test: test-rebalance test-worker-processing test-upkeep-retry test-upkeep-expiry test-upkeep-delay test-failed-tasks ## Run all integration tests
+integration-test: test-rebalance test-worker-processing test-upkeep-retry test-upkeep-expiry test-upkeep-delay test-failed-tasks test-heartbeat ## Run all integration tests
 .PHONY: integration-test
 
 

clients/python/pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 dependencies = [
     "sentry-arroyo>=2.41.1",
     "sentry-sdk[http2]>=2.73.0",
-    "sentry-protos>=0.26.1",
+    "sentry-protos>=0.106.0",
     "confluent_kafka>=2.3.0",
     "cronsim>=2.6",
     "grpcio>=1.67.1",

clients/python/src/taskbroker_client/constants.py

@@ -11,6 +11,11 @@
 will use if neither the TaskNamespace or Task define a deadline
 """
 
+MIN_HEARTBEAT = 10
+"""
+The smallest heartbeat (in seconds) a task can set.
+"""
+
 DEFAULT_REBALANCE_AFTER = 32
 """
 The number of tasks a worker will process before it

clients/python/src/taskbroker_client/registry.py

@@ -93,6 +93,7 @@ def _validate_raw_mode(
         processing_deadline_duration: int | datetime.timedelta | None,
         at_most_once: bool,
         compression_type: CompressionType,
+        heartbeat: int | datetime.timedelta | None,
     ) -> None:
         """
         Reject task options that taskbroker ignores in raw mode.
@@ -137,6 +138,11 @@ def _validate_raw_mode(
                 "taskbroker hardcodes at_most_once=false for raw activations, so the task "
                 "would still be retried past its processing deadline."
             )
+        if heartbeat is not None:
+            raise ValueError(
+                f"Task {name!r} in raw-mode namespace {self.name!r} sets heartbeat. "
+                "taskbroker never sets a heartbeat on raw activations, so it would be ignored."
+            )
         if retry is not None and retry._times_exceeded == LastAction.Deadletter:
             raise ValueError(
                 f"Task {name!r} in raw-mode namespace {self.name!r} sets "
@@ -157,6 +163,7 @@ def register(
         report_timeout_errors: bool = True,
         silenced_exceptions: tuple[type[BaseException], ...] | None = None,
         pass_headers: bool = False,
+        heartbeat: int | datetime.timedelta | None = None,
     ) -> Callable[[Callable[P, R]], Task[P, R]]:
         """
         Register a task.
@@ -191,6 +198,10 @@ def register(
         pass_headers: bool
             If True, the task function will receive task activation headers
             as a keyword argument named `headers` (dict[str, str]).
+        heartbeat: int | datetime.timedelta | None
+            If set, the worker renews this task's heartbeat while it runs, and
+            the broker retries it once a heartbeat is missed for this many seconds.
+            Must be at least 10s and less than processing_deadline_duration.
         """
 
         def wrapped(func: Callable[P, R]) -> Task[P, R]:
@@ -205,6 +216,7 @@ def wrapped(func: Callable[P, R]) -> Task[P, R]:
                     processing_deadline_duration=processing_deadline_duration,
                     at_most_once=at_most_once,
                     compression_type=compression_type,
+                    heartbeat=heartbeat,
                 )
             task = Task(
                 name=name,
@@ -221,6 +233,7 @@ def wrapped(func: Callable[P, R]) -> Task[P, R]:
                 report_timeout_errors=report_timeout_errors,
                 silenced_exceptions=silenced_exceptions,
                 pass_headers=pass_headers,
+                heartbeat=heartbeat,
             )
             # TODO(taskworker) tasks should be registered into the registry
             # so that we can ensure task names are globally unique
@@ -320,6 +333,7 @@ def register(
         report_timeout_errors: bool = True,
         silenced_exceptions: tuple[type[BaseException], ...] | None = None,
         pass_headers: bool = False,
+        heartbeat: int | datetime.timedelta | None = None,
     ) -> Callable[[Callable[P, R]], ExternalTask[P, R]]:
         """
         Register an external task stub.
@@ -372,6 +386,7 @@ def wrapped(func: Callable[P, R]) -> ExternalTask[P, R]:
                 report_timeout_errors=report_timeout_errors,
                 silenced_exceptions=silenced_exceptions,
                 pass_headers=pass_headers,
+                heartbeat=heartbeat,
             )
             self._registered_tasks[name] = task
             return task

clients/python/src/taskbroker_client/task.py

@@ -22,6 +22,7 @@
 from taskbroker_client.constants import (
     DEFAULT_PROCESSING_DEADLINE,
     MAX_PARAMETER_BYTES_BEFORE_COMPRESSION,
+    MIN_HEARTBEAT,
     CompressionType,
 )
 from taskbroker_client.retry import Retry
@@ -106,6 +107,7 @@ def __init__(
         report_timeout_errors: bool = True,
         silenced_exceptions: tuple[type[BaseException], ...] | None = None,
         pass_headers: bool = False,
+        heartbeat: int | datetime.timedelta | None = None,
     ):
         self.name = name
         self._func = func
@@ -120,6 +122,22 @@ def __init__(
                 Having retries enabled means that a task supports being executed
                 multiple times and thus cannot be idempotent.
                 """)
+        if isinstance(heartbeat, datetime.timedelta):
+            heartbeat = int(heartbeat.total_seconds())
+        if heartbeat is not None:
+            deadline = self._processing_deadline_duration
+            if isinstance(deadline, datetime.timedelta):
+                deadline = int(deadline.total_seconds())
+            if heartbeat < MIN_HEARTBEAT:
+                raise ValueError(
+                    f"Task {name!r} heartbeat must be at least {MIN_HEARTBEAT}s, got {heartbeat}s"
+                )
+            if heartbeat >= deadline:
+                raise ValueError(
+                    f"Task {name!r} heartbeat ({heartbeat}s) must be less than its "
+                    f"processing_deadline_duration ({deadline}s)"
+                )
+        self.heartbeat = heartbeat
         self._retry = retry
         self.at_most_once = at_most_once
         self.wait_for_delivery = wait_for_delivery
@@ -410,6 +428,7 @@ def create_activation(
             processing_deadline_duration=processing_deadline,
             expires=expires,
             delay=countdown,
+            heartbeat=self.heartbeat,
         )
 
     def create_testing_activation(

clients/python/src/taskbroker_client/worker/client.py

@@ -13,6 +13,7 @@
 import grpc
 from google.protobuf.message import Message
 from sentry_protos.taskbroker.v1.taskbroker_pb2 import (
+    BatchUpdateHeartbeatRequest,
     FetchNextTask,
     GetTaskRequest,
     SetTaskStatusRequest,
@@ -39,6 +40,9 @@
 _AUTH_FAILED = object()
 
 MAX_ACTIVATION_SIZE = 1024 * 1024 * 10
+
+# Well below the 10s minimum heartbeat, so a hung call can't eat a whole interval.
+HEARTBEAT_RPC_TIMEOUT_SEC = 5
 """Max payload size we will process."""
 
 GRPC_HEALTH_CHECK_METHOD = "/grpc.health.v1.Health/Check"
@@ -522,3 +526,23 @@ def update_task(
                 receive_timestamp=time.monotonic(),
             )
         return None
+
+    def update_heartbeats(self, host: str, ids: list[str]) -> list[str]:
+        """
+        Renew the heartbeat of activations processing on `host`. Returns the ids
+        that host no longer has processing.
+        """
+        tags = {"processing_pool": self._processing_pool_name}
+        try:
+            response = self._get_stub(host).BatchUpdateHeartbeat(
+                BatchUpdateHeartbeatRequest(ids=ids), timeout=HEARTBEAT_RPC_TIMEOUT_SEC
+            )
+        except grpc.RpcError as err:
+            self._metrics.incr(
+                "taskworker.client.heartbeat.rpc", tags={"status": err.code().name, **tags}
+            )
+            raise
+
+        self._metrics.incr("taskworker.client.heartbeat.rpc", tags={"status": "OK", **tags})
+        self._metrics.distribution("taskworker.client.heartbeat.batch_size", len(ids), tags=tags)
+        return list(response.not_found_ids)