Repository: getsentry/sentry
Pull request: feat(taskworker) Add concurrent worker
Move the taskworker process to be a multiprocess concurrent worker. This will help enable higher CPU usage in worker pods, as we can pack more concurrent CPU operations into each pod (at the cost of memory).

The main process is responsible for:

- Spawning children
- Making RPC requests to fill child queues and submit results.

Each child process handles:

- Resolving task names
- Checking at_most_once keys
- Enforcing processing deadlines
- Executing task functions

Instead of using more child processes to enforce timeouts, I've used SIGALRM. I've verified that tasks like 

```python
@exampletasks.register(name="examples.infinite", retry=Retry(times=2))
def infinite_task() -> None:
    try:
        while True:
            pass
    except Exception as e:
        print("haha caught exception", e)
```

Do not paralyze workers with infinite loops. 

When a worker is terminated, it uses an `Event` to have children exit, and then drains any results. If there are tasks in the `_child_tasks` queue will not be completed, and instead will sent to another worker when the `processing_deadline` on the activations expires.

Problem statement:
Make taskworker worker support concurrent work
The current worker implementation is single threaded and will have limited throughput. Expand the worker to support concurrent processing based on CLI flags. `--concurrency=12` should create 12 worker threads/processes allowing a worker to process up to 12 tasks concurrently.

Rough thinking would be to have a single thread responsible for fetching tasks and sending rpc responses back and a pool of worker threads that execute tasks. Processing timeouts could be tricky to enforce though as python threads don't have a way to do that. Perhaps `celery worker` has some tricks we could replicate.

Hints from issue discussion:

Commit message: More coverage for fetching next task

Patch to review:
diff --git a/src/sentry/runner/commands/run.py b/src/sentry/runner/commands/run.py
index 193a5ec171eb12..6a15b17673ed76 100644
--- a/src/sentry/runner/commands/run.py
+++ b/src/sentry/runner/commands/run.py
@@ -242,6 +242,7 @@ def worker(ignore_unknown_queues: bool, **options: Any) -> None:
 @click.option(
     "--max-task-count", help="Number of tasks this worker should run before exiting", default=10000
 )
+@click.option("--concurrency", help="Number of child worker processes to create.", default=1)
 @click.option(
     "--namespace", help="The dedicated task namespace that taskworker operates on", default=None
 )
@@ -258,7 +259,7 @@ def taskworker(**options: Any) -> None:
 
 
 def run_taskworker(
-    rpc_host: str, max_task_count: int, namespace: str | None, **options: Any
+    rpc_host: str, max_task_count: int, namespace: str | None, concurrency: int, **options: Any
 ) -> None:
     """
     taskworker factory that can be reloaded
@@ -267,7 +268,11 @@ def run_taskworker(
 
     with managed_bgtasks(role="taskworker"):
         worker = TaskWorker(
-            rpc_host=rpc_host, max_task_count=max_task_count, namespace=namespace, **options
+            rpc_host=rpc_host,
+            max_task_count=max_task_count,
+            namespace=namespace,
+            concurrency=concurrency,
+            **options,
         )
         exitcode = worker.start()
         raise SystemExit(exitcode)
diff --git a/src/sentry/taskworker/client.py b/src/sentry/taskworker/client.py
index 6beb6886b30e51..b185885df71846 100644
--- a/src/sentry/taskworker/client.py
+++ b/src/sentry/taskworker/client.py
@@ -44,7 +44,10 @@ def get_task(self, namespace: str | None = None) -> TaskActivation | None:
         return None
 
     def update_task(
-        self, task_id: str, status: TaskActivationStatus.ValueType, fetch_next_task: FetchNextTask
+        self,
+        task_id: str,
+        status: TaskActivationStatus.ValueType,
+        fetch_next_task: FetchNextTask | None = None,
     ) -> TaskActivation | None:
         """
         Update the status for a given task activation.
diff --git a/src/sentry/taskworker/worker.py b/src/sentry/taskworker/worker.py
index 5153dafe6de5b2..c4571283ab6afc 100644
--- a/src/sentry/taskworker/worker.py
+++ b/src/sentry/taskworker/worker.py
@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+import dataclasses
 import logging
 import multiprocessing
+import queue
+import signal
+import sys
 import time
-from multiprocessing.context import TimeoutError
-from multiprocessing.pool import Pool
+from multiprocessing.synchronize import Event
+from types import FrameType
 from typing import Any
 from uuid import uuid4
 
@@ -19,6 +23,7 @@
     TASK_ACTIVATION_STATUS_RETRY,
     FetchNextTask,
     TaskActivation,
+    TaskActivationStatus,
 )
 
 from sentry.taskworker.client import TaskworkerClient
@@ -29,32 +34,13 @@
 
 logger = logging.getLogger("sentry.taskworker.worker")
 
-mp_context = multiprocessing.get_context("spawn")
 
+@dataclasses.dataclass
+class ProcessingResult:
+    """Result structure from child processess to parent"""
 
-def _init_pool_process() -> None:
-    """initialize pool workers by loading all task modules"""
-    for module in settings.TASKWORKER_IMPORTS:
-        __import__(module)
-
-
-def _process_activation(activation: TaskActivation) -> None:
-    """multiprocess worker method"""
-    parameters = orjson.loads(activation.parameters)
-    args = parameters.get("args", [])
-    kwargs = parameters.get("kwargs", {})
-    headers = {k: v for k, v in activation.headers.items()}
-
-    transaction = sentry_sdk.continue_trace(
-        environ_or_headers=headers,
-        op="task.taskworker",
-        name=f"{activation.namespace}:{activation.taskname}",
-    )
-    with (
-        track_memory_usage("taskworker.worker.memory_change"),
-        sentry_sdk.start_transaction(transaction),
-    ):
-        taskregistry.get(activation.namespace).get(activation.taskname)(*args, **kwargs)
+    task_id: str
+    status: TaskActivationStatus.ValueType
 
 
 AT_MOST_ONCE_TIMEOUT = 60 * 60 * 24  # 1 day
@@ -65,137 +51,84 @@ def get_at_most_once_key(namespace: str, taskname: str, task_id: str) -> str:
     return f"tw:amo:{namespace}:{taskname}:{task_id}"
 
 
-class TaskWorker:
-    """
-    A TaskWorker fetches tasks from a taskworker RPC host and handles executing task activations.
-
-    Tasks are executed in a forked process so that processing timeouts can be enforced.
-    As tasks are completed status changes will be sent back to the RPC host and new tasks
-    will be fetched.
-
-    Taskworkers can be run with `sentry run taskworker`
-    """
+def _get_known_task(activation: TaskActivation) -> Task[Any, Any] | None:
+    if not taskregistry.contains(activation.namespace):
+        logger.error(
+            "taskworker.invalid_namespace",
+            extra={"namespace": activation.namespace, "taskname": activation.taskname},
+        )
+        return None
 
-    def __init__(
-        self,
-        rpc_host: str,
-        max_task_count: int | None = None,
-        namespace: str | None = None,
-        **options: dict[str, Any],
-    ) -> None:
-        self.options = options
-        self._execution_count = 0
-        self._worker_id = uuid4().hex
-        self._max_task_count = max_task_count
-        self._namespace = namespace
-        self.client = TaskworkerClient(rpc_host)
-        self._pool: Pool | None = None
-        self._build_pool()
+    namespace = taskregistry.get(activation.namespace)
+    if not namespace.contains(activation.taskname):
+        logger.error(
+            "taskworker.invalid_taskname",
+            extra={"namespace": activation.namespace, "taskname": activation.taskname},
+        )
+        return None
+    return namespace.get(activation.taskname)
 
-    def __del__(self) -> None:
-        if self._pool:
-            self._pool.terminate()
 
-    def _build_pool(self) -> None:
-        if self._pool:
-            self._pool.terminate()
-        self._pool = mp_context.Pool(processes=1, initializer=_init_pool_process)
+def child_worker(
+    child_tasks: queue.Queue[TaskActivation],
+    processed_tasks: queue.Queue[ProcessingResult],
+    shutdown_event: Event,
+    max_task_count: int | None,
+) -> None:
+    for module in settings.TASKWORKER_IMPORTS:
+        __import__(module)
 
-    def do_imports(self) -> None:
-        for module in settings.TASKWORKER_IMPORTS:
-            __import__(module)
-        self._build_pool()
+    current_task_id: str | None = None
+    processed_task_count = 0
 
-    def start(self) -> int:
+    def handle_alarm(signum: int, frame: Any) -> None:
         """
-        Run the worker main loop
+        Handle SIGALRM
 
-        Once started a Worker will loop until it is killed, or
-        completes its max_task_count when it shuts down.
+        If we hit an alarm in a child, we need to push a result
+        and terminate the child.
         """
-        self.do_imports()
-        next_task: TaskActivation | None = None
-        task: TaskActivation | None = None
-        try:
-            while True:
-                if next_task:
-                    task = next_task
-                    next_task = None
-                else:
-                    task = self.fetch_task()
-
-                if not task:
-                    metrics.incr("taskworker.worker.no_task.pause")
-                    time.sleep(1)
-                    continue
-
-                next_task = self.process_task(task)
-                if (
-                    self._max_task_count is not None
-                    and self._max_task_count <= self._execution_count
-                ):
-                    metrics.incr(
-                        "taskworker.worker.max_task_count_reached",
-                        tags={"count": self._execution_count},
-                    )
-                    logger.info("Max task execution count reached. Terminating")
-                    return 0
-
-        except KeyboardInterrupt:
-            return 1
-        except Exception:
-            logger.exception("Worker process crashed")
-            return 2
+        nonlocal current_task_id, processed_tasks
 
-    def fetch_task(self) -> TaskActivation | None:
-        try:
-            activation = self.client.get_task(self._namespace)
-        except grpc.RpcError:
-            metrics.incr("taskworker.worker.get_task.failed")
-            logger.info("get_task failed. Retrying in 1 second")
-            return None
+        if current_task_id:
+            processed_tasks.put(
+                ProcessingResult(task_id=current_task_id, status=TASK_ACTIVATION_STATUS_FAILURE)
+            )
+        metrics.incr("taskworker.worker.processing_deadline_exceeded")
+        sys.exit(1)
 
-        if not activation:
-            metrics.incr("taskworker.worker.get_task.not_found")
-            logger.debug("No task fetched")
-            return None
+    signal.signal(signal.SIGALRM, handle_alarm)
 
-        metrics.incr("taskworker.worker.get_task.success")
-        return activation
-
-    def _get_known_task(self, activation: TaskActivation) -> Task[Any, Any] | None:
-        if not taskregistry.contains(activation.namespace):
-            logger.error(
-                "taskworker.invalid_namespace",
-                extra={"namespace": activation.namespace, "taskname": activation.taskname},
+    while True:
+        if max_task_count and processed_task_count >= max_task_count:
+            metrics.incr(
+                "taskworker.worker.max_task_count_reached",
+                tags={"count": processed_task_count},
             )
-            return None
+            logger.info("taskwworker.max_task_count_reached")
+            break
 
-        namespace = taskregistry.get(activation.namespace)
-        if not namespace.contains(activation.taskname):
-            logger.error(
-                "taskworker.invalid_taskname",
-                extra={"namespace": activation.namespace, "taskname": activation.taskname},
-            )
-            return None
-        return namespace.get(activation.taskname)
+        if shutdown_event.is_set():
+            logger.info("taskworker.worker.shutdown_event")
+            break
 
-    def process_task(self, activation: TaskActivation) -> TaskActivation | None:
-        assert self._pool
-        task = self._get_known_task(activation)
-        if not task:
+        try:
+            activation = child_tasks.get_nowait()
+        except queue.Empty:
+            continue
+
+        task_func = _get_known_task(activation)
+        if not task_func:
             metrics.incr(
                 "taskworker.worker.unknown_task",
                 tags={"namespace": activation.namespace, "taskname": activation.taskname},
             )
-            self._execution_count += 1
-            return self.client.update_task(
-                task_id=activation.id,
-                status=TASK_ACTIVATION_STATUS_FAILURE,
-                fetch_next_task=FetchNextTask(namespace=self._namespace),
+            processed_tasks.put(
+                ProcessingResult(task_id=activation.id, status=TASK_ACTIVATION_STATUS_FAILURE)
             )
+            continue
 
-        if task.at_most_once:
+        if task_func.at_most_once:
             key = get_at_most_once_key(activation.namespace, activation.taskname, activation.id)
             if cache.add(key, "1", timeout=AT_MOST_ONCE_TIMEOUT):  # The key didn't exist
                 metrics.incr(
@@ -205,47 +138,35 @@ def process_task(self, activation: TaskActivation) -> TaskActivation | None:
                 metrics.incr(
                     "taskworker.worker.at_most_once.skipped", tags={"task": activation.taskname}
                 )
-                return None
+                continue
 
-        processing_timeout = activation.processing_deadline_duration
-        namespace = taskregistry.get(activation.namespace)
-        next_state = TASK_ACTIVATION_STATUS_FAILURE
-        result = None
-        execution_start_time = 0.0
-        try:
-            execution_start_time = time.time()
+        current_task_id = activation.id
 
-            result = self._pool.apply_async(
-                func=_process_activation,
-                args=(activation,),
-            )
-            # Will trigger a TimeoutError if the task execution runs long
-            result.get(timeout=processing_timeout)
+        # 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
+        try:
+            _execute_activation(task_func, activation)
             next_state = TASK_ACTIVATION_STATUS_COMPLETE
-        except TimeoutError:
-            logger.info(
-                "taskworker.task_execution_timeout",
-                extra={
-                    "taskname": activation.taskname,
-                    "processing_deadline": processing_timeout,
-                },
-            )
-            # When a task times out we kill the entire pool. This is necessary because
-            # stdlib multiprocessing.Pool doesn't expose ways to terminate individual tasks.
-            self._build_pool()
+            # Clear the alarm
+            signal.alarm(0)
         except Exception as err:
-            if namespace.get(activation.taskname).should_retry(activation.retry_state, err):
+            if task_func.should_retry(activation.retry_state, err):
                 logger.info("taskworker.task.retry", extra={"task": activation.taskname})
                 next_state = TASK_ACTIVATION_STATUS_RETRY
 
             if next_state != TASK_ACTIVATION_STATUS_RETRY:
                 logger.info(
-                    "taskworker.task_errored", extra={"type": str(err.__class__), "error": str(err)}
+                    "taskworker.task.errored", extra={"type": str(err.__class__), "error": str(err)}
                 )
 
+        processed_task_count += 1
+
+        # Get completion time before pushing to queue to avoid inflating latency metrics.
         execution_complete_time = time.time()
-        self._execution_count += 1
+        processed_tasks.put(ProcessingResult(task_id=activation.id, status=next_state))
 
         task_added_time = activation.received_at.ToDatetime().timestamp()
         execution_duration = execution_complete_time - execution_start_time
@@ -277,8 +198,189 @@ def process_task(self, activation: TaskActivation) -> TaskActivation | None:
             tags={"namespace": activation.namespace},
         )
 
-        return self.client.update_task(
-            task_id=activation.id,
-            status=next_state,
-            fetch_next_task=FetchNextTask(namespace=self._namespace),
+
+def _execute_activation(task_func: Task[Any, Any], activation: TaskActivation) -> None:
+    """Invoke a task function with the activation parameters."""
+    parameters = orjson.loads(activation.parameters)
+    args = parameters.get("args", [])
+    kwargs = parameters.get("kwargs", {})
+    headers = {k: v for k, v in activation.headers.items()}
+
+    transaction = sentry_sdk.continue_trace(
+        environ_or_headers=headers,
+        op="task.taskworker",
+        name=f"{activation.namespace}:{activation.taskname}",
+    )
+    with (
+        track_memory_usage("taskworker.worker.memory_change"),
+        sentry_sdk.start_transaction(transaction),
+    ):
+        task_func(*args, **kwargs)
+
+
+class TaskWorker:
+    """
+    A TaskWorker fetches tasks from a taskworker RPC host and handles executing task activations.
+
+    Tasks are executed in a forked process so that processing timeouts can be enforced.
+    As tasks are completed status changes will be sent back to the RPC host and new tasks
+    will be fetched.
+
+    Taskworkers can be run with `sentry run taskworker`
+    """
+
+    def __init__(
+        self,
+        rpc_host: str,
+        max_task_count: int | None = None,
+        namespace: str | None = None,
+        concurrency: int = 1,
+        **options: dict[str, Any],
+    ) -> None:
+        self.options = options
+        self._execution_count = 0
+        self._worker_id = uuid4().hex
+        self._max_task_count = max_task_count
+        self._namespace = namespace
+        self._concurrency = concurrency
+        self.client = TaskworkerClient(rpc_host)
+        self._child_tasks: multiprocessing.Queue[TaskActivation] = multiprocessing.Queue(
+            maxsize=(concurrency * 10)
+        )
+        self._processed_tasks: multiprocessing.Queue[ProcessingResult] = multiprocessing.Queue(
+            maxsize=(concurrency * 10)
+        )
+        self._children: list[multiprocessing.Process] = []
+        self._shutdown_event = multiprocessing.Event()
+
+    def __del__(self) -> None:
+        self._shutdown()
+
+    def do_imports(self) -> None:
+        for module in settings.TASKWORKER_IMPORTS:
+            __import__(module)
+
+    def start(self) -> int:
+        """
+        Run the worker main loop
+
+        Once started a Worker will loop until it is killed, or
+        completes its max_task_count when it shuts down.
+        """
+        self.do_imports()
+        self._spawn_children()
+
+        signal.signal(signal.SIGINT, self._handle_sigint)
+
+        while True:
+            self.run_once()
+
+    def run_once(self) -> None:
+        """Access point for tests to run a single worker loop"""
+        self._add_task()
+        self._drain_result()
+        self._spawn_children()
+
+    def _handle_sigint(self, signum: int, frame: FrameType | None) -> None:
+        logger.info("taskworker.worker.sigint_received")
+        self._shutdown()
+        raise KeyboardInterrupt("Shutdown complete")
+
+    def _shutdown(self) -> None:
+        """
+        Shutdown cleanly
+        Activate the shutdown event and drain results before terminating children.
+        """
+        self._shutdown_event.set()
+
+        while True:
+            more = self._drain_result(fetch=False)
+            if not more:
+                break
+        for child in self._children:
+            child.terminate()
+            child.join()
+
+    def _add_task(self) -> None:
+        """
+        Add a task to child tasks queue
+        """
+        if self._child_tasks.full():
+            return
+
+        task = self.fetch_task()
+        if task:
+            try:
+                self._child_tasks.put(task, block=False)
+            except queue.Full:
+                logger.exception(
+                    "taskworker.add_task.child_task_queue_full", extra={"task_id": task.id}
+                )
+
+    def _drain_result(self, fetch: bool = True) -> bool:
+        """
+        Consume results from children and update taskbroker
+        """
+        try:
+            result = self._processed_tasks.get_nowait()
+        except queue.Empty:
+            return False
+
+        if fetch:
+            next_task = self.client.update_task(
+                task_id=result.task_id,
+                status=result.status,
+                fetch_next_task=FetchNextTask(namespace=self._namespace),
+            )
+            if next_task:
+                try:
+                    self._child_tasks.put(next_task, block=False)
+                except queue.Full:
+                    logger.exception(
+                        "taskworker.drain_result.child_task_queue_full",
+                        extra={"task_id": next_task.id},
+                    )
+            return True
+
+        self.client.update_task(
+            task_id=result.task_id,
+            status=result.status,
         )
+        return True
+
+    def _spawn_children(self) -> None:
+        active_children = [child for child in self._children if child.is_alive()]
+        if len(active_children) >= self._concurrency:
+            return
+        for _ in range(self._concurrency - len(active_children)):
+            process = multiprocessing.Process(
+                target=child_worker,
+                args=(
+                    self._child_tasks,
+                    self._processed_tasks,
+                    self._shutdown_event,
+                    self._max_task_count,
+                ),
+            )
+            process.start()
+            active_children.append(process)
+        self._children = active_children
+
+    def fetch_task(self) -> TaskActivation | None:
+        try:
+            activation = self.client.get_task(self._namespace)
+        except grpc.RpcError:
+            metrics.incr("taskworker.worker.fetch_task", tags={"status": "failed"})
+            logger.info("taskworker.fetch_task.failed")
+            return None
+
+        if not activation:
+            metrics.incr("taskworker.worker.fetch_task", tags={"status": "notfound"})
+            logger.debug("taskworker.fetch_task.not_found")
+            return None
+
+        metrics.incr(
+            "taskworker.worker.fetch_task",
+            tags={"status": "success", "namespace": activation.namespace},
+        )
+        return activation
diff --git a/tests/sentry/taskworker/test_worker.py b/tests/sentry/taskworker/test_worker.py
index 5d174958e6c1de..711e70375955b6 100644
--- a/tests/sentry/taskworker/test_worker.py
+++ b/tests/sentry/taskworker/test_worker.py
@@ -1,3 +1,6 @@
+import queue
+import time
+from multiprocessing import Event
 from unittest import mock
 
 from sentry_protos.sentry.v1.taskworker_pb2 import (
@@ -9,7 +12,7 @@
 )
 
 import sentry.taskworker.tasks.examples as example_tasks
-from sentry.taskworker.worker import TaskWorker
+from sentry.taskworker.worker import ProcessingResult, TaskWorker, child_worker
 from sentry.testutils.cases import TestCase
 
 SIMPLE_TASK = TaskActivation(
@@ -79,129 +82,153 @@ def test_fetch_no_task(self) -> None:
             mock_get.assert_called_once()
         assert task is None
 
-    def test_process_task_complete(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=100)
-        with mock.patch.object(taskworker.client, "update_task") as mock_update:
-            mock_update.return_value = RETRY_TASK
+    def test_run_once_no_next_task(self) -> None:
+        max_runtime = 5
+        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=1)
+        with mock.patch.object(taskworker, "client") as mock_client:
+            mock_client.get_task.return_value = SIMPLE_TASK
+            # No next_task returned
+            mock_client.update_task.return_value = None
 
-            result = taskworker.process_task(SIMPLE_TASK)
+            # Run once to add the task and then poll until the task is complete.
+            taskworker.run_once()
+            start = time.time()
+            while True:
+                taskworker.run_once()
+                if mock_client.update_task.called:
+                    break
+                if time.time() - start > max_runtime:
+                    raise AssertionError("Timeout waiting for get_task to be called")
 
-            mock_update.assert_called_with(
+            assert mock_client.get_task.called
+            mock_client.update_task.assert_called_with(
                 task_id=SIMPLE_TASK.id,
                 status=TASK_ACTIVATION_STATUS_COMPLETE,
                 fetch_next_task=FetchNextTask(namespace=None),
             )
 
-            assert result
-            assert result.id == RETRY_TASK.id
+    def test_run_once_with_next_task(self) -> None:
+        # Cover the scenario where update_task returns the next task which should
+        # be processed.
+        max_runtime = 5
+        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=1)
+        with mock.patch.object(taskworker, "client") as mock_client:
 
-    def test_process_task_retry(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=100)
-        with mock.patch.object(taskworker.client, "update_task") as mock_update:
-            mock_update.return_value = SIMPLE_TASK
-            result = taskworker.process_task(RETRY_TASK)
+            def update_task_response(*args, **kwargs):
+                if mock_client.update_task.call_count >= 1:
+                    return None
+                return SIMPLE_TASK
+
+            mock_client.update_task.side_effect = update_task_response
+            mock_client.get_task.return_value = SIMPLE_TASK
 
-            mock_update.assert_called_with(
-                task_id=RETRY_TASK.id,
-                status=TASK_ACTIVATION_STATUS_RETRY,
+            # Run until two tasks have been processed
+            start = time.time()
+            while True:
+                taskworker.run_once()
+                if mock_client.update_task.call_count >= 2:
+                    break
+                if time.time() - start > max_runtime:
+                    raise AssertionError("Timeout waiting for get_task to be called")
+
+            assert mock_client.get_task.called
+            assert mock_client.update_task.call_count == 2
+            mock_client.update_task.assert_called_with(
+                task_id=SIMPLE_TASK.id,
+                status=TASK_ACTIVATION_STATUS_COMPLETE,
                 fetch_next_task=FetchNextTask(namespace=None),
             )
 
-            assert result
-            assert result.id == SIMPLE_TASK.id
 
-    def test_process_task_failure(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=100)
-        with mock.patch.object(taskworker.client, "update_task") as mock_update:
-            mock_update.return_value = SIMPLE_TASK
-            result = taskworker.process_task(FAIL_TASK)
+def test_child_worker_complete() -> None:
+    todo: queue.Queue[TaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
 
-            mock_update.assert_called_with(
-                task_id=FAIL_TASK.id,
-                status=TASK_ACTIVATION_STATUS_FAILURE,
-                fetch_next_task=FetchNextTask(namespace=None),
-            )
-            assert result
-            assert result.id == SIMPLE_TASK.id
+    todo.put(SIMPLE_TASK)
+    child_worker(todo, processed, shutdown, max_task_count=1)
 
-    def test_process_task_at_most_once(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=100)
-        with mock.patch.object(taskworker.client, "update_task") as mock_update:
-            mock_update.return_value = SIMPLE_TASK
-            result = taskworker.process_task(AT_MOST_ONCE_TASK)
+    assert todo.empty()
+    result = processed.get()
+    assert result.task_id == SIMPLE_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_COMPLETE
 
-            mock_update.assert_called_with(
-                task_id=AT_MOST_ONCE_TASK.id,
-                status=TASK_ACTIVATION_STATUS_COMPLETE,
-                fetch_next_task=FetchNextTask(namespace=None),
-            )
-        assert taskworker.process_task(AT_MOST_ONCE_TASK) is None
-        assert result
-        assert result.id == SIMPLE_TASK.id
 
-        result = taskworker.process_task(AT_MOST_ONCE_TASK)
-        assert result is None
+def test_child_worker_retry_task() -> None:
+    todo: queue.Queue[TaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
 
-    def test_start_max_task_count(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=1)
-        with mock.patch.object(taskworker, "client") as mock_client:
-            mock_client.get_task.return_value = SIMPLE_TASK
-            mock_client.update_task.return_value = None
+    todo.put(RETRY_TASK)
+    child_worker(todo, processed, shutdown, max_task_count=1)
 
-            result = taskworker.start()
+    assert todo.empty()
+    result = processed.get()
+    assert result.task_id == RETRY_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_RETRY
 
-            # start should exit after completing the one task
-            assert result == 0
-            assert mock_client.get_task.called
-            mock_client.update_task.assert_called_with(
-                task_id=SIMPLE_TASK.id,
-                status=TASK_ACTIVATION_STATUS_COMPLETE,
-                fetch_next_task=FetchNextTask(namespace=None),
-            )
 
-    def test_start_loop(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=2)
-        with mock.patch.object(taskworker, "client") as mock_client:
-            mock_client.get_task.return_value = SIMPLE_TASK
-            mock_client.update_task.side_effect = [RETRY_TASK, None]
+def test_child_worker_failure_task() -> None:
+    todo: queue.Queue[TaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
 
-            # Because complete_task is returning another task, we should get
-            # two executions. One successful and one retry
-            result = taskworker.start()
+    todo.put(FAIL_TASK)
+    child_worker(todo, processed, shutdown, max_task_count=1)
 
-            # start should exit after completing the both task
-            assert result == 0
-            assert mock_client.get_task.call_count == 1
-            assert mock_client.update_task.call_count == 2
+    assert todo.empty()
+    result = processed.get()
+    assert result.task_id == FAIL_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_FAILURE
 
-            mock_client.update_task.assert_any_call(
-                task_id=SIMPLE_TASK.id,
-                status=TASK_ACTIVATION_STATUS_COMPLETE,
-                fetch_next_task=FetchNextTask(namespace=None),
-            )
-            mock_client.update_task.assert_any_call(
-                task_id=RETRY_TASK.id,
-                status=TASK_ACTIVATION_STATUS_RETRY,
-                fetch_next_task=FetchNextTask(namespace=None),
-            )
 
-    def test_start_keyboard_interrupt(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=2)
-        with mock.patch.object(taskworker, "client") as mock_client:
-            mock_client.get_task.side_effect = KeyboardInterrupt()
+def test_child_worker_shutdown() -> None:
+    todo: queue.Queue[TaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
+    shutdown.set()
 
-            result = taskworker.start()
-            assert result == 1, "Exit non-zero"
+    todo.put(SIMPLE_TASK)
+    child_worker(todo, processed, shutdown, max_task_count=1)
 
-    def test_start_unknown_task(self) -> None:
-        taskworker = TaskWorker(rpc_host="127.0.0.1:50051", max_task_count=2)
-        with mock.patch.object(taskworker, "client") as mock_client:
-            mock_client.get_task.return_value = UNDEFINED_TASK
+    # When shutdown has been set, the child should not process more tasks.
+    assert todo.qsize() == 1
+    assert processed.qsize() == 0
 
-            result = taskworker.start()
-            assert result == 0, "Exit zero, all tasks complete"
-            mock_client.update_task.assert_any_call(
-                task_id=UNDEFINED_TASK.id,
-                status=TASK_ACTIVATION_STATUS_FAILURE,
-                fetch_next_task=FetchNextTask(namespace=None),
-            )
+
+def test_child_worker_unknown_task() -> None:
+    todo: queue.Queue[TaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
+
+    todo.put(UNDEFINED_TASK)
+    todo.put(SIMPLE_TASK)
+    child_worker(todo, processed, shutdown, max_task_count=1)
+
+    result = processed.get()
+    assert result.task_id == UNDEFINED_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_FAILURE
+
+    result = processed.get()
+    assert result.task_id == SIMPLE_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_COMPLETE
+
+
+def test_child_worker_at_most_once() -> None:
+    todo: queue.Queue[TaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
+
+    todo.put(AT_MOST_ONCE_TASK)
+    todo.put(AT_MOST_ONCE_TASK)
+    todo.put(SIMPLE_TASK)
+    child_worker(todo, processed, shutdown, max_task_count=2)
+
+    assert todo.empty()
+    result = processed.get(block=False)
+    assert result.task_id == AT_MOST_ONCE_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_COMPLETE
+
+    result = processed.get(block=False)
+    assert result.task_id == SIMPLE_TASK.id
+    assert result.status == TASK_ACTIVATION_STATUS_COMPLETE


Code review comments:
diff_hunk: @@ -65,137 +51,84 @@ def get_at_most_once_key(namespace: str, taskname: str, task_id: str) -> str:
     return f"tw:amo:{namespace}:{taskname}:{task_id}"
 
 
-class TaskWorker:
-    """
-    A TaskWorker fetches tasks from a taskworker RPC host and handles executing task activations.
-
-    Tasks are executed in a forked process so that processing timeouts can be enforced.
-    As tasks are completed status changes will be sent back to the RPC host and new tasks
-    will be fetched.
-
-    Taskworkers can be run with `sentry run taskworker`
-    """
+def _get_known_task(activation: TaskActivation) -> Task[Any, Any] | None:
+    if not taskregistry.contains(activation.namespace):
+        logger.error(
+            "taskworker.invalid_namespace",
+            extra={"namespace": activation.namespace, "taskname": activation.taskname},
+        )
+        return None
 
-    def __init__(
-        self,
-        rpc_host: str,
-        max_task_count: int | None = None,
-        namespace: str | None = None,
-        **options: dict[str, Any],
-    ) -> None:
-        self.options = options
-        self._execution_count = 0
-        self._worker_id = uuid4().hex
-        self._max_task_count = max_task_count
-        self._namespace = namespace
-        self.client = TaskworkerClient(rpc_host)
-        self._pool: Pool | None = None
-        self._build_pool()
+    namespace = taskregistry.get(activation.namespace)
+    if not namespace.contains(activation.taskname):
+        logger.error(
+            "taskworker.invalid_taskname",
+            extra={"namespace": activation.namespace, "taskname": activation.taskname},
+        )
+        return None
+    return namespace.get(activation.taskname)
 
-    def __del__(self) -> None:
-        if self._pool:
-            self._pool.terminate()
 
-    def _build_pool(self) -> None:
-        if self._pool:
-            self._pool.terminate()
-        self._pool = mp_context.Pool(processes=1, initializer=_init_pool_process)
+def child_worker(
+    child_tasks: queue.Queue[TaskActivation],
+    processed_tasks: queue.Queue[ProcessingResult],
+    shutdown_event: Event,
+    max_task_count: int | None,
+) -> None:
+    for module in settings.TASKWORKER_IMPORTS:
+        __import__(module)
 
-    def do_imports(self) -> None:
-        for module in settings.TASKWORKER_IMPORTS:
-            __import__(module)
-        self._build_pool()
+    current_task_id: str | None = None
+    processed_task_count = 0
 
-    def start(self) -> int:
+    def handle_alarm(signum: int, frame: Any) -> None:
         """
-        Run the worker main loop
+        Handle SIGALRM
 
-        Once started a Worker will loop until it is killed, or
-        completes its max_task_count when it shuts down.
+        If we hit an alarm in a child, we need to push a result
+        and terminate the child.
         """
-        self.do_imports()
-        next_task: TaskActivation | None = None
-        task: TaskActivation | None = None
-        try:
-            while True:
-                if next_task:
-                    task = next_task
-                    next_task = None
-                else:
-                    task = self.fetch_task()
-
-                if not task:
-                    metrics.incr("taskworker.worker.no_task.pause")
-                    time.sleep(1)
-                    continue
-
-                next_task = self.process_task(task)
-                if (
-                    self._max_task_count is not None
-                    and self._max_task_count <= self._execution_count
-                ):
-                    metrics.incr(
-                        "taskworker.worker.max_task_count_reached",
-                        tags={"count": self._execution_count},
-                    )
-                    logger.info("Max task execution count reached. Terminating")
-                    return 0
-
-        except KeyboardInterrupt:
-            return 1
-        except Exception:
-            logger.exception("Worker process crashed")
-            return 2
+        nonlocal current_task_id, processed_tasks
 
-    def fetch_task(self) -> TaskActivation | None:
-        try:
-            activation = self.client.get_task(self._namespace)
-        except grpc.RpcError:
-            metrics.incr("taskworker.worker.get_task.failed")
-            logger.info("get_task failed. Retrying in 1 second")
-            return None
+        if current_task_id:
+            processed_tasks.put(
+                ProcessingResult(task_id=current_task_id, status=TASK_ACTIVATION_STATUS_FAILURE)
+            )
+        metrics.incr("taskworker.worker.processing_deadline_exceeded")
+        sys.exit(1)
 
-        if not activation:
-            metrics.incr("taskworker.worker.get_task.not_found")
-            logger.debug("No task fetched")
-            return None
+    signal.signal(signal.SIGALRM, handle_alarm)
 
-        metrics.incr("taskworker.worker.get_task.success")
-        return activation
-
-    def _get_known_task(self, activation: TaskActivation) -> Task[Any, Any] | None:
-        if not taskregistry.contains(activation.namespace):
-            logger.error(
-                "taskworker.invalid_namespace",
-                extra={"namespace": activation.namespace, "taskname": activation.taskname},
+    while True:
+        if max_task_count and processed_task_count >= max_task_count:
+            metrics.incr(
+                "taskworker.worker.max_task_count_reached",
+                tags={"count": processed_task_count},
             )
-            return None
+            logger.info("taskwworker.max_task_count_reached")
+            break
 
-        namespace = taskregistry.get(activation.namespace)
-        if not namespace.contains(activation.taskname):
-            logger.error(
-                "taskworker.invalid_taskname",
-                extra={"namespace": activation.namespace, "taskname": activation.taskname},
-            )
-            return None
-        return namespace.get(activation.taskname)
+        if shutdown_event.is_set():
+            logger.info("taskworker.worker.shutdown_event")
+            break
 
-    def process_task(self, activation: TaskActivation) -> TaskActivation | None:
-        assert self._pool
-        task = self._get_known_task(activation)
-        if not task:
+        try:
+            activation = child_tasks.get_nowait()
path: src/sentry/taskworker/worker.py
text: @user1:
I don't think `get_nowait` is necessarily correct here. I understand this is ensuring that the process doesn't block while waiting for a task before checking for the shutdown, but I think some kind of timeout/delay would good to here to avoid spiking the CPU. Maybe like 100ms or something like that?

@author:
Good point about the potential CPU burn on an empty queue. I'll put a blocking get with a timeout in.

diff_hunk: @@ -65,137 +51,84 @@ def get_at_most_once_key(namespace: str, taskname: str, task_id: str) -> str:
     return f"tw:amo:{namespace}:{taskname}:{task_id}"
 
 
-class TaskWorker:
-    """
-    A TaskWorker fetches tasks from a taskworker RPC host and handles executing task activations.
-
-    Tasks are executed in a forked process so that processing timeouts can be enforced.
-    As tasks are completed status changes will be sent back to the RPC host and new tasks
-    will be fetched.
-
-    Taskworkers can be run with `sentry run taskworker`
-    """
+def _get_known_task(activation: TaskActivation) -> Task[Any, Any] | None:
+    if not taskregistry.contains(activation.namespace):
+        logger.error(
+            "taskworker.invalid_namespace",
+            extra={"namespace": activation.namespace, "taskname": activation.taskname},
+        )
+        return None
 
-    def __init__(
-        self,
-        rpc_host: str,
-        max_task_count: int | None = None,
-        namespace: str | None = None,
-        **options: dict[str, Any],
-    ) -> None:
-        self.options = options
-        self._execution_count = 0
-        self._worker_id = uuid4().hex
-        self._max_task_count = max_task_count
-        self._namespace = namespace
-        self.client = TaskworkerClient(rpc_host)
-        self._pool: Pool | None = None
-        self._build_pool()
+    namespace = taskregistry.get(activation.namespace)
+    if not namespace.contains(activation.taskname):
+        logger.error(
+            "taskworker.invalid_taskname",
+            extra={"namespace": activation.namespace, "taskname": activation.taskname},
+        )
+        return None
+    return namespace.get(activation.taskname)
 
-    def __del__(self) -> None:
-        if self._pool:
-            self._pool.terminate()
 
-    def _build_pool(self) -> None:
-        if self._pool:
-            self._pool.terminate()
-        self._pool = mp_context.Pool(processes=1, initializer=_init_pool_process)
+def child_worker(
+    child_tasks: queue.Queue[TaskActivation],
+    processed_tasks: queue.Queue[ProcessingResult],
+    shutdown_event: Event,
+    max_task_count: int | None,
+) -> None:
+    for module in settings.TASKWORKER_IMPORTS:
+        __import__(module)
 
-    def do_imports(self) -> None:
-        for module in settings.TASKWORKER_IMPORTS:
-            __import__(module)
-        self._build_pool()
+    current_task_id: str | None = None
+    processed_task_count = 0
 
-    def start(self) -> int:
+    def handle_alarm(signum: int, frame: Any) -> None:
         """
-        Run the worker main loop
+        Handle SIGALRM
 
-        Once started a Worker will loop until it is killed, or
-        completes its max_task_count when it shuts down.
+        If we hit an alarm in a child, we need to push a result
+        and terminate the child.
         """
-        self.do_imports()
-        next_task: TaskActivation | None = None
-        task: TaskActivation | None = None
-        try:
-            while True:
-                if next_task:
-                    task = next_task
-                    next_task = None
-                else:
-                    task = self.fetch_task()
-
-                if not task:
-                    metrics.incr("taskworker.worker.no_task.pause")
-                    time.sleep(1)
-                    continue
-
-                next_task = self.process_task(task)
-                if (
-                    self._max_task_count is not None
-                    and self._max_task_count <= self._execution_count
-                ):
-                    metrics.incr(
-                        "taskworker.worker.max_task_count_reached",
-                        tags={"count": self._execution_count},
-                    )
-                    logger.info("Max task execution count reached. Terminating")
-                    return 0
-
-        except KeyboardInterrupt:
-            return 1
-        except Exception:
-            logger.exception("Worker process crashed")
-            return 2
+        nonlocal current_task_id, processed_tasks
 
-    def fetch_task(self) -> TaskActivation | None:
-        try:
-            activation = self.client.get_task(self._namespace)
-        except grpc.RpcError:
-            metrics.incr("taskworker.worker.get_task.failed")
-            logger.info("get_task failed. Retrying in 1 second")
-            return None
+        if current_task_id:
+            processed_tasks.put(
+                ProcessingResult(task_id=current_task_id, status=TASK_ACTIVATION_STATUS_FAILURE)
+            )
+        metrics.incr("taskworker.worker.processing_deadline_exceeded")
+        sys.exit(1)
 
-        if not activation:
-            metrics.incr("taskworker.worker.get_task.not_found")
-            logger.debug("No task fetched")
-            return None
+    signal.signal(signal.SIGALRM, handle_alarm)
 
-        metrics.incr("taskworker.worker.get_task.success")
-        return activation
-
-    def _get_known_task(self, activation: TaskActivation) -> Task[Any, Any] | None:
-        if not taskregistry.contains(activation.namespace):
-            logger.error(
-                "taskworker.invalid_namespace",
-                extra={"namespace": activation.namespace, "taskname": activation.taskname},
+    while True:
+        if max_task_count and processed_task_count >= max_task_count:
+            metrics.incr(
+                "taskworker.worker.max_task_count_reached",
+                tags={"count": processed_task_count},
             )
-            return None
+            logger.info("taskwworker.max_task_count_reached")
path: src/sentry/taskworker/worker.py
text: @user1:
```suggestion
            logger.info("taskworker.max_task_count_reached")
```