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