From 84c32eb55ebc022094af1370aa60ffcc55edb59c Mon Sep 17 00:00:00 2001 From: "sentry-junior[bot]" <264270552+sentry-junior[bot]@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:02:45 +0000 Subject: [PATCH] feat(client): Add completion hooks and detect masked processing deadlines The processing deadline is raised as a BaseException from SIGALRM and can interrupt library cleanup at any point. In Django, a deadline that lands in Atomic.__exit__ leaves the connection in manual transaction management and the driver raises its own error, which hides the deadline. The child then runs every later task on the broken connection. Treat an exception whose chain contains the deadline as a deadline, and add an optional TaskCompletionHook that runs after every task once the alarm is disarmed so applications can reset process state the deadline broke. Co-Authored-By: Rudra Patel --- clients/python/src/examples/tasks.py | 13 + clients/python/src/taskbroker_client/app.py | 11 +- clients/python/src/taskbroker_client/types.py | 25 ++ .../taskbroker_client/worker/workerchild.py | 233 ++++++++++++------ clients/python/tests/worker/test_worker.py | 154 +++++++++++- 5 files changed, 358 insertions(+), 78 deletions(-) diff --git a/clients/python/src/examples/tasks.py b/clients/python/src/examples/tasks.py index c3474660..fd5a144d 100644 --- a/clients/python/src/examples/tasks.py +++ b/clients/python/src/examples/tasks.py @@ -114,6 +114,19 @@ def will_retry_on_deadline_exceeded() -> None: timed_task(sleep_seconds=2) +@exampletasks.register( + name="examples.will_mask_deadline_exceeded", + processing_deadline_duration=1, +) +def will_mask_deadline_exceeded() -> None: + # Mimics a database driver that raises its own error from cleanup code + # while the deadline is unwinding the task. + try: + timed_task(sleep_seconds=2) + finally: + raise RuntimeError("cleanup failed while the deadline unwound") + + @exampletasks.register(name="examples.task_with_headers", pass_headers=True) def task_with_headers(value: str, headers: dict[str, str]) -> None: redis = StrictRedis(host="localhost", port=6379, decode_responses=True) diff --git a/clients/python/src/taskbroker_client/app.py b/clients/python/src/taskbroker_client/app.py index 6517a8f0..27cdb05f 100644 --- a/clients/python/src/taskbroker_client/app.py +++ b/clients/python/src/taskbroker_client/app.py @@ -13,7 +13,12 @@ from taskbroker_client.retry import Retry from taskbroker_client.router import TaskRouter from taskbroker_client.task import Task -from taskbroker_client.types import AtMostOnceStore, ContextHook, ProducerFactory +from taskbroker_client.types import ( + AtMostOnceStore, + ContextHook, + ProducerFactory, + TaskCompletionHook, +) class TaskbrokerApp: @@ -29,10 +34,14 @@ def __init__( metrics_class: str | MetricsBackend = "taskbroker_client.metrics.NoOpMetricsBackend", at_most_once_store: AtMostOnceStore | None = None, context_hooks: list[ContextHook] | None = None, + completion_hooks: list[TaskCompletionHook] | None = None, ) -> None: self.name = name self.metrics = self._build_metrics(metrics_class) self.context_hooks: list[ContextHook] = context_hooks if context_hooks is not None else [] + self.completion_hooks: list[TaskCompletionHook] = ( + completion_hooks if completion_hooks is not None else [] + ) self._config = { "rpc_secret": None, "grpc_config": None, diff --git a/clients/python/src/taskbroker_client/types.py b/clients/python/src/taskbroker_client/types.py index 1388bc1a..da449c57 100644 --- a/clients/python/src/taskbroker_client/types.py +++ b/clients/python/src/taskbroker_client/types.py @@ -27,6 +27,31 @@ def on_dispatch(self, headers: MutableMapping[str, Any]) -> None: ... def on_execute(self, headers: dict[str, str]) -> contextlib.AbstractContextManager[None]: ... +class TaskCompletionHook(Protocol): + """ + Hook called in the worker child after every activation, once the + processing deadline alarm has been disarmed. + + The deadline is enforced by raising a BaseException from a signal + handler, so it can interrupt library code at an arbitrary point and leave + process-wide state (database connections, client sessions) broken for + the next task on the same child. This hook is where an application can + check that state and reset it before the next task runs. + + deadline_exceeded is True when the processing deadline fired during the + task, including when another exception replaced it. Hooks should not rely + on it alone: a deadline can be swallowed entirely by code that catches + BaseException. + """ + + def after_task( + self, + activation: TaskActivation, + status: TaskActivationStatus.ValueType, + deadline_exceeded: bool, + ) -> None: ... + + class AtMostOnceStore(Protocol): """ Interface for the at_most_once store used for idempotent task execution. diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index 7ca5ee9e..ce70659b 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -49,6 +49,26 @@ class ProcessingDeadlineExceeded(BaseException): pass +def find_deadline_in_chain(err: BaseException) -> ProcessingDeadlineExceeded | None: + """ + Return the ProcessingDeadlineExceeded that `err` was raised from or while + handling, if any. + """ + seen: set[int] = set() + pending: list[BaseException] = [err] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, ProcessingDeadlineExceeded): + return current + # `raise X from Y` sets __cause__ but the deadline can still sit on + # __context__, so follow both links. + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + return None + + @dataclass class ActivationWithPendingFutures: """ @@ -265,6 +285,130 @@ def handle_alarm(signum: int, frame: FrameType | None) -> None: f"execution deadline of {deadline} seconds exceeded by {taskname}" ) + def handle_deadline_exceeded( + inflight: InflightTaskActivation, + task_func: Task[Any, Any], + err: BaseException, + deadline: ProcessingDeadlineExceeded, + ) -> TaskActivationStatus.ValueType: + """ + Report a processing deadline and decide the next state. `err` is the + exception that reached the worker, which is `deadline` itself unless + another exception replaced it while the task unwound. + """ + if task_func.report_timeout_errors: + with sentry_sdk.isolation_scope() as scope: + scope.fingerprint = [ + "taskworker.processing_deadline_exceeded", + inflight.activation.namespace, + inflight.activation.taskname, + ] + # Override transaction names inherited from worker startup. + sentry_sdk.get_current_scope().set_transaction_name( + inflight.activation.taskname + ) + _log_task_failed(inflight.activation, err, processing_pool_name) + metrics.incr( + "taskworker.worker.processing_deadline_exceeded", + tags={ + "processing_pool": processing_pool_name, + "namespace": inflight.activation.namespace, + "taskname": inflight.activation.taskname, + "masked": str(err is not deadline).lower(), + }, + ) + retry = task_func.retry + if retry and retry.should_retry(inflight.activation.retry_state, deadline): + return TASK_ACTIVATION_STATUS_RETRY + return TASK_ACTIVATION_STATUS_FAILURE + + def handle_task_exception( + inflight: InflightTaskActivation, + task_func: Task[Any, Any], + err: Exception, + ) -> TaskActivationStatus.ValueType: + next_state = TASK_ACTIVATION_STATUS_FAILURE + retry = task_func.retry + captured_error = False + should_capture_error = not isinstance(err, task_func.silenced_exceptions) + if retry: + if retry.should_retry(inflight.activation.retry_state, err): + logger.info( + "taskworker.task.retry", + extra={ + "namespace": inflight.activation.namespace, + "taskname": inflight.activation.taskname, + "processing_pool": processing_pool_name, + "error": str(err), + }, + ) + next_state = TASK_ACTIVATION_STATUS_RETRY + elif retry.max_attempts_reached(inflight.activation.retry_state): + with sentry_sdk.isolation_scope() as scope: + # Only report to Sentry when the underlying error is not silenced. + if should_capture_error: + retry_error = NoRetriesRemainingError( + f"{inflight.activation.taskname} has consumed all of its retries" + ) + retry_error.__cause__ = err + scope.fingerprint = [ + "taskworker.no_retries_remaining", + inflight.activation.namespace, + inflight.activation.taskname, + ] + sentry_sdk.get_current_scope().set_transaction_name( + inflight.activation.taskname + ) + sentry_sdk.capture_exception(retry_error) + # Emit a structured worker log without logger.exception so this + # branch does not create a second Sentry error event. + _log_task_retry_exhausted(inflight.activation, err, processing_pool_name) + # In this branch, all exceptions should be either + # captured or silenced. + captured_error = True + + if ( + should_capture_error + and not captured_error + and next_state != TASK_ACTIVATION_STATUS_RETRY + ): + with sentry_sdk.isolation_scope(): + sentry_sdk.get_current_scope().set_transaction_name( + inflight.activation.taskname + ) + _log_task_failed(inflight.activation, err, processing_pool_name) + + return next_state + + def run_completion_hooks( + activation: TaskActivation, + status: TaskActivationStatus.ValueType, + deadline_exceeded: bool, + ) -> None: + """ + Runs after timeout_alarm has exited, so the deadline cannot interrupt + a hook that is repairing state the deadline broke. + """ + for hook in app.completion_hooks: + try: + hook.after_task(activation, status, deadline_exceeded) + except Exception: + logger.exception( + "taskworker.worker.completion_hook_failed", + extra={ + "hook": type(hook).__name__, + "taskname": activation.taskname, + "processing_pool": processing_pool_name, + }, + ) + metrics.incr( + "taskworker.worker.completion_hook_failed", + tags={ + "hook": type(hook).__name__, + "processing_pool": processing_pool_name, + }, + ) + def await_task_futures(task: ActivationWithPendingFutures) -> None: """ Blocks on the result of each producer futures in the given task. @@ -501,6 +645,7 @@ def check_task_future_completion( set_current_task(inflight.activation) next_state = TASK_ACTIVATION_STATUS_FAILURE + deadline_exceeded = False # Use time.time() so we can measure against activation.received_at execution_start_time = time.time() try: @@ -508,86 +653,22 @@ def check_task_future_completion( _execute_activation(task_func, inflight.activation, app.context_hooks) next_state = TASK_ACTIVATION_STATUS_COMPLETE except ProcessingDeadlineExceeded as err: - if task_func.report_timeout_errors: - with sentry_sdk.isolation_scope() as scope: - scope.fingerprint = [ - "taskworker.processing_deadline_exceeded", - inflight.activation.namespace, - inflight.activation.taskname, - ] - # Override transaction names inherited from worker startup. - sentry_sdk.get_current_scope().set_transaction_name( - inflight.activation.taskname - ) - _log_task_failed(inflight.activation, err, processing_pool_name) - metrics.incr( - "taskworker.worker.processing_deadline_exceeded", - tags={ - "processing_pool": processing_pool_name, - "namespace": inflight.activation.namespace, - "taskname": inflight.activation.taskname, - }, - ) - retry = task_func.retry - if retry and retry.should_retry(inflight.activation.retry_state, err): - next_state = TASK_ACTIVATION_STATUS_RETRY - else: - next_state = TASK_ACTIVATION_STATUS_FAILURE + deadline_exceeded = True + next_state = handle_deadline_exceeded(inflight, task_func, err, err) except Exception as err: - retry = task_func.retry - captured_error = False - should_capture_error = not isinstance(err, task_func.silenced_exceptions) - if retry: - if retry.should_retry(inflight.activation.retry_state, err): - logger.info( - "taskworker.task.retry", - extra={ - "namespace": inflight.activation.namespace, - "taskname": inflight.activation.taskname, - "processing_pool": processing_pool_name, - "error": str(err), - }, - ) - next_state = TASK_ACTIVATION_STATUS_RETRY - elif retry.max_attempts_reached(inflight.activation.retry_state): - with sentry_sdk.isolation_scope() as scope: - # Only report to Sentry when the underlying error is not silenced. - if should_capture_error: - retry_error = NoRetriesRemainingError( - f"{inflight.activation.taskname} has consumed all of its retries" - ) - retry_error.__cause__ = err - scope.fingerprint = [ - "taskworker.no_retries_remaining", - inflight.activation.namespace, - inflight.activation.taskname, - ] - sentry_sdk.get_current_scope().set_transaction_name( - inflight.activation.taskname - ) - sentry_sdk.capture_exception(retry_error) - # Emit a structured worker log without logger.exception so this - # branch does not create a second Sentry error event. - _log_task_retry_exhausted( - inflight.activation, err, processing_pool_name - ) - # In this branch, all exceptions should be either - # captured or silenced. - captured_error = True - - if ( - should_capture_error - and not captured_error - and next_state != TASK_ACTIVATION_STATUS_RETRY - ): - with sentry_sdk.isolation_scope(): - sentry_sdk.get_current_scope().set_transaction_name( - inflight.activation.taskname - ) - _log_task_failed(inflight.activation, err, processing_pool_name) + # Cleanup code that runs while the deadline unwinds (for example a + # database driver refusing to reset a connection) can raise a new + # exception that replaces it. Treat that as the deadline it is. + masked_deadline = find_deadline_in_chain(err) + if masked_deadline is not None: + deadline_exceeded = True + next_state = handle_deadline_exceeded(inflight, task_func, err, masked_deadline) + else: + next_state = handle_task_exception(inflight, task_func, err) clear_current_task() processed_task_count += 1 + run_completion_hooks(inflight.activation, next_state, deadline_exceeded) # To have Taskworker track futures, set the env var `ARROYO_TRACK_PRODUCER_FUTURES = True` # in the worker process diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 56079841..0c88ff8c 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -40,6 +40,7 @@ PushTaskResponse, RetryState, TaskActivation, + TaskActivationStatus, ) from sentry_sdk.crons import MonitorStatus from sentry_sdk.integrations.logging import LoggingIntegration @@ -73,8 +74,9 @@ TrackedChild, WorkerServicer, ) -from taskbroker_client.worker.workerchild import ChildMessage +from taskbroker_client.worker.workerchild import ChildMessage, ProcessingDeadlineExceeded from taskbroker_client.worker.workerchild import child_process as _child_process +from taskbroker_client.worker.workerchild import find_deadline_in_chain SIMPLE_TASK = InflightTaskActivation( host="localhost:50051", @@ -262,6 +264,18 @@ ), ) +MASKED_DEADLINE_TASK = InflightTaskActivation( + host="localhost:50051", + receive_timestamp=0, + activation=TaskActivation( + id="masked-deadline", + taskname="examples.will_mask_deadline_exceeded", + namespace="examples", + parameters_bytes=msgpack.packb({"args": [], "kwargs": {}}, use_bin_type=True), + processing_deadline_duration=1, + ), +) + # Task set to retry on deadline exceeded exceptions RETRY_TASK_ON_DEADLINE_EXCEEDED = InflightTaskActivation( host="localhost:50051", @@ -3830,3 +3844,141 @@ def test_stale_drain_file_cleared_on_startup(tmp_path: Path) -> None: assert not drain_file.exists() assert pool._should_publish_occupancy() + + +def _run_child(*activations: InflightTaskActivation) -> list[ProcessingResult]: + todo: queue.Queue[InflightTaskActivation] = queue.Queue() + processed: queue.Queue[ProcessingResult] = queue.Queue() + for activation in activations: + todo.put(activation) + child_process( + "examples.app:app", + todo, + processed, + Event(), + max_task_count=len(activations), + processing_pool_name="test", + process_type="fork", + skip_awaiting_futures=False, + future_checking_frequency=0.1, + ) + return [processed.get() for _ in activations] + + +def test_find_deadline_in_chain() -> None: + deadline = ProcessingDeadlineExceeded("deadline") + assert find_deadline_in_chain(deadline) is deadline + assert find_deadline_in_chain(ValueError("unrelated")) is None + + # psycopg2-style: the driver error is raised while handling the deadline, + # then wrapped with `raise ... from` by the ORM. + driver_error = RuntimeError("set_session cannot be used inside a transaction") + driver_error.__context__ = deadline + wrapped = ValueError("wrapped") + wrapped.__cause__ = driver_error + wrapped.__context__ = driver_error + assert find_deadline_in_chain(wrapped) is deadline + + # Explicit cause that hides the deadline on __context__. + other = KeyError("other") + hidden = ValueError("hidden") + hidden.__cause__ = other + hidden.__context__ = deadline + assert find_deadline_in_chain(hidden) is deadline + + cyclic_a = ValueError("a") + cyclic_b = ValueError("b") + cyclic_a.__context__ = cyclic_b + cyclic_b.__context__ = cyclic_a + assert find_deadline_in_chain(cyclic_a) is None + + +@mock.patch("taskbroker_client.worker.workerchild.logger") +def test_child_process_reports_masked_deadline(mock_logger: mock.Mock) -> None: + """A deadline replaced by another exception is still handled as a deadline.""" + from examples.app import app + + with mock.patch.object(app.metrics, "incr") as mock_incr: + (result,) = _run_child(MASKED_DEADLINE_TASK) + + assert result.status == TASK_ACTIVATION_STATUS_FAILURE + failed_calls = [ + c + for c in mock_logger.exception.call_args_list + if c.args and c.args[0] == "taskworker.task.failed" + ] + assert len(failed_calls) == 1 + assert failed_calls[0].kwargs["extra"]["exception_type"] == "RuntimeError" + deadline_calls = [ + c + for c in mock_incr.call_args_list + if c.args and c.args[0] == "taskworker.worker.processing_deadline_exceeded" + ] + assert len(deadline_calls) == 1 + assert deadline_calls[0].kwargs["tags"]["masked"] == "true" + + +def test_child_process_completion_hooks() -> None: + """Completion hooks run after every task, outside the deadline alarm.""" + calls: list[tuple[str, int, bool, float]] = [] + + class RecordingHook: + def after_task( + self, + activation: TaskActivation, + status: TaskActivationStatus.ValueType, + deadline_exceeded: bool, + ) -> None: + remaining, _ = signal.getitimer(signal.ITIMER_REAL) + calls.append((activation.taskname, status, deadline_exceeded, remaining)) + + from examples.app import app + + hook = RecordingHook() + app.completion_hooks.append(hook) + try: + results = _run_child(SIMPLE_TASK, FAIL_TASK, MASKED_DEADLINE_TASK) + finally: + app.completion_hooks.remove(hook) + + assert [r.status for r in results] == [ + TASK_ACTIVATION_STATUS_COMPLETE, + TASK_ACTIVATION_STATUS_FAILURE, + TASK_ACTIVATION_STATUS_FAILURE, + ] + assert calls == [ + ("examples.simple_task", TASK_ACTIVATION_STATUS_COMPLETE, False, 0.0), + ("examples.fail_task", TASK_ACTIVATION_STATUS_FAILURE, False, 0.0), + ("examples.will_mask_deadline_exceeded", TASK_ACTIVATION_STATUS_FAILURE, True, 0.0), + ] + + +@mock.patch("taskbroker_client.worker.workerchild.logger") +def test_child_process_completion_hook_failure_does_not_stop_child( + mock_logger: mock.Mock, +) -> None: + class BrokenHook: + def after_task( + self, + activation: TaskActivation, + status: TaskActivationStatus.ValueType, + deadline_exceeded: bool, + ) -> None: + raise RuntimeError("hook failed") + + from examples.app import app + + hook = BrokenHook() + app.completion_hooks.append(hook) + try: + results = _run_child(SIMPLE_TASK, SIMPLE_TASK) + finally: + app.completion_hooks.remove(hook) + + assert [r.status for r in results] == [TASK_ACTIVATION_STATUS_COMPLETE] * 2 + hook_failures = [ + c + for c in mock_logger.exception.call_args_list + if c.args and c.args[0] == "taskworker.worker.completion_hook_failed" + ] + assert len(hook_failures) == 2