Repository: getsentry/sentry
Pull request: fix(taskworker) remove signal handler safely
There are some tests for `child_worker` that are not in child processes, and we're currently leaking signal handlers. Use a context manager to add/cancel the signal handler.

Fixes #88742

Problem statement:
Fix signal leak in taskworker child_worker
```diff
diff --git a/tests/conftest.py b/tests/conftest.py
index 74ff8dbb8de..788cf0a6a9e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,7 @@
 import os
+import signal
 import sys
-from collections.abc import MutableMapping
+from collections.abc import Generator, MutableMapping
 
 import psutil
 import pytest
@@ -146,3 +147,29 @@ def check_leaked_responses_mocks():
             f"`responses` were leaked outside of the test context:\n{leaked_s}"
             f"(make sure to use `@responses.activate` or `with responses.mock:`)"
         )
+
+
+def _leaked_signals() -> list[str]:
+    leaked = []
+    for signum in signal.Signals:
+        got = signal.getsignal(signum)
+        if (
+            got is not None
+            and not isinstance(got, int)
+            and got is not signal.default_int_handler
+            # prevent the debugger from interrupting teardown
+            and "Pdb.sigint_handler" not in str(got)
+        ):
+            leaked.append(f"- {signum}: {got}")
+    return leaked
+
+
+@pytest.fixture(autouse=True)
+def check_leaked_signals() -> Generator[None]:
+    before = frozenset(_leaked_signals())
+    yield
+    after = _leaked_signals()
+
+    leaked = [msg for msg in after if msg not in before]
+    if leaked:
+        raise AssertionError(f"test leaked os signal handlers:\n{'\n'.join(leaked)}")
```

when applied can result in tests failing

```
$ pytest tests/sentry/taskworker/test_worker.py::test_child_worker_failure_task
============================= test session starts ==============================
platform darwin -- Python 3.13.1, pytest-8.1.2, pluggy-1.5.0
django: version: 5.1.7
rootdir: /Users/asottile/workspace/sentry
configfile: pyproject.toml
plugins: fail-slow-0.3.0, repeat-0.9.3, time-machine-2.16.0, json-report-1.5.0, metadata-3.1.1, xdist-3.0.2, django-4.9.0, pytest_sentry-0.3.0, anyio-3.7.1, rerunfailures-15.0, cov-4.0.0
collected 1 item                                                               

tests/sentry/taskworker/test_worker.py .E                                [100%]

==================================== ERRORS ====================================
_____________ ERROR at teardown of test_child_worker_failure_task ______________
tests/conftest.py:175: in check_leaked_signals
    raise AssertionError(f"test leaked os signal handlers:\n{'\n'.join(leaked)}")
E   AssertionError: test leaked os signal handlers:
E   - 14: <function child_worker.<locals>.handle_alarm at 0x12d0c65c0>
=========================== short test summary info ============================
ERROR tests/sentry/taskworker/test_worker.py::test_child_worker_failure_task - AssertionError: test leaked os signal handlers:
========================== 1 passed, 1 error in 4.56s ==========================
````

Hints from issue discussion:

Commit message: fix(taskworker) remove signal handler safely

There are some tests for `child_worker` that are not in child processes,
and we're currently leaking signal handlers. Use a context manager to
add/cancel the signal handler.

Fixes #88742

Patch to review:
diff --git a/src/sentry/taskworker/worker.py b/src/sentry/taskworker/worker.py
index 1daec357e068f1..9b8fea1b883878 100644
--- a/src/sentry/taskworker/worker.py
+++ b/src/sentry/taskworker/worker.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import atexit
+import contextlib
 import dataclasses
 import logging
 import multiprocessing
@@ -9,6 +10,7 @@
 import sys
 import threading
 import time
+from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 from multiprocessing.context import ForkProcess
 from multiprocessing.synchronize import Event
@@ -39,6 +41,8 @@
 mp_context = multiprocessing.get_context("fork")
 logger = logging.getLogger("sentry.taskworker.worker")
 
+AT_MOST_ONCE_TIMEOUT = 60 * 60 * 24  # 1 day
+
 
 @dataclasses.dataclass
 class ProcessingResult:
@@ -48,7 +52,22 @@ class ProcessingResult:
     status: TaskActivationStatus.ValueType
 
 
-AT_MOST_ONCE_TIMEOUT = 60 * 60 * 24  # 1 day
+@contextlib.contextmanager
+def timeout_alarm(seconds: int, handler: Callable[[int, Any], None]):
+    """
+    Context manager to handle SIGALRM handlers
+
+    To prevent tasks from consuming a worker forever, we set a timeout
+    alarm that will interrupt tasks that run longer than
+    their processing_deadline.
+    """
+    original = signal.signal(signal.SIGALRM, handler)
+    try:
+        signal.alarm(seconds)
+        yield
+    finally:
+        signal.alarm(0)
+        signal.signal(signal.SIGALRM, original)
 
 
 def get_at_most_once_key(namespace: str, taskname: str, task_id: str) -> str:
@@ -106,8 +125,6 @@ def handle_alarm(signum: int, frame: Any) -> None:
             )
         sys.exit(1)
 
-    signal.signal(signal.SIGALRM, handle_alarm)
-
     while True:
         if max_task_count and processed_task_count >= max_task_count:
             metrics.incr(
@@ -152,16 +169,12 @@ def handle_alarm(signum: int, frame: Any) -> None:
 
         set_current_task(activation)
 
-        # Set an alarm for the processing_deadline_duration
-        signal.alarm(activation.processing_deadline_duration)
-
-        execution_start_time = time.time()
         next_state = TASK_ACTIVATION_STATUS_FAILURE
+        execution_start_time = time.time()
         try:
-            _execute_activation(task_func, activation)
+            with timeout_alarm(activation.processing_deadline_duration, handle_alarm):
+                _execute_activation(task_func, activation)
             next_state = TASK_ACTIVATION_STATUS_COMPLETE
-            # Clear the alarm
-            signal.alarm(0)
         except Exception as err:
             if task_func.should_retry(activation.retry_state, err):
                 logger.info("taskworker.task.retry", extra={"taskname": activation.taskname})


Code review comments:
diff_hunk: @@ -152,16 +169,12 @@ def handle_alarm(signum: int, frame: Any) -> None:
 
         set_current_task(activation)
 
-        # Set an alarm for the processing_deadline_duration
-        signal.alarm(activation.processing_deadline_duration)
-
-        execution_start_time = time.time()
         next_state = TASK_ACTIVATION_STATUS_FAILURE
+        execution_start_time = time.time()
path: src/sentry/taskworker/worker.py
text: @user1:
`time.time` should really never be used for timing as it is subjected to smear / etc.

you probably want `time.monotonic()` https://youtu.be/hb-12mgQrzM

@author:
While I can switch a few usage of `time.time()`, I don't think all of them can be changed. We need to do some time differences against the message produce time which is a protobuf timestamp. These won't be precise measurements as clock times can vary from one system to another, but there aren't many other portable time representations.

diff_hunk: @@ -48,7 +52,22 @@ class ProcessingResult:
     status: TaskActivationStatus.ValueType
 
 
-AT_MOST_ONCE_TIMEOUT = 60 * 60 * 24  # 1 day
+@contextlib.contextmanager
+def timeout_alarm(seconds: int, handler: Callable[[int, Any], None]):
path: src/sentry/taskworker/worker.py
text: @user1:
since the actual handler doesn't care about either of its arguments it can be `Callable[[object, object], None]` and this should `-> Generator[None]`