From afcfb0a3fdd77ccadd37ac7cf0d5fa12aa10f083 Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Tue, 10 Nov 2026 14:37:12 -0500 Subject: [PATCH 1/2] feat(taskworker): Send heartbeats for running tasks Co-Authored-By: Claude Opus 5.5 --- clients/python/pyproject.toml | 2 +- .../python/src/taskbroker_client/constants.py | 5 + .../python/src/taskbroker_client/registry.py | 15 ++ clients/python/src/taskbroker_client/task.py | 19 ++ .../src/taskbroker_client/worker/heartbeat.py | 168 ++++++++++++++++++ .../taskbroker_client/worker/push_clients.py | 24 +++ .../src/taskbroker_client/worker/worker.py | 31 +++- clients/python/tests/test_registry.py | 53 ++++++ clients/python/tests/worker/test_heartbeat.py | 143 +++++++++++++++ uv.lock | 6 +- 10 files changed, 460 insertions(+), 6 deletions(-) create mode 100644 clients/python/src/taskbroker_client/worker/heartbeat.py create mode 100644 clients/python/tests/worker/test_heartbeat.py diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index 2e6ad250..c208138e 100644 --- a/clients/python/pyproject.toml +++ b/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", diff --git a/clients/python/src/taskbroker_client/constants.py b/clients/python/src/taskbroker_client/constants.py index dbbd2128..5de395dd 100644 --- a/clients/python/src/taskbroker_client/constants.py +++ b/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 diff --git a/clients/python/src/taskbroker_client/registry.py b/clients/python/src/taskbroker_client/registry.py index 9e0c7699..1aaaa246 100644 --- a/clients/python/src/taskbroker_client/registry.py +++ b/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 diff --git a/clients/python/src/taskbroker_client/task.py b/clients/python/src/taskbroker_client/task.py index c3429e54..e46f2909 100644 --- a/clients/python/src/taskbroker_client/task.py +++ b/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( diff --git a/clients/python/src/taskbroker_client/worker/heartbeat.py b/clients/python/src/taskbroker_client/worker/heartbeat.py new file mode 100644 index 00000000..98b6bab5 --- /dev/null +++ b/clients/python/src/taskbroker_client/worker/heartbeat.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import logging +import multiprocessing +import queue +import signal +import threading +from multiprocessing.context import ForkContext, ForkServerContext, SpawnContext +from multiprocessing.process import BaseProcess +from multiprocessing.synchronize import Event +from typing import Literal + +import grpc + +from taskbroker_client.app import import_app +from taskbroker_client.metrics import MetricsBackend +from taskbroker_client.worker.push_clients import PushTaskbrokerClient + +logger = logging.getLogger(__name__) + +HEARTBEAT_INTERVAL_SEC = 1.0 + +HeartbeatCommand = tuple[Literal["track", "untrack"], list[str]] + + +def heartbeat_process( + app_module: str, + broker_service: str, + processing_pool_name: str, + ids: list[str], + commands: queue.Queue[HeartbeatCommand], + shutdown_event: Event, +) -> None: + """Renew the heartbeat of every tracked activation once per interval.""" + # The parent owns shutdown and tells us through shutdown_event. + signal.signal(signal.SIGINT, signal.SIG_IGN) + signal.signal(signal.SIGTERM, signal.SIG_IGN) + + app = import_app(app_module) + metrics = app.metrics + client = PushTaskbrokerClient( + service=broker_service, + application=app.name, + metrics=metrics, + rpc_secret=app.config["rpc_secret"], + grpc_config=app.config["grpc_config"], + processing_pool_name=processing_pool_name, + ) + tags = {"processing_pool": processing_pool_name} + + tracked = set(ids) + while not shutdown_event.wait(HEARTBEAT_INTERVAL_SEC): + while True: + try: + action, command_ids = commands.get_nowait() + except queue.Empty: + break + if action == "track": + tracked.update(command_ids) + else: + tracked.difference_update(command_ids) + + metrics.gauge("taskworker.worker.heartbeat.tracked", len(tracked), tags=tags) + if not tracked: + continue + + try: + not_found = client.update_heartbeats(list(tracked)) + except grpc.RpcError as e: + logger.warning("taskworker.worker.heartbeat.failed", extra={"error": e, **tags}) + continue + + if not_found: + tracked.difference_update(not_found) + metrics.incr("taskworker.worker.heartbeat.not_found", value=len(not_found), tags=tags) + + +class HeartbeatTracker: + """ + Runs `heartbeat_process` once there is something to track, and restarts it if it dies. + + The parent keeps the set of tracked ids so a restarted process starts from it. + """ + + def __init__( + self, + app_module: str, + broker_service: str, + processing_pool_name: str, + mp_context: ForkContext | SpawnContext | ForkServerContext, + metrics: MetricsBackend, + ) -> None: + self._metrics = metrics + self._app_module = app_module + self._broker_service = broker_service + self._processing_pool_name = processing_pool_name + self._mp_context = mp_context + self._shutdown_event = mp_context.Event() + self._lock = threading.Lock() + self._ids: set[str] = set() + self._commands: multiprocessing.Queue[HeartbeatCommand] | None = None + self._process: BaseProcess | None = None + self._supervisor: threading.Thread | None = None + + def track(self, task_id: str) -> None: + with self._lock: + self._ids.add(task_id) + if self._commands is not None: + self._commands.put(("track", [task_id])) + + def untrack(self, task_ids: list[str]) -> None: + with self._lock: + removed = [task_id for task_id in task_ids if task_id in self._ids] + if not removed: + return + self._ids.difference_update(removed) + if self._commands is not None: + self._commands.put(("untrack", removed)) + + def _spawn(self) -> None: + with self._lock: + # A fresh queue, so commands stranded in a dead process's queue are dropped + # and the new process starts from the parent's set instead. + self._commands = self._mp_context.Queue() + self._process = self._mp_context.Process( + name="taskworker-heartbeat", + target=heartbeat_process, + args=( + self._app_module, + self._broker_service, + self._processing_pool_name, + list(self._ids), + self._commands, + self._shutdown_event, + ), + daemon=True, + ) + self._process.start() + + def _supervise(self) -> None: + while not self._shutdown_event.wait(HEARTBEAT_INTERVAL_SEC): + # Start lazily, so workers without heartbeat tasks don't pay for the process. + if self._process is None: + if self._ids: + self._spawn() + elif not self._process.is_alive(): + self._metrics.incr( + "taskworker.worker.heartbeat.restarted", + tags={"processing_pool": self._processing_pool_name}, + ) + logger.warning( + "taskworker.worker.heartbeat.restarted", + extra={"processing_pool": self._processing_pool_name}, + ) + self._spawn() + + def start(self) -> None: + self._supervisor = threading.Thread( + name="heartbeat-supervisor", target=self._supervise, daemon=True + ) + self._supervisor.start() + + def shutdown(self) -> None: + self._shutdown_event.set() + if self._process is not None: + self._process.join(timeout=HEARTBEAT_INTERVAL_SEC * 2) + if self._process.is_alive(): + self._process.terminate() diff --git a/clients/python/src/taskbroker_client/worker/push_clients.py b/clients/python/src/taskbroker_client/worker/push_clients.py index 31d7d5d8..e6f4963c 100644 --- a/clients/python/src/taskbroker_client/worker/push_clients.py +++ b/clients/python/src/taskbroker_client/worker/push_clients.py @@ -5,6 +5,7 @@ import grpc from sentry_protos.taskbroker.v1.taskbroker_pb2 import ( + BatchUpdateHeartbeatRequest, SetBatchActivationStatusRequest, SetTaskStatusRequest, ) @@ -26,6 +27,9 @@ logger = logging.getLogger(__name__) +# Well below the 10s minimum heartbeat, so a hung call can't eat a whole interval. +HEARTBEAT_RPC_TIMEOUT_SEC = 5 + def service_authority(service: str) -> str: """The ``:authority`` to send for a push broker service target. @@ -162,3 +166,23 @@ def update_tasks( if exception: raise exception + + def update_heartbeats(self, ids: list[str]) -> list[str]: + """ + Renew the heartbeat of processing activations. Returns the ids the broker + no longer has processing. + """ + tags = {"processing_pool": self._processing_pool_name} + try: + response = self._stub.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) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index cdd71f71..7c01377d 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -50,6 +50,7 @@ TaskbrokerClient, parse_rpc_secret_list, ) +from taskbroker_client.worker.heartbeat import HeartbeatTracker from taskbroker_client.worker.push_clients import PushTaskbrokerClient from taskbroker_client.worker.workerchild import ChildMessage, child_process @@ -160,9 +161,15 @@ class WorkerServicer(taskbroker_pb2_grpc.WorkerServiceServicer): gRPC servicer that receives task activations pushed from the broker """ - def __init__(self, worker: TaskWorkerProcessingPool, push_task_timeout: float = 5) -> None: + def __init__( + self, + worker: TaskWorkerProcessingPool, + push_task_timeout: float = 5, + heartbeats: HeartbeatTracker | None = None, + ) -> None: self.worker_pool = worker self.push_task_timeout = push_task_timeout + self.heartbeats = heartbeats def PushTask( self, @@ -183,8 +190,15 @@ def PushTask( receive_timestamp=time.monotonic(), ) + # Track before queueing, so a fast result can't untrack it before it is tracked. + heartbeats = self.heartbeats if request.task.HasField("heartbeat") else None + if heartbeats: + heartbeats.track(request.task.id) + # Push the task to the worker queue (wait at most N seconds) if not self.worker_pool.push_task(inflight, timeout=self.push_task_timeout): + if heartbeats: + heartbeats.untrack([request.task.id]) self.worker_pool._metrics.incr( "taskworker.worker.push_rpc", tags={"result": "busy", "processing_pool": self.worker_pool._processing_pool_name}, @@ -327,6 +341,14 @@ def __init__( self._warmup_timeout = warmup_timeout + self._heartbeats = HeartbeatTracker( + app_module=app_module, + broker_service=broker_service, + processing_pool_name=self._processing_pool_name, + mp_context=self._mp_context, + metrics=app.metrics, + ) + def _create_client( self, service: str, @@ -377,6 +399,7 @@ def _send_results( try: self.client.update_tasks(results) + self._heartbeats.untrack([result.task_id for result in results]) self._setstatus_backoff_seconds = 0 return None except grpc.RpcError as e: @@ -388,6 +411,7 @@ def _send_results( if e.code() != grpc.StatusCode.NOT_FOUND: # If the task was not found, we can't update it, so we should just return None raise RequeueException(f"Failed to update task batch: {e}") + self._heartbeats.untrack([result.task_id for result in results]) except HostTemporarilyUnavailable as e: self._setstatus_backoff_seconds = min( self._setstatus_backoff_seconds + 4, MAX_BACKOFF_SECONDS_WHEN_HOST_UNAVAILABLE @@ -493,6 +517,7 @@ def start(self) -> int: self.worker_pool.start_metrics_thread() self.worker_pool.start_result_thread() self.worker_pool.start_spawn_children_thread() + self._heartbeats.start() server: grpc.Server | None = None server_started = False @@ -535,7 +560,8 @@ def signal_handler(*args: Any) -> None: ) taskbroker_pb2_grpc.add_WorkerServiceServicer_to_server( - WorkerServicer(self.worker_pool, self._push_task_timeout), server + WorkerServicer(self.worker_pool, self._push_task_timeout, self._heartbeats), + server, ) # The health service is used by the K8s readiness check @@ -603,6 +629,7 @@ def shutdown(self) -> None: self._stop_health_check_thread() self._shutdown_signal.set() self.worker_pool.shutdown() + self._heartbeats.shutdown() class TaskWorker: diff --git a/clients/python/tests/test_registry.py b/clients/python/tests/test_registry.py index 34fab8ec..30c926ff 100644 --- a/clients/python/tests/test_registry.py +++ b/clients/python/tests/test_registry.py @@ -515,3 +515,56 @@ def test_registry_create_namespace_is_raw_mode() -> None: ) assert not registry.create_namespace("tests").is_raw_mode assert registry.create_namespace("tests.raw", is_raw_mode=True).is_raw_mode + + +def test_register_heartbeat_sets_activation_field() -> None: + namespace = TaskNamespace( + name="tests", + application="acme", + producer_factory=producer_factory, + router=DefaultRouter(), + metrics=NoOpMetricsBackend(), + retry=None, + ) + + @namespace.register(name="tests.heartbeat", processing_deadline_duration=60, heartbeat=10) + def heartbeat_task() -> None: + raise NotImplementedError + + assert heartbeat_task.create_activation([], {}).heartbeat == 10 + + +@pytest.mark.parametrize( + "heartbeat,match", + [ + pytest.param(9, "at least 10s", id="too_small"), + pytest.param(60, "less than its processing_deadline_duration", id="not_below_deadline"), + ], +) +def test_register_rejects_invalid_heartbeat(heartbeat: int, match: str) -> None: + namespace = TaskNamespace( + name="tests", + application="acme", + producer_factory=producer_factory, + router=DefaultRouter(), + metrics=NoOpMetricsBackend(), + retry=None, + ) + + with pytest.raises(ValueError, match=match): + + @namespace.register( + name="tests.heartbeat", processing_deadline_duration=60, heartbeat=heartbeat + ) + def heartbeat_task() -> None: + raise NotImplementedError + + +def test_raw_mode_namespace_rejects_heartbeat() -> None: + namespace = raw_namespace() + + with pytest.raises(ValueError, match="sets heartbeat"): + + @namespace.register(name="tests.raw.consume", heartbeat=10) + def consume(message_bytes: bytes) -> None: + raise NotImplementedError diff --git a/clients/python/tests/worker/test_heartbeat.py b/clients/python/tests/worker/test_heartbeat.py new file mode 100644 index 00000000..ccd149fa --- /dev/null +++ b/clients/python/tests/worker/test_heartbeat.py @@ -0,0 +1,143 @@ +import queue +from unittest import mock + +import grpc +from sentry_protos.taskbroker.v1.taskbroker_pb2 import ( + BatchUpdateHeartbeatResponse, + PushTaskRequest, + TaskActivation, +) + +from taskbroker_client.metrics import NoOpMetricsBackend +from taskbroker_client.worker.heartbeat import HeartbeatCommand, HeartbeatTracker, heartbeat_process +from taskbroker_client.worker.push_clients import PushTaskbrokerClient +from taskbroker_client.worker.worker import PushTaskWorker, WorkerServicer + + +class FakeEvent: + """Lets `heartbeat_process` and `_supervise` run a fixed number of iterations.""" + + def __init__(self, iterations: int) -> None: + self.iterations = iterations + + def wait(self, timeout: float) -> bool: + self.iterations -= 1 + return self.iterations < 0 + + def is_set(self) -> bool: + return self.iterations < 0 + + +class UnavailableError(grpc.RpcError): + def code(self) -> grpc.StatusCode: + return grpc.StatusCode.UNAVAILABLE + + +def test_update_heartbeats_returns_not_found_ids() -> None: + client = PushTaskbrokerClient( + service="localhost:50051", application="tests", metrics=NoOpMetricsBackend() + ) + with mock.patch.object(client, "_stub") as stub: + stub.BatchUpdateHeartbeat.return_value = BatchUpdateHeartbeatResponse(not_found_ids=["b"]) + assert client.update_heartbeats(["a", "b"]) == ["b"] + + (request,), _ = stub.BatchUpdateHeartbeat.call_args + assert list(request.ids) == ["a", "b"] + + +def test_heartbeat_process_applies_commands_and_drops_not_found() -> None: + commands: queue.Queue[HeartbeatCommand] = queue.Queue() + commands.put(("track", ["b", "c"])) + commands.put(("untrack", ["c"])) + + with ( + mock.patch("taskbroker_client.worker.heartbeat.signal.signal"), + mock.patch.object( + PushTaskbrokerClient, + "update_heartbeats", + side_effect=[["a"], UnavailableError(), []], + ) as update, + ): + heartbeat_process( + "examples.app:app", "localhost:50051", "test", ["a"], commands, FakeEvent(3) # type: ignore[arg-type] + ) + + assert [sorted(c.args[0]) for c in update.call_args_list] == [["a", "b"], ["b"], ["b"]] + + +def _tracker(mp_context: mock.MagicMock) -> HeartbeatTracker: + return HeartbeatTracker( + app_module="examples.app:app", + broker_service="localhost:50051", + processing_pool_name="test", + mp_context=mp_context, + metrics=NoOpMetricsBackend(), + ) + + +def test_tracker_restarts_dead_process_with_tracked_ids() -> None: + mp_context = mock.MagicMock() + tracker = _tracker(mp_context) + tracker.track("a") + tracker.track("b") + tracker.untrack(["a", "missing"]) + + # First tick starts it, second tick sees it alive (MagicMock) and leaves it. + tracker._shutdown_event = FakeEvent(2) # type: ignore[assignment] + tracker._supervise() + assert mp_context.Process.call_count == 1 + assert mp_context.Process.call_args.kwargs["args"][3] == ["b"] + + mp_context.Process.return_value.is_alive.return_value = False + tracker._shutdown_event = FakeEvent(1) # type: ignore[assignment] + tracker._supervise() + assert mp_context.Process.call_count == 2 + + +def test_tracker_does_not_start_without_tracked_ids() -> None: + mp_context = mock.MagicMock() + tracker = _tracker(mp_context) + tracker._shutdown_event = FakeEvent(2) # type: ignore[assignment] + tracker._supervise() + mp_context.Process.assert_not_called() + + +def _push_worker() -> PushTaskWorker: + return PushTaskWorker( + app_module="examples.app:app", + broker_service="127.0.0.1:50051", + max_child_task_count=100, + process_type="fork", + drain_file_path=None, + ) + + +def test_push_task_tracks_heartbeat_tasks_only() -> None: + taskworker = _push_worker() + servicer = WorkerServicer(taskworker.worker_pool, heartbeats=taskworker._heartbeats) + + with mock.patch.object(taskworker.worker_pool, "push_task", return_value=True): + servicer.PushTask( + PushTaskRequest(task=TaskActivation(id="beat", heartbeat=10)), mock.MagicMock() + ) + servicer.PushTask(PushTaskRequest(task=TaskActivation(id="plain")), mock.MagicMock()) + + assert taskworker._heartbeats._ids == {"beat"} + + with mock.patch.object(taskworker.worker_pool, "push_task", return_value=False): + servicer.PushTask( + PushTaskRequest(task=TaskActivation(id="busy", heartbeat=10)), mock.MagicMock() + ) + + assert taskworker._heartbeats._ids == {"beat"} + + +def test_send_results_untracks_sent_tasks() -> None: + taskworker = _push_worker() + taskworker._heartbeats.track("a") + result = mock.MagicMock(task_id="a", receive_timestamp=0) + + with mock.patch.object(taskworker.client, "update_tasks"): + taskworker._send_results([result]) + + assert taskworker._heartbeats._ids == set() diff --git a/uv.lock b/uv.lock index 737b7c65..f06e8287 100644 --- a/uv.lock +++ b/uv.lock @@ -690,7 +690,7 @@ wheels = [ [[package]] name = "sentry-protos" -version = "0.26.1" +version = "0.106.0" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ { name = "grpc-stubs" }, @@ -698,7 +698,7 @@ dependencies = [ { name = "protobuf" }, ] wheels = [ - { url = "https://pypi.devinfra.sentry.io/wheels/sentry_protos-0.26.1-py3-none-any.whl", hash = "sha256:66bc22be8ac3efeceeb09af30a3ca55943ad83a3d28905263255f74a02206822" }, + { url = "https://pypi.devinfra.sentry.io/wheels/sentry_protos-0.106.0-py3-none-any.whl", hash = "sha256:120a7845a82c63e100bcba88ee265a1964f1f648094f8513d2a6e1dcf40ab38e" }, ] [[package]] @@ -855,7 +855,7 @@ requires-dist = [ { name = "redis", specifier = ">=3.4.1" }, { name = "redis-py-cluster", marker = "extra == 'cluster'", specifier = ">=2.1.0" }, { name = "sentry-arroyo", specifier = ">=2.41.1" }, - { name = "sentry-protos", specifier = ">=0.26.1" }, + { name = "sentry-protos", specifier = ">=0.106.0" }, { name = "sentry-sdk", extras = ["http2"], specifier = ">=2.73.0" }, { name = "setuptools", marker = "extra == 'examples'", specifier = ">=80.0" }, { name = "zstandard", specifier = ">=0.18.0" }, From b7fa36a22369b115b42c40c1d90cc09126da380b Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Tue, 10 Nov 2026 15:00:09 -0500 Subject: [PATCH 2/2] feat(taskworker): Send heartbeats from pull workers and add integration test Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 38 +++++ Makefile | 7 +- .../src/taskbroker_client/worker/client.py | 24 ++++ .../src/taskbroker_client/worker/heartbeat.py | 125 +++++++++++------ .../taskbroker_client/worker/push_clients.py | 4 +- .../src/taskbroker_client/worker/worker.py | 23 ++- clients/python/tests/worker/test_heartbeat.py | 115 +++++++++++---- .../integration_tests/test_heartbeat.py | 132 ++++++++++++++++++ integration_tests/integration_tests/worker.py | 6 + 9 files changed, 400 insertions(+), 74 deletions(-) create mode 100644 integration_tests/integration_tests/test_heartbeat.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6ffa4e0b..a5c1377b 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/Makefile b/Makefile index 3a809312..006a8ac6 100644 --- a/Makefile +++ b/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 diff --git a/clients/python/src/taskbroker_client/worker/client.py b/clients/python/src/taskbroker_client/worker/client.py index 98aa1fa0..a32fe728 100644 --- a/clients/python/src/taskbroker_client/worker/client.py +++ b/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) diff --git a/clients/python/src/taskbroker_client/worker/heartbeat.py b/clients/python/src/taskbroker_client/worker/heartbeat.py index 98b6bab5..9d11a1c7 100644 --- a/clients/python/src/taskbroker_client/worker/heartbeat.py +++ b/clients/python/src/taskbroker_client/worker/heartbeat.py @@ -5,6 +5,8 @@ import queue import signal import threading +from collections import defaultdict +from collections.abc import Callable from multiprocessing.context import ForkContext, ForkServerContext, SpawnContext from multiprocessing.process import BaseProcess from multiprocessing.synchronize import Event @@ -14,42 +16,66 @@ from taskbroker_client.app import import_app from taskbroker_client.metrics import MetricsBackend +from taskbroker_client.types import InflightTaskActivation +from taskbroker_client.worker.client import TaskbrokerClient from taskbroker_client.worker.push_clients import PushTaskbrokerClient logger = logging.getLogger(__name__) HEARTBEAT_INTERVAL_SEC = 1.0 -HeartbeatCommand = tuple[Literal["track", "untrack"], list[str]] +# ("track", {id: host}) or ("untrack", {id: ""}) +HeartbeatCommand = tuple[Literal["track", "untrack"], dict[str, str]] def heartbeat_process( app_module: str, - broker_service: str, + broker_service: str | None, + broker_hosts: list[str] | None, processing_pool_name: str, - ids: list[str], + tracked: dict[str, str], commands: queue.Queue[HeartbeatCommand], shutdown_event: Event, ) -> None: - """Renew the heartbeat of every tracked activation once per interval.""" + """ + Renew the heartbeat of every tracked activation once per interval. + + `tracked` maps an activation id to the broker host it came from. Push mode sends + every id to `broker_service`; pull mode sends each id to its own host. + """ # The parent owns shutdown and tells us through shutdown_event. signal.signal(signal.SIGINT, signal.SIG_IGN) signal.signal(signal.SIGTERM, signal.SIG_IGN) app = import_app(app_module) metrics = app.metrics - client = PushTaskbrokerClient( - service=broker_service, - application=app.name, - metrics=metrics, - rpc_secret=app.config["rpc_secret"], - grpc_config=app.config["grpc_config"], - processing_pool_name=processing_pool_name, - ) tags = {"processing_pool": processing_pool_name} - tracked = set(ids) - while not shutdown_event.wait(HEARTBEAT_INTERVAL_SEC): + send: Callable[[str, list[str]], list[str]] + if broker_service is not None: + push_client = PushTaskbrokerClient( + service=broker_service, + application=app.name, + metrics=metrics, + rpc_secret=app.config["rpc_secret"], + grpc_config=app.config["grpc_config"], + processing_pool_name=processing_pool_name, + ) + send = lambda host, ids: push_client.update_heartbeats(ids) # noqa: E731 + else: + assert broker_hosts, "Pull mode heartbeats need broker_hosts" + pull_client = TaskbrokerClient( + hosts=broker_hosts, + application=app.name, + metrics=metrics, + rpc_secret=app.config["rpc_secret"], + grpc_config=app.config["grpc_config"], + processing_pool_name=processing_pool_name, + ) + send = pull_client.update_heartbeats + + tracked = dict(tracked) + while not shutdown_event.is_set(): while True: try: action, command_ids = commands.get_nowait() @@ -58,69 +84,87 @@ def heartbeat_process( if action == "track": tracked.update(command_ids) else: - tracked.difference_update(command_ids) + for task_id in command_ids: + tracked.pop(task_id, None) metrics.gauge("taskworker.worker.heartbeat.tracked", len(tracked), tags=tags) - if not tracked: - continue - try: - not_found = client.update_heartbeats(list(tracked)) - except grpc.RpcError as e: - logger.warning("taskworker.worker.heartbeat.failed", extra={"error": e, **tags}) - continue + by_host: defaultdict[str, list[str]] = defaultdict(list) + for task_id, host in tracked.items(): + by_host["" if broker_service is not None else host].append(task_id) + + for host, ids in by_host.items(): + try: + not_found = send(host, ids) + except grpc.RpcError as e: + logger.warning( + "taskworker.worker.heartbeat.failed", extra={"error": e, "host": host, **tags} + ) + continue + + for task_id in not_found: + tracked.pop(task_id, None) + if not_found: + metrics.incr( + "taskworker.worker.heartbeat.not_found", value=len(not_found), tags=tags + ) - if not_found: - tracked.difference_update(not_found) - metrics.incr("taskworker.worker.heartbeat.not_found", value=len(not_found), tags=tags) + shutdown_event.wait(HEARTBEAT_INTERVAL_SEC) class HeartbeatTracker: """ - Runs `heartbeat_process` once there is something to track, and restarts it if it dies. + Runs `heartbeat_process` and restarts it if it dies. - The parent keeps the set of tracked ids so a restarted process starts from it. + The parent keeps the tracked ids so a restarted process starts from them. + Pass `broker_service` for push mode or `broker_hosts` for pull mode. """ def __init__( self, app_module: str, - broker_service: str, processing_pool_name: str, mp_context: ForkContext | SpawnContext | ForkServerContext, metrics: MetricsBackend, + broker_service: str | None = None, + broker_hosts: list[str] | None = None, ) -> None: self._metrics = metrics self._app_module = app_module self._broker_service = broker_service + self._broker_hosts = broker_hosts self._processing_pool_name = processing_pool_name self._mp_context = mp_context self._shutdown_event = mp_context.Event() self._lock = threading.Lock() - self._ids: set[str] = set() + self._tracked: dict[str, str] = {} self._commands: multiprocessing.Queue[HeartbeatCommand] | None = None self._process: BaseProcess | None = None self._supervisor: threading.Thread | None = None - def track(self, task_id: str) -> None: + def track(self, inflight: InflightTaskActivation) -> None: + """Track `inflight` if its task has a heartbeat.""" + if not inflight.activation.HasField("heartbeat"): + return with self._lock: - self._ids.add(task_id) + self._tracked[inflight.activation.id] = inflight.host if self._commands is not None: - self._commands.put(("track", [task_id])) + self._commands.put(("track", {inflight.activation.id: inflight.host})) def untrack(self, task_ids: list[str]) -> None: with self._lock: - removed = [task_id for task_id in task_ids if task_id in self._ids] + removed = {task_id: "" for task_id in task_ids if task_id in self._tracked} if not removed: return - self._ids.difference_update(removed) + for task_id in removed: + del self._tracked[task_id] if self._commands is not None: self._commands.put(("untrack", removed)) def _spawn(self) -> None: with self._lock: # A fresh queue, so commands stranded in a dead process's queue are dropped - # and the new process starts from the parent's set instead. + # and the new process starts from the parent's ids instead. self._commands = self._mp_context.Queue() self._process = self._mp_context.Process( name="taskworker-heartbeat", @@ -128,8 +172,9 @@ def _spawn(self) -> None: args=( self._app_module, self._broker_service, + self._broker_hosts, self._processing_pool_name, - list(self._ids), + dict(self._tracked), self._commands, self._shutdown_event, ), @@ -139,11 +184,7 @@ def _spawn(self) -> None: def _supervise(self) -> None: while not self._shutdown_event.wait(HEARTBEAT_INTERVAL_SEC): - # Start lazily, so workers without heartbeat tasks don't pay for the process. - if self._process is None: - if self._ids: - self._spawn() - elif not self._process.is_alive(): + if self._process is not None and not self._process.is_alive(): self._metrics.incr( "taskworker.worker.heartbeat.restarted", tags={"processing_pool": self._processing_pool_name}, @@ -155,6 +196,8 @@ def _supervise(self) -> None: self._spawn() def start(self) -> None: + # Start now, so the app import is done before the first task needs a heartbeat. + self._spawn() self._supervisor = threading.Thread( name="heartbeat-supervisor", target=self._supervise, daemon=True ) diff --git a/clients/python/src/taskbroker_client/worker/push_clients.py b/clients/python/src/taskbroker_client/worker/push_clients.py index e6f4963c..40850f8e 100644 --- a/clients/python/src/taskbroker_client/worker/push_clients.py +++ b/clients/python/src/taskbroker_client/worker/push_clients.py @@ -14,6 +14,7 @@ from taskbroker_client.metrics import MetricsBackend from taskbroker_client.types import ProcessingResult from taskbroker_client.worker.client import ( + HEARTBEAT_RPC_TIMEOUT_SEC, MAX_ACTIVATION_SIZE, HealthCheckSettings, RequestSignatureInterceptor, @@ -27,9 +28,6 @@ logger = logging.getLogger(__name__) -# Well below the 10s minimum heartbeat, so a hung call can't eat a whole interval. -HEARTBEAT_RPC_TIMEOUT_SEC = 5 - def service_authority(service: str) -> str: """The ``:authority`` to send for a push broker service target. diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 7c01377d..14049b07 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -191,14 +191,13 @@ def PushTask( ) # Track before queueing, so a fast result can't untrack it before it is tracked. - heartbeats = self.heartbeats if request.task.HasField("heartbeat") else None - if heartbeats: - heartbeats.track(request.task.id) + if self.heartbeats: + self.heartbeats.track(inflight) # Push the task to the worker queue (wait at most N seconds) if not self.worker_pool.push_task(inflight, timeout=self.push_task_timeout): - if heartbeats: - heartbeats.untrack([request.task.id]) + if self.heartbeats: + self.heartbeats.untrack([request.task.id]) self.worker_pool._metrics.incr( "taskworker.worker.push_rpc", tags={"result": "busy", "processing_pool": self.worker_pool._processing_pool_name}, @@ -714,6 +713,14 @@ def __init__( self._processing_pool_name: str = processing_pool_name or "unknown" + self._heartbeats = HeartbeatTracker( + app_module=app_module, + broker_hosts=broker_hosts, + processing_pool_name=self._processing_pool_name, + mp_context=self._mp_context, + metrics=app.metrics, + ) + def start(self) -> int: """ This starts a loop that runs until the worker completes its `max_task_count` or it is killed. @@ -721,6 +728,7 @@ def start(self) -> int: self.worker_pool.start_metrics_thread() self.worker_pool.start_result_thread() self.worker_pool.start_spawn_children_thread() + self._heartbeats.start() # Record the request and let the loop below act on it. Raising from a # handler unwinds at an arbitrary bytecode and can leave the locks held @@ -804,6 +812,9 @@ def _send_update_task( try: next_task = self.client.update_task(result, fetch_next) + self._heartbeats.untrack([result.task_id]) + if next_task: + self._heartbeats.track(next_task) self._setstatus_backoff_seconds = 0 return next_task except grpc.RpcError as e: @@ -872,6 +883,7 @@ def fetch_task(self) -> InflightTaskActivation | None: ) return None + self._heartbeats.track(activation) return activation def shutdown(self) -> None: @@ -880,6 +892,7 @@ def shutdown(self) -> None: """ self._shutdown_signal.set() self.worker_pool.shutdown() + self._heartbeats.shutdown() class TaskWorkerProcessingPool: diff --git a/clients/python/tests/worker/test_heartbeat.py b/clients/python/tests/worker/test_heartbeat.py index ccd149fa..1c5070bb 100644 --- a/clients/python/tests/worker/test_heartbeat.py +++ b/clients/python/tests/worker/test_heartbeat.py @@ -1,4 +1,5 @@ import queue +import time from unittest import mock import grpc @@ -9,9 +10,11 @@ ) from taskbroker_client.metrics import NoOpMetricsBackend +from taskbroker_client.types import InflightTaskActivation +from taskbroker_client.worker.client import TaskbrokerClient from taskbroker_client.worker.heartbeat import HeartbeatCommand, HeartbeatTracker, heartbeat_process from taskbroker_client.worker.push_clients import PushTaskbrokerClient -from taskbroker_client.worker.worker import PushTaskWorker, WorkerServicer +from taskbroker_client.worker.worker import PushTaskWorker, TaskWorker, WorkerServicer class FakeEvent: @@ -33,7 +36,18 @@ def code(self) -> grpc.StatusCode: return grpc.StatusCode.UNAVAILABLE -def test_update_heartbeats_returns_not_found_ids() -> None: +def _inflight( + task_id: str, host: str = "broker:50051", heartbeat: bool = True +) -> InflightTaskActivation: + activation = ( + TaskActivation(id=task_id, heartbeat=10) if heartbeat else TaskActivation(id=task_id) + ) + return InflightTaskActivation( + activation=activation, host=host, receive_timestamp=time.monotonic() + ) + + +def test_push_client_update_heartbeats_returns_not_found_ids() -> None: client = PushTaskbrokerClient( service="localhost:50051", application="tests", metrics=NoOpMetricsBackend() ) @@ -45,10 +59,23 @@ def test_update_heartbeats_returns_not_found_ids() -> None: assert list(request.ids) == ["a", "b"] -def test_heartbeat_process_applies_commands_and_drops_not_found() -> None: +def test_pull_client_update_heartbeats_uses_the_host_stub() -> None: + client = TaskbrokerClient( + hosts=["broker-0:50051"], application="tests", metrics=NoOpMetricsBackend() + ) + with mock.patch.object(client, "_get_stub") as get_stub: + get_stub.return_value.BatchUpdateHeartbeat.return_value = BatchUpdateHeartbeatResponse( + not_found_ids=["b"] + ) + assert client.update_heartbeats("broker-1:50051", ["a", "b"]) == ["b"] + + get_stub.assert_called_once_with("broker-1:50051") + + +def test_push_heartbeat_process_applies_commands_and_drops_not_found() -> None: commands: queue.Queue[HeartbeatCommand] = queue.Queue() - commands.put(("track", ["b", "c"])) - commands.put(("untrack", ["c"])) + commands.put(("track", {"b": "", "c": ""})) + commands.put(("untrack", {"c": ""})) with ( mock.patch("taskbroker_client.worker.heartbeat.signal.signal"), @@ -59,12 +86,39 @@ def test_heartbeat_process_applies_commands_and_drops_not_found() -> None: ) as update, ): heartbeat_process( - "examples.app:app", "localhost:50051", "test", ["a"], commands, FakeEvent(3) # type: ignore[arg-type] + "examples.app:app", + "localhost:50051", + None, + "test", + {"a": ""}, + commands, + FakeEvent(2), # type: ignore[arg-type] ) assert [sorted(c.args[0]) for c in update.call_args_list] == [["a", "b"], ["b"], ["b"]] +def test_pull_heartbeat_process_sends_each_id_to_its_host() -> None: + commands: queue.Queue[HeartbeatCommand] = queue.Queue() + + with ( + mock.patch("taskbroker_client.worker.heartbeat.signal.signal"), + mock.patch.object(TaskbrokerClient, "update_heartbeats", return_value=[]) as update, + ): + heartbeat_process( + "examples.app:app", + None, + ["broker-0:50051", "broker-1:50051"], + "test", + {"a": "broker-0:50051", "b": "broker-1:50051", "c": "broker-0:50051"}, + commands, + FakeEvent(0), # type: ignore[arg-type] + ) + + sent = {c.args[0]: sorted(c.args[1]) for c in update.call_args_list} + assert sent == {"broker-0:50051": ["a", "c"], "broker-1:50051": ["b"]} + + def _tracker(mp_context: mock.MagicMock) -> HeartbeatTracker: return HeartbeatTracker( app_module="examples.app:app", @@ -78,28 +132,23 @@ def _tracker(mp_context: mock.MagicMock) -> HeartbeatTracker: def test_tracker_restarts_dead_process_with_tracked_ids() -> None: mp_context = mock.MagicMock() tracker = _tracker(mp_context) - tracker.track("a") - tracker.track("b") + tracker._spawn() + tracker.track(_inflight("a")) + tracker.track(_inflight("b", host="broker-1:50051")) + tracker.track(_inflight("plain", heartbeat=False)) tracker.untrack(["a", "missing"]) - # First tick starts it, second tick sees it alive (MagicMock) and leaves it. - tracker._shutdown_event = FakeEvent(2) # type: ignore[assignment] + # Alive (MagicMock default): left alone. + tracker._shutdown_event = FakeEvent(1) # type: ignore[assignment] tracker._supervise() assert mp_context.Process.call_count == 1 - assert mp_context.Process.call_args.kwargs["args"][3] == ["b"] + # Dead: restarted with the parent's tracked ids. mp_context.Process.return_value.is_alive.return_value = False tracker._shutdown_event = FakeEvent(1) # type: ignore[assignment] tracker._supervise() assert mp_context.Process.call_count == 2 - - -def test_tracker_does_not_start_without_tracked_ids() -> None: - mp_context = mock.MagicMock() - tracker = _tracker(mp_context) - tracker._shutdown_event = FakeEvent(2) # type: ignore[assignment] - tracker._supervise() - mp_context.Process.assert_not_called() + assert mp_context.Process.call_args.kwargs["args"][4] == {"b": "broker-1:50051"} def _push_worker() -> PushTaskWorker: @@ -122,22 +171,40 @@ def test_push_task_tracks_heartbeat_tasks_only() -> None: ) servicer.PushTask(PushTaskRequest(task=TaskActivation(id="plain")), mock.MagicMock()) - assert taskworker._heartbeats._ids == {"beat"} + assert set(taskworker._heartbeats._tracked) == {"beat"} with mock.patch.object(taskworker.worker_pool, "push_task", return_value=False): servicer.PushTask( PushTaskRequest(task=TaskActivation(id="busy", heartbeat=10)), mock.MagicMock() ) - assert taskworker._heartbeats._ids == {"beat"} + assert set(taskworker._heartbeats._tracked) == {"beat"} -def test_send_results_untracks_sent_tasks() -> None: +def test_push_send_results_untracks_sent_tasks() -> None: taskworker = _push_worker() - taskworker._heartbeats.track("a") + taskworker._heartbeats.track(_inflight("a")) result = mock.MagicMock(task_id="a", receive_timestamp=0) with mock.patch.object(taskworker.client, "update_tasks"): taskworker._send_results([result]) - assert taskworker._heartbeats._ids == set() + assert taskworker._heartbeats._tracked == {} + + +def test_pull_worker_tracks_fetched_and_next_tasks() -> None: + taskworker = TaskWorker( + app_module="examples.app:app", + broker_hosts=["broker-0:50051"], + max_child_task_count=100, + process_type="fork", + ) + + with mock.patch.object(taskworker.client, "get_task", return_value=_inflight("a")): + taskworker.fetch_task() + assert taskworker._heartbeats._tracked == {"a": "broker:50051"} + + result = mock.MagicMock(task_id="a") + with mock.patch.object(taskworker.client, "update_task", return_value=_inflight("b")): + taskworker._send_update_task(result, None) + assert taskworker._heartbeats._tracked == {"b": "broker:50051"} diff --git a/integration_tests/integration_tests/test_heartbeat.py b/integration_tests/integration_tests/test_heartbeat.py new file mode 100644 index 00000000..957023bf --- /dev/null +++ b/integration_tests/integration_tests/test_heartbeat.py @@ -0,0 +1,132 @@ +import signal +import sqlite3 +import subprocess +import time +from uuid import uuid4 + +import orjson +import yaml +from google.protobuf.timestamp_pb2 import Timestamp +from sentry_protos.taskbroker.v1.taskbroker_pb2 import ( + OnAttemptsExceeded, + RetryState, + TaskActivation, +) + +from integration_tests.helpers import ( + TASKBROKER_BIN, + TESTS_OUTPUT_ROOT, + TaskbrokerConfig, + create_topic, + get_available_ports, + get_num_tasks_in_sqlite, + send_custom_messages_to_topic, +) +from integration_tests.worker import TaskWorkerClient + +TEST_OUTPUT_PATH = TESTS_OUTPUT_ROOT / "test_heartbeat" + +# Short enough to keep the test fast. The client enforces a 10s minimum at task +# registration, but the broker accepts any value. +HEARTBEAT_SEC = 2 + + +def heartbeat_activation() -> TaskActivation: + return TaskActivation( + id=uuid4().hex, + namespace="integration_tests", + taskname="integration_tests.say_hello", + parameters=orjson.dumps({"args": [], "kwargs": {}}).decode("utf-8"), + retry_state=RetryState( + attempts=0, + max_attempts=3, + on_attempts_exceeded=OnAttemptsExceeded.ON_ATTEMPTS_EXCEEDED_DISCARD, + ), + processing_deadline_duration=3000, + received_at=Timestamp(seconds=int(time.time())), + heartbeat=HEARTBEAT_SEC, + ) + + +def task_state(config: TaskbrokerConfig, task_id: str) -> tuple[str, int]: + con = sqlite3.connect(config.db_path) + try: + status, attempts = con.execute( + "SELECT status, processing_attempts FROM inflight_taskactivations WHERE id = ?", + (task_id,), + ).fetchone() + finally: + con.close() + return status, attempts + + +def test_heartbeat() -> None: + """ + A pull worker fetches two tasks that have a heartbeat. It renews one of them with + `BatchUpdateHeartbeat` and lets the other one lapse. The broker + should keep the renewed task processing and reset the lapsed one to pending, and + report the lapsed one in `not_found_ids` from then on. + """ + curr_time = int(time.time()) + topic_name = f"taskworker-heartbeat-{curr_time}" + (grpc_port,) = get_available_ports(1) + host = f"127.0.0.1:{grpc_port}" + + create_topic(topic_name, 1) + TEST_OUTPUT_PATH.mkdir(parents=True, exist_ok=True) + db_name = f"db_0_{curr_time}_test_heartbeat" + config = TaskbrokerConfig( + db_name=db_name, + db_path=str(TEST_OUTPUT_PATH / f"{db_name}.sqlite"), + max_pending_count=1000, + kafka_topic=topic_name, + kafka_deadletter_topic=f"{topic_name}-dlq", + kafka_consumer_group=topic_name, + kafka_auto_offset_reset="earliest", + grpc_port=grpc_port, + ) + config_path = TEST_OUTPUT_PATH / "config_0_test_heartbeat.yml" + with open(config_path, "w") as f: + # Upkeep skips deadline resets right after startup by default. + yaml.safe_dump({**config.to_dict(), "upkeep_deadline_reset_skip_after_startup_sec": 0}, f) + + send_custom_messages_to_topic(topic_name, [heartbeat_activation(), heartbeat_activation()]) + + log_path = TEST_OUTPUT_PATH / f"taskbroker_0_{curr_time}_test_heartbeat.log" + with open(log_path, "a") as log_file: + process = subprocess.Popen( + [str(TASKBROKER_BIN), "-c", str(config_path)], + stderr=subprocess.STDOUT, + stdout=log_file, + ) + try: + time.sleep(3) # Let the broker create the sqlite file before we open it. + end = time.time() + 60 + while time.time() < end: + if get_num_tasks_in_sqlite(config) == 2: + break + time.sleep(1) + else: + raise AssertionError("Broker did not write both tasks to sqlite") + + worker = TaskWorkerClient(host) + renewed = worker.get_task() + lapsed = worker.get_task() + assert renewed is not None and lapsed is not None + + # Renew one task each second (like the worker's heartbeat process) for a few + # heartbeat periods, which is long enough for upkeep to reset the other one. + end = time.time() + HEARTBEAT_SEC * 4 + while time.time() < end: + assert worker.update_heartbeats([renewed.id]) == [] + time.sleep(1) + + assert task_state(config, renewed.id) == ("Processing", 0) + assert task_state(config, lapsed.id) == ("Pending", 1) + assert worker.update_heartbeats([renewed.id, lapsed.id]) == [lapsed.id] + finally: + process.send_signal(signal.SIGINT) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() diff --git a/integration_tests/integration_tests/worker.py b/integration_tests/integration_tests/worker.py index 63125c5f..f1c80850 100644 --- a/integration_tests/integration_tests/worker.py +++ b/integration_tests/integration_tests/worker.py @@ -7,6 +7,7 @@ TASK_ACTIVATION_STATUS_COMPLETE, TASK_ACTIVATION_STATUS_FAILURE, TASK_ACTIVATION_STATUS_RETRY, + BatchUpdateHeartbeatRequest, FetchNextTask, GetTaskRequest, SetTaskStatusRequest, @@ -70,6 +71,11 @@ def update_task( return response.task return None + def update_heartbeats(self, ids: list[str]) -> list[str]: + """Renew the heartbeat of processing activations. Returns the ids that are not processing.""" + response = self._stub.BatchUpdateHeartbeat(BatchUpdateHeartbeatRequest(ids=ids)) + return list(response.not_found_ids) + class ConfigurableTaskWorker: """