Skip to content

[WIP] feat(client): Add completion hooks and detect masked processing deadlines #840

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
sentry-junior wants to merge 1 commit into from
Closed

[WIP] feat(client): Add completion hooks and detect masked processing deadlines #840

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions clients/python/src/examples/tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
11 changes: 10 additions & 1 deletion clients/python/src/taskbroker_client/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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,
Expand Down
25 changes: 25 additions & 0 deletions clients/python/src/taskbroker_client/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
233 changes: 157 additions & 76 deletions clients/python/src/taskbroker_client/worker/workerchild.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -501,93 +645,30 @@ 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:
with timeout_alarm(inflight.activation.processing_deadline_duration, handle_alarm):
_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
Expand Down
Loading
Loading