diff --git a/devservices/config.yml b/devservices/config.yml index 514bb94432d7..8bfa916cfe5f 100644 --- a/devservices/config.yml +++ b/devservices/config.yml @@ -400,9 +400,9 @@ x-programs: devserver: command: sentry devserver taskworker: - command: sentry run taskworker + command: sentry run taskworker --autoreload taskworker-sentry: - command: sentry run taskworker --rpc-host localhost:50052 + command: sentry run taskworker --autoreload --rpc-host localhost:50052 taskworker-scheduler: command: sentry run taskworker-scheduler ingest-events: diff --git a/src/sentry/runner/commands/devserver.py b/src/sentry/runner/commands/devserver.py index b218b58c25e5..0e524fab267d 100644 --- a/src/sentry/runner/commands/devserver.py +++ b/src/sentry/runner/commands/devserver.py @@ -39,6 +39,13 @@ def _get_daemon(name: str) -> tuple[str, list[str]]: return name, _DEFAULT_DAEMONS[name] +def _get_taskworker_daemon(reload: bool) -> tuple[str, list[str]]: + name, command = _get_daemon("taskworker") + if reload: + command = [*command, "--autoreload"] + return name, command + + @click.command() @click.option( "--reload/--no-reload", @@ -303,7 +310,7 @@ def devserver( "Disable TASKWORKER_ALWAYS_EAGER in your settings file to spawn workers." ) - daemons.append(_get_daemon("taskworker")) + daemons.append(_get_taskworker_daemon(reload)) from sentry import eventstream @@ -504,7 +511,10 @@ def devserver( control_services.append("taskworker-scheduler") for service in control_services: - name, cmd = _get_daemon(service) + if service == "taskworker": + name, cmd = _get_taskworker_daemon(reload) + else: + name, cmd = _get_daemon(service) name = f"control.{name}" quiet = bool( name not in (settings.DEVSERVER_LOGS_ALLOWLIST or ()) diff --git a/src/sentry/runner/commands/run.py b/src/sentry/runner/commands/run.py index fef6376bbb29..1cd92f0afedd 100644 --- a/src/sentry/runner/commands/run.py +++ b/src/sentry/runner/commands/run.py @@ -1,10 +1,13 @@ from __future__ import annotations import logging +import multiprocessing import os import random import signal +import threading import time +from multiprocessing.synchronize import Event from typing import Any import click @@ -234,6 +237,12 @@ def taskworker_scheduler(redis_cluster: str, **options: Any) -> None: default=None, type=int, ) +@click.option( + "--autoreload", + is_flag=True, + default=False, + help="Restart the worker when Python files change. For local development only.", +) @log_options() @configuration def taskworker(**options: Any) -> None: @@ -241,8 +250,71 @@ def taskworker(**options: Any) -> None: Run a taskworker worker """ os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0" - # TODO(mark) restore autoreload - run_taskworker(**options) + if options.pop("autoreload"): + _run_taskworker_with_autoreload(**options) + else: + run_taskworker(**options) + + +def _run_taskworker_with_autoreload(**options: Any) -> None: + # watchfiles pulls in anyio, which CLI startup must not import. + import watchfiles + + import sentry + from sentry.runner.settings import discover_configs + + # Run the worker in its own process so its signal handlers get the main thread. + # Known gap: a task still running 5s after SIGINT orphans its child and gets redelivered. + ready = multiprocessing.get_context("spawn").Event() + try: + watchfiles.run_process( + os.path.dirname(sentry.__file__), + discover_configs()[0], + target=_run_reloaded_taskworker, + kwargs={"ready": ready, **options}, + callback=lambda _: _wait_for_reloaded_taskworker(ready), + watch_filter=watchfiles.PythonFilter(), + ) + except KeyboardInterrupt: + # Stopped while waiting on a restart. watchfiles has already stopped + # the worker on its way out. + pass + + +def _wait_for_reloaded_taskworker(ready: Event) -> None: + # A SIGINT before the worker registers its handler skips the shutdown that + # stops its children, which are then orphaned when watchfiles kills it. + # So hold off restarting until the handler is in place. + while not ready.wait(timeout=0.1): + if not multiprocessing.active_children(): + break + ready.clear() + + +def _run_reloaded_taskworker(ready: Event, **options: Any) -> None: + # Runs in a fresh spawned process, which has to configure Sentry again. + from sentry.runner import configure + + threading.Thread( + target=_set_when_sigint_handled, + args=(ready, signal.getsignal(signal.SIGINT)), + daemon=True, + ).start() + try: + configure() + run_taskworker(**options) + except KeyboardInterrupt: + # A Ctrl-C arrived before the worker registered its SIGINT handler. + # Ignore the second SIGINT watchfiles sends while we exit. + signal.signal(signal.SIGINT, signal.SIG_IGN) + + +def _set_when_sigint_handled(ready: Event, initial_handler: Any) -> None: + # The worker registers its handler deep inside start(), with no hook to + # tell us, so watch for it to replace the initial one. + while signal.getsignal(signal.SIGINT) is initial_handler: + time.sleep(0.05) + ready.set() def run_taskworker(