From c8166a2ffe997b506af6ad767a96848a55946d9a Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Thu, 30 Jul 2026 14:50:27 -0700 Subject: [PATCH 1/6] feat(operator): control server and generic runner --- sentry_streams/Cargo.lock | 25 +-- sentry_streams/Cargo.toml | 1 - .../sentry_streams/adapters/arroyo/adapter.py | 21 ++- .../adapters/arroyo/rust_arroyo.py | 14 +- .../sentry_streams/adapters/stream_adapter.py | 91 +++++++++- sentry_streams/sentry_streams/control.py | 104 ++++++++++++ .../sentry_streams/dummy/dummy_adapter.py | 5 +- sentry_streams/sentry_streams/runner.py | 160 +++++++++++++++++- .../sentry_streams/server/__init__.py | 0 .../sentry_streams/server/control_server.py | 97 +++++++++++ sentry_streams/src/consumer.rs | 61 ++++--- sentry_streams/src/run.rs | 10 +- .../tests/adapters/arroyo/test_adapter.py | 63 ++++++- sentry_streams/tests/adapters/fake_adapter.py | 85 ++++++++++ sentry_streams/tests/test_control.py | 160 ++++++++++++++++++ sentry_streams/tests/test_control_server.py | 92 ++++++++++ sentry_streams/tests/test_runner.py | 40 ++++- 17 files changed, 956 insertions(+), 73 deletions(-) create mode 100644 sentry_streams/sentry_streams/control.py create mode 100644 sentry_streams/sentry_streams/server/__init__.py create mode 100644 sentry_streams/sentry_streams/server/control_server.py create mode 100644 sentry_streams/tests/adapters/fake_adapter.py create mode 100644 sentry_streams/tests/test_control.py create mode 100644 sentry_streams/tests/test_control_server.py diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index e21a054a..d61d9533 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -495,16 +495,6 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" -[[package]] -name = "ctrlc" -version = "3.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "697b5419f348fd5ae2478e8018cb016c00a5881c7f46c717de98ffd135a5651c" -dependencies = [ - "nix 0.29.0", - "windows-sys 0.59.0", -] - [[package]] name = "debugid" version = "0.8.0" @@ -1351,18 +1341,6 @@ dependencies = [ "smallvec", ] -[[package]] -name = "nix" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" -dependencies = [ - "bitflags", - "cfg-if", - "cfg_aliases", - "libc", -] - [[package]] name = "nix" version = "0.30.1" @@ -1660,7 +1638,7 @@ checksum = "e4022a17595a00d6a369236fdae483f0de7f0a339960a53118b818238e132224" dependencies = [ "android_system_properties", "log", - "nix 0.30.1", + "nix", "objc2", "objc2-foundation", "objc2-ui-kit", @@ -2138,7 +2116,6 @@ dependencies = [ "anyhow", "chrono", "clap", - "ctrlc", "gcp_auth", "log", "metrics", diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index 07847758..088c188c 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -10,7 +10,6 @@ sentry_arroyo = { version = "2.40.0", features = ["ssl"] } chrono = "0.4.40" tracing = "0.1.40" tracing-subscriber = "0.3.20" -ctrlc = "3.4.6" rdkafka = { version = "0.37.0", features = ["cmake-build", "tracing"] } anyhow = "1.0.98" reqwest = "0.12.15" diff --git a/sentry_streams/sentry_streams/adapters/arroyo/adapter.py b/sentry_streams/sentry_streams/adapters/arroyo/adapter.py index e7e31fe5..a3f829bc 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/adapter.py +++ b/sentry_streams/sentry_streams/adapters/arroyo/adapter.py @@ -33,7 +33,11 @@ RouterStep, StreamSinkStep, ) -from sentry_streams.adapters.stream_adapter import PipelineConfig, StreamAdapter +from sentry_streams.adapters.stream_adapter import ( + PipelineConfig, + RuntimeState, + StreamAdapter, +) from sentry_streams.config_types import ( KafkaConsumerConfig, KafkaProducerConfig, @@ -322,7 +326,7 @@ def create_processors(self) -> None: for source, consumer in self.__consumers.items() } - def run(self) -> None: + def _run(self) -> None: """ Starts the pipeline """ @@ -332,14 +336,17 @@ def run(self) -> None: source = next(iter(self.__consumers)) processor = self.__processors[source] + if self.status.state is RuntimeState.STOPPING: + processor.signal_shutdown() + else: + self._set_status(RuntimeState.CONSUMING) + processor.run() - def shutdown(self) -> None: + def _shutdown(self) -> None: """ Shutdown the arroyo processors allowing them to terminate the inflight work. """ - assert len(self.__consumers) == 1, "Only one consumer is supported" - source = next(iter(self.__consumers)) - processor = self.__processors[source] - processor.signal_shutdown() + for processor in self.__processors.values(): + processor.signal_shutdown() diff --git a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py b/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py index 94d4e67b..9c0ff705 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py +++ b/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py @@ -27,7 +27,11 @@ from sentry_streams.adapters.arroyo.routers import build_branches from sentry_streams.adapters.arroyo.routes import Route from sentry_streams.adapters.arroyo.steps_chain import TransformChains -from sentry_streams.adapters.stream_adapter import PipelineConfig, StreamAdapter +from sentry_streams.adapters.stream_adapter import ( + PipelineConfig, + RuntimeState, + StreamAdapter, +) from sentry_streams.config_types import ( KafkaConsumerConfig, KafkaProducerConfig, @@ -587,18 +591,20 @@ def routing_function(msg: Message[Any]) -> str: ) return build_branches(stream, step.routing_table.values()) - def run(self) -> None: + def _run(self) -> None: """ Starts the pipeline """ # TODO: Support multiple consumers assert len(self.__consumers) == 1, "Multiple consumers not supported yet" consumer = next(iter(self.__consumers.values())) + self._set_status(RuntimeState.CONSUMING) consumer.run() - def shutdown(self) -> None: + def _shutdown(self) -> None: """ Shutdown the arroyo processors allowing them to terminate the inflight work. """ - raise NotImplementedError + for consumer in self.__consumers.values(): + consumer.shutdown() diff --git a/sentry_streams/sentry_streams/adapters/stream_adapter.py b/sentry_streams/sentry_streams/adapters/stream_adapter.py index f034771b..f0e83700 100644 --- a/sentry_streams/sentry_streams/adapters/stream_adapter.py +++ b/sentry_streams/sentry_streams/adapters/stream_adapter.py @@ -1,6 +1,9 @@ from __future__ import annotations +import threading from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import StrEnum from typing import ( Any, Callable, @@ -41,6 +44,46 @@ StreamSinkT = TypeVar("StreamSinkT") +class RuntimeState(StrEnum): + IDLE = "idle" + STARTING = "starting" + CONSUMING = "consuming" + STOPPING = "stopping" + STOPPED = "stopped" + ERRORED = "errored" + + @property + def is_terminal(self) -> bool: + return self in (RuntimeState.STOPPED, RuntimeState.ERRORED) + + +# To simplify lifecycle and state transitions we can use a forward-only ranking. +# A higher-ranked state should realistically never transition to a lower one. +# Would need to be updated if we allow restarting (but we don't for now). + +_STATE_RANKS: Mapping[RuntimeState, int] = { + RuntimeState.IDLE: 0, + RuntimeState.STARTING: 1, + RuntimeState.CONSUMING: 2, + RuntimeState.STOPPING: 3, + RuntimeState.STOPPED: 4, + RuntimeState.ERRORED: 4, +} + + +@dataclass(frozen=True) +class RuntimeStatus: + state: RuntimeState + error: str | None = None + + @property + def is_terminal(self) -> bool: + return self.state.is_terminal + + def as_dict(self) -> dict[str, str | None]: + return {"state": self.state.value, "error": self.error} + + class StreamAdapter(ABC, Generic[StreamT, StreamSinkT]): """ A generic adapter for mapping sentry_streams APIs @@ -48,6 +91,50 @@ class StreamAdapter(ABC, Generic[StreamT, StreamSinkT]): be extended to specific runtimes. """ + def __init__(self) -> None: + self.__status_lock = threading.Lock() + self.__status = RuntimeStatus(RuntimeState.IDLE) + + @property + def status(self) -> RuntimeStatus: + with self.__status_lock: + return self.__status + + def _set_status(self, state: RuntimeState, error: str | None = None) -> None: + with self.__status_lock: + if _STATE_RANKS[state] <= _STATE_RANKS[self.__status.state]: + return + self.__status = RuntimeStatus(state, error) + + def begin_start(self) -> RuntimeStatus: + with self.__status_lock: + if self.__status.state is RuntimeState.IDLE: + self.__status = RuntimeStatus(RuntimeState.STARTING) + return self.__status + + def run(self) -> None: + if self.begin_start().state is not RuntimeState.STARTING: + self._set_status(RuntimeState.STOPPED) + return + + try: + self._run() + except BaseException as exc: + self._set_status(RuntimeState.ERRORED, str(exc)) + raise + else: + self._set_status(RuntimeState.STOPPED) + + def shutdown(self) -> None: + with self.__status_lock: + state = self.__status.state + if not state.is_terminal: + self.__status = RuntimeStatus( + RuntimeState.STOPPED if state is RuntimeState.IDLE else RuntimeState.STOPPING + ) + + self._shutdown() + @classmethod @abstractmethod def build(cls, config: PipelineConfig) -> Self: @@ -145,14 +232,14 @@ def broadcast( raise NotImplementedError @abstractmethod - def run(self) -> None: + def _run(self) -> None: """ Starts the pipeline """ raise NotImplementedError @abstractmethod - def shutdown(self) -> None: + def _shutdown(self) -> None: """ Cleanly shutdown the application. """ diff --git a/sentry_streams/sentry_streams/control.py b/sentry_streams/sentry_streams/control.py new file mode 100644 index 00000000..5a7f9c43 --- /dev/null +++ b/sentry_streams/sentry_streams/control.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import logging +import threading +from typing import Any + +from sentry_streams.adapters.stream_adapter import ( + RuntimeState, + RuntimeStatus, + StreamAdapter, +) + +logger = logging.getLogger(__name__) + + +class PipelineStateError(RuntimeError): + """ + Raised when a request conflicts with the pipeline's state. + Used to differentiate between a fatal error and rejected request. + """ + + +class PipelineController: + """ + Control panel for one streaming pipeline. Starts the adapter's blocking run loop + in a background thread and allows concurrent requests to start/stop the pipeline. + + A pipeline is single-use: after it stops or fails, it cannot be restarted. + A replacement deployment must create a new controller and adapter. + """ + + def __init__(self, runtime: StreamAdapter[Any, Any]) -> None: + self._runtime = runtime + self._lock = threading.Lock() + self._thread: threading.Thread | None = None + self._finished = threading.Event() + + @property + def snapshot(self) -> RuntimeStatus: + return self._runtime.status + + def request_start(self) -> RuntimeStatus: + """ + Ask the pipeline to start (non-blocking). + """ + with self._lock: + state = self._runtime.status.state + if state is RuntimeState.STOPPING: + raise PipelineStateError("cannot start a pipeline that is stopping") + if state.is_terminal: + raise PipelineStateError(f"cannot restart a pipeline that is {state}") + + status = self._runtime.begin_start() + if self._thread is None: + self._thread = threading.Thread( + target=self._run_runtime, + name="pipeline-run", + daemon=False, + ) + self._thread.start() + return status + + def request_stop(self) -> RuntimeStatus: + """ + Ask the pipeline to stop (non-blocking). + """ + with self._lock: + status = self._runtime.status + + if status.state is RuntimeState.STOPPING or status.is_terminal: + return status + + self._runtime.shutdown() + if self._thread is None: + self._finished.set() + + return self._runtime.status + + def wait_until_finished(self, timeout: float | None = None) -> RuntimeStatus: + """ + Wait until this controller is completely finished. + """ + self._finished.wait(timeout) + return self._runtime.status + + def wait_until_stopped(self, timeout: float | None = None) -> RuntimeStatus: + """ + Wait for a started pipeline's background thread to exit. + """ + with self._lock: + thread = self._thread + + if thread is not None: + thread.join(timeout) + + return self._runtime.status + + def _run_runtime(self) -> None: + try: + self._runtime.run() + except Exception: + logger.exception("pipeline run loop failed") + finally: + self._finished.set() diff --git a/sentry_streams/sentry_streams/dummy/dummy_adapter.py b/sentry_streams/sentry_streams/dummy/dummy_adapter.py index 9463dcc3..b7651e3e 100644 --- a/sentry_streams/sentry_streams/dummy/dummy_adapter.py +++ b/sentry_streams/sentry_streams/dummy/dummy_adapter.py @@ -33,6 +33,7 @@ class DummyAdapter(StreamAdapter[DummyInput, DummyOutput]): """ def __init__(self, _: PipelineConfig) -> None: + super().__init__() self.input_streams: list[str] = [] self.branches: list[str] = [] @@ -94,8 +95,8 @@ def router(self, step: Router[RoutingFuncReturnType, Any], stream: Any) -> Any: ret[branch.root.name] = branch return ret - def run(self) -> None: + def _run(self) -> None: pass - def shutdown(self) -> None: + def _shutdown(self) -> None: pass diff --git a/sentry_streams/sentry_streams/runner.py b/sentry_streams/sentry_streams/runner.py index de2569c7..230a145e 100644 --- a/sentry_streams/sentry_streams/runner.py +++ b/sentry_streams/sentry_streams/runner.py @@ -1,17 +1,25 @@ import logging import multiprocessing +import signal import sys -from typing import Any, Mapping, Optional, cast +import threading +from http.server import ThreadingHTTPServer +from types import FrameType +from typing import Any, Callable, Mapping, Optional, cast import click import sentry_sdk from sentry_streams.adapters.loader import load_adapter from sentry_streams.adapters.stream_adapter import ( + RuntimeState, + RuntimeStatus, RuntimeTranslator, + StreamAdapter, StreamSinkT, StreamT, ) +from sentry_streams.control import PipelineController from sentry_streams.metrics import ( DatadogMetricsConfig, MetricsConfig, @@ -23,9 +31,95 @@ WithInput, ) from sentry_streams.pipeline.validation import validate_all_branches_have_sinks +from sentry_streams.server.control_server import make_server +from sentry_streams.server.control_server import serve as serve_control_server logger = logging.getLogger(__name__) +SHUTDOWN_TIMEOUT_SEC = 60.0 + + +def _install_signal_handlers(shutdown_requested: threading.Event) -> None: + """Turn SIGINT/SIGTERM into a shutdown request.""" + + def _handle_termination(signum: int, _frame: FrameType | None) -> None: + logger.info("received signal %d; requesting pipeline shutdown", signum) + shutdown_requested.set() + + for signum in (signal.SIGINT, signal.SIGTERM): + signal.signal(signum, _handle_termination) + + +def _raise_on_error(snapshot: RuntimeStatus) -> None: + if snapshot.state is RuntimeState.ERRORED: + raise RuntimeError(snapshot.error or "pipeline run loop failed") + + +def _run_pipeline( + controller: PipelineController, + shutdown_requested: threading.Event, + server: ThreadingHTTPServer | None = None, + serve: Callable[[ThreadingHTTPServer], None] = serve_control_server, +) -> RuntimeStatus: + def _stop_on_shutdown_request() -> None: + shutdown_requested.wait() + controller.request_stop() + + shutdown_thread = threading.Thread( + target=_stop_on_shutdown_request, + name="pipeline-signal", + daemon=False, + ) + shutdown_thread.start() + + serve_thread: threading.Thread | None = None + serve_failure: list[BaseException] = [] + + if server is not None: + + def _serve_until_shutdown() -> None: + try: + serve(server) + except BaseException as exc: + serve_failure.append(exc) + logger.exception("control server failed") + finally: + shutdown_requested.set() + + serve_thread = threading.Thread( + target=_serve_until_shutdown, + name="control-server", + daemon=True, + ) + serve_thread.start() + else: + controller.request_start() + + try: + controller.wait_until_finished() + finally: + shutdown_requested.set() + shutdown_thread.join(SHUTDOWN_TIMEOUT_SEC) + controller.request_stop() + snapshot = controller.wait_until_stopped(SHUTDOWN_TIMEOUT_SEC) + if not snapshot.is_terminal: + logger.warning("pipeline did not stop within %ss, exiting anyway", SHUTDOWN_TIMEOUT_SEC) + if server is not None and serve_thread is not None: + if serve_thread.is_alive(): + server.shutdown() + serve_thread.join(SHUTDOWN_TIMEOUT_SEC) + + if serve_failure: + raise serve_failure[0] + + return controller.snapshot + + +def run_runtime(runtime: StreamAdapter[Any, Any]) -> None: + shutdown_requested = threading.Event() + _install_signal_handlers(shutdown_requested) + _raise_on_error(_run_pipeline(PipelineController(runtime), shutdown_requested)) + def iterate_edges( p_graph: Pipeline[Any], translator: RuntimeTranslator[StreamT, StreamSinkT] @@ -179,14 +273,34 @@ def run_with_config_file( NOTE: This function is separate from load_runtime_with_config_file() for a reason: - load_runtime_with_config_file() returns the runtime WITHOUT calling .run() - - This allows the Rust CLI (run.rs) to call .run() itself - - Do NOT combine these functions - it would break the Rust CLI which needs to - control when .run() is called + - This allows the Rust CLI (run.rs) to pass that runtime to run_runtime() + - Do NOT combine these functions: both CLIs need the runner-owned controller + to decide when .run() is called """ runtime = load_runtime_with_config_file( name, log_level, adapter, config, segment_id, application ) - runtime.run() + run_runtime(runtime) + + +def serve_with_config_file( + name: str, + log_level: str, + adapter: str, + config: str, + segment_id: Optional[str], + application: str, + control_host: str, + control_port: int, +) -> None: + runtime = load_runtime_with_config_file( + name, log_level, adapter, config, segment_id, application + ) + controller = PipelineController(runtime) + server = make_server(controller, control_host, control_port) + shutdown_requested = threading.Event() + _install_signal_handlers(shutdown_requested) + _raise_on_error(_run_pipeline(controller, shutdown_requested, server)) @click.command() @@ -232,6 +346,24 @@ def run_with_config_file( type=str, help="The segment id to run the pipeline for", ) +@click.option( + "--control-host", + type=str, + default=None, + help=( + "Runs in operator-controlled mode and serves the control server on this host." + "Required with --control-port." + ), +) +@click.option( + "--control-port", + type=int, + default=None, + help=( + "Runs in operator-controlled mode and serves the control server on this port." + "Required with --control-host." + ), +) @click.argument( "application", required=True, @@ -242,9 +374,25 @@ def main( adapter: str, config: str, segment_id: Optional[str], + control_host: Optional[str], + control_port: Optional[int], application: str, ) -> None: - run_with_config_file(name, log_level, adapter, config, segment_id, application) + if control_host is not None or control_port is not None: + if control_host is None or control_port is None: + raise click.UsageError("--control-host and --control-port must be provided together") + serve_with_config_file( + name, + log_level, + adapter, + config, + segment_id, + application, + control_host, + control_port, + ) + else: + run_with_config_file(name, log_level, adapter, config, segment_id, application) if __name__ == "__main__": diff --git a/sentry_streams/sentry_streams/server/__init__.py b/sentry_streams/sentry_streams/server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/sentry_streams/sentry_streams/server/control_server.py b/sentry_streams/sentry_streams/server/control_server.py new file mode 100644 index 00000000..72fe28af --- /dev/null +++ b/sentry_streams/sentry_streams/server/control_server.py @@ -0,0 +1,97 @@ +""" +Operator control server for a consumer process. +""" + +from __future__ import annotations + +import json +import logging +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, cast +from urllib.parse import urlparse + +from sentry_streams.adapters.stream_adapter import RuntimeStatus +from sentry_streams.control import ( + PipelineController, + PipelineStateError, +) + +logger = logging.getLogger(__name__) + + +class ControlHandler(BaseHTTPRequestHandler): + @property + def _controller(self) -> PipelineController: + return cast(ControlServer, self.server).controller + + def log_message(self, format: str, *args: Any) -> None: + logger.debug("control-server %s - %s", self.address_string(), format % args) + + def _respond(self, code: int, payload: dict[str, Any]) -> None: + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: + try: + path = urlparse(self.path).path + if path == "/readyz": + snapshot = self._controller.snapshot + code = 503 if snapshot.is_terminal else 200 + self._respond(code, snapshot.as_dict()) + elif path == "/status": + self._respond(200, self._controller.snapshot.as_dict()) + else: + self._respond(404, {"error": "not found"}) + except Exception as exc: + self._respond_to_failure(exc) + + def do_POST(self) -> None: + try: + path = urlparse(self.path).path + if path == "/start": + self._respond(202, self._controller.request_start().as_dict()) + elif path == "/stop": + self._respond_to_stop(self._controller.request_stop()) + else: + self._respond(404, {"error": "not found"}) + except Exception as exc: + self._respond_to_failure(exc) + + def _respond_to_stop(self, snapshot: RuntimeStatus) -> None: + code = 200 if snapshot.is_terminal else 202 + self._respond(code, snapshot.as_dict()) + + def _respond_to_failure(self, exc: Exception) -> None: + if isinstance(exc, PipelineStateError): + logger.info("control-server rejected %s: %s", self.path, exc) + self._respond(409, {"error": str(exc)}) + else: + logger.exception("control-server request failed: %s", self.path) + self._respond(500, {"error": str(exc)}) + + +class ControlServer(ThreadingHTTPServer): + def __init__(self, address: tuple[str, int], controller: PipelineController) -> None: + self.controller = controller + super().__init__(address, ControlHandler) + + +def make_server(controller: PipelineController, host: str, port: int) -> ControlServer: + return ControlServer((host, port), controller) + + +def serve(server: ThreadingHTTPServer) -> None: + host, port = server.server_address[:2] + logger.info( + "Streams control server listening on %s:%d (pipeline idle, awaiting /start)", + host, + port, + ) + try: + server.serve_forever() + finally: + server.server_close() diff --git a/sentry_streams/src/consumer.rs b/sentry_streams/src/consumer.rs index 5ce4c884..5ef2aa52 100644 --- a/sentry_streams/src/consumer.rs +++ b/sentry_streams/src/consumer.rs @@ -16,6 +16,7 @@ use crate::routes::Route; use crate::routes::RoutedValue; use crate::utils::traced_with_gil; use crate::watermark::WatermarkEmitter; +use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use rdkafka::message::{Header, Headers, OwnedHeaders}; use sentry_arroyo::backends::kafka::producer::KafkaProducer; @@ -31,7 +32,8 @@ use sentry_arroyo::processing::ProcessorHandle; use sentry_arroyo::processing::StreamProcessor; use sentry_arroyo::types::{Message, Topic}; use std::collections::HashMap; -use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; /// Default path for the healthcheck file touched when write_healthcheck is enabled. /// Matches Arroyo docs for Kubernetes liveness probes. @@ -85,9 +87,13 @@ pub struct ArroyoConsumer { steps: Vec>, + /// Set by `shutdown()` so a stop that arrives before `run()` builds the + /// processor is not lost. `run()` checks it before entering the run loop. + shutdown_requested: AtomicBool, + /// The ProcessorHandle allows the main thread to stop the StreamingProcessor /// from a different thread. - handle: Option, + handle: Mutex>, // this variable must live for the lifetime of the entire consumer. // This is a requirement of Arroyo Rust. @@ -133,7 +139,8 @@ impl ArroyoConsumer { source, schema, steps: Vec::new(), - handle: None, + shutdown_requested: AtomicBool::new(false), + handle: Mutex::new(None), concurrency_config: Arc::new(ConcurrencyConfig::new(1)), step_concurrency_configs: HashMap::new(), metric_config, @@ -162,8 +169,9 @@ impl ArroyoConsumer { /// Runs the consumer. /// This method is blocking and will run until the consumer - /// is stopped via SIGTERM or SIGINT. - fn run(&mut self) { + /// is stopped via shutdown(). Signals are handled by the Python + /// runner, which turns them into a shutdown call. + fn run(&self, py: Python<'_>) -> PyResult<()> { tracing_subscriber::fmt::init(); println!("Running Arroyo Consumer..."); @@ -196,25 +204,40 @@ impl ArroyoConsumer { let processor = StreamProcessor::with_kafka(config, factory, Topic::new(&self.topic), dlq_policy); - self.handle = Some(processor.get_handle()); - let mut handle = processor.get_handle(); - ctrlc::set_handler(move || { - println!("\nCtrl+C pressed!"); - handle.signal_shutdown(); - }) - .expect("Error setting Ctrl+C handler"); + *self.handle.lock().unwrap() = Some(processor.get_handle()); + + if self.shutdown_requested.load(Ordering::Acquire) { + processor.get_handle().signal_shutdown(); + } + + // The GIL is released around the run loop so that a Python control + // thread can call shutdown while this consumer is running: - if let Err(e) = processor.run() { - tracing::error!("StreamProcessor error: {:?}", e); - sentry::capture_error(&e); + let run_error = py.detach(|| match processor.run() { + Ok(()) => None, + Err(error) => { + tracing::error!("StreamProcessor error: {:?}", error); + sentry::capture_error(&error); + Some(error.to_string()) + } + }); + + *self.handle.lock().unwrap() = None; + + if let Some(error) = run_error { + return Err(PyRuntimeError::new_err(format!( + "StreamProcessor error: {error}" + ))); } + + Ok(()) } - fn shutdown(&mut self) { - match self.handle.take() { - Some(mut handle) => handle.signal_shutdown(), - None => println!("No handle to shut down."), + fn shutdown(&self) { + self.shutdown_requested.store(true, Ordering::Release); + if let Some(mut handle) = self.handle.lock().unwrap().take() { + handle.signal_shutdown(); } } } diff --git a/sentry_streams/src/run.rs b/sentry_streams/src/run.rs index 85031758..11256815 100644 --- a/sentry_streams/src/run.rs +++ b/sentry_streams/src/run.rs @@ -84,11 +84,11 @@ pub fn run(args: Args) -> Result<(), Box> { })?; traced_with_gil!(|py| { - runtime - .bind(py) - .call_method0("run") - .expect("Unable to start runtime"); - }); + py.import("sentry_streams.runner")? + .getattr("run_runtime")? + .call1((runtime,))?; + PyResult::Ok(()) + })?; Ok(()) } diff --git a/sentry_streams/tests/adapters/arroyo/test_adapter.py b/sentry_streams/tests/adapters/arroyo/test_adapter.py index 5094d79e..28ca596c 100644 --- a/sentry_streams/tests/adapters/arroyo/test_adapter.py +++ b/sentry_streams/tests/adapters/arroyo/test_adapter.py @@ -12,7 +12,10 @@ ArroyoAdapter, StreamSources, ) -from sentry_streams.adapters.stream_adapter import RuntimeTranslator +from sentry_streams.adapters.stream_adapter import ( + RuntimeState, + RuntimeTranslator, +) from sentry_streams.config_types import KafkaConsumerConfig from sentry_streams.pipeline.pipeline import ( Pipeline, @@ -99,3 +102,61 @@ def test_adapter( msg2 = broker.consume(Partition(topic, 0), 1) assert msg2 is not None and msg2.payload.value == json.dumps(transformed_metric).encode("utf-8") assert broker.consume(Partition(topic, 0), 2) is None + + +def _adapter_with_stub_processor(processor: mock.Mock) -> ArroyoAdapter: + adapter = ArroyoAdapter({}) + setattr(adapter, "_ArroyoAdapter__consumers", {"source": mock.Mock()}) + setattr( + adapter, + "create_processors", + lambda: setattr(adapter, "_ArroyoAdapter__processors", {"source": processor}), + ) + return adapter + + +def test_shutdown_while_building_the_processor_is_applied_to_it() -> None: + processor = mock.Mock() + adapter = _adapter_with_stub_processor(processor) + processor.run.side_effect = lambda: ( + adapter.status.state is RuntimeState.STOPPING + or pytest.fail(f"unexpected status: {adapter.status.state}") + ) + + def shutdown_while_building() -> None: + adapter.shutdown() + setattr(adapter, "_ArroyoAdapter__processors", {"source": processor}) + + setattr(adapter, "create_processors", shutdown_while_building) + adapter.begin_start() + adapter.run() + + processor.signal_shutdown.assert_called_once_with() + processor.run.assert_called_once_with() + assert adapter.status.state is RuntimeState.STOPPED + + +def test_shutdown_before_start_never_builds_a_processor() -> None: + processor = mock.Mock() + adapter = _adapter_with_stub_processor(processor) + + adapter.shutdown() + assert adapter.status.state is RuntimeState.STOPPED + + adapter.run() + + processor.run.assert_not_called() + assert adapter.status.state is RuntimeState.STOPPED + + +def test_shutdown_during_startup_never_builds_a_processor() -> None: + processor = mock.Mock() + adapter = _adapter_with_stub_processor(processor) + + adapter.begin_start() + adapter.shutdown() + adapter.shutdown() + adapter.run() + + processor.run.assert_not_called() + assert adapter.status.state is RuntimeState.STOPPED diff --git a/sentry_streams/tests/adapters/fake_adapter.py b/sentry_streams/tests/adapters/fake_adapter.py new file mode 100644 index 00000000..7f6328d8 --- /dev/null +++ b/sentry_streams/tests/adapters/fake_adapter.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import threading +from typing import Any, Callable, Mapping, Self, Type + +from sentry_streams.adapters.stream_adapter import ( + PipelineConfig, + RuntimeState, + StreamAdapter, +) +from sentry_streams.pipeline.function_template import InputType, OutputType +from sentry_streams.pipeline.pipeline import ( + Broadcast, + ComplexStep, + Filter, + FlatMap, + Map, + Reduce, + Router, + RoutingFuncReturnType, + Sink, + Source, +) +from sentry_streams.pipeline.window import MeasurementUnit + + +class FakeAdapter(StreamAdapter[Any, Any]): + def __init__(self, fail: bool = False, block_before_consuming: bool = False) -> None: + super().__init__() + self._fail = fail + self._block_before_consuming = block_before_consuming + self._stop = threading.Event() + self.allow_consume = threading.Event() + self.run_started = threading.Event() + self.run_finished = threading.Event() + self.run_calls = 0 + self.shutdown_calls = 0 + + def _run(self) -> None: + self.run_calls += 1 + self.run_started.set() + if self._fail: + raise RuntimeError("runtime failed") + if self._block_before_consuming: + assert self.allow_consume.wait(3.0) + self._set_status(RuntimeState.CONSUMING) + self._stop.wait() + self.run_finished.set() + + def _shutdown(self) -> None: + self.shutdown_calls += 1 + self._stop.set() + + @classmethod + def build(cls, config: PipelineConfig) -> Self: + return cls() + + def complex_step_override( + self, + ) -> dict[Type[ComplexStep[Any, Any]], Callable[[ComplexStep[Any, Any]], Any]]: + return {} + + def source(self, step: Source[Any]) -> Any: + raise NotImplementedError + + def sink(self, step: Sink[Any], stream: Any) -> Any: + raise NotImplementedError + + def map(self, step: Map[Any, Any], stream: Any) -> Any: + raise NotImplementedError + + def flat_map(self, step: FlatMap[Any, Any], stream: Any) -> Any: + raise NotImplementedError + + def filter(self, step: Filter[Any], stream: Any) -> Any: + raise NotImplementedError + + def reduce(self, step: Reduce[MeasurementUnit, InputType, OutputType], stream: Any) -> Any: + raise NotImplementedError + + def router(self, step: Router[RoutingFuncReturnType, Any], stream: Any) -> Mapping[str, Any]: + raise NotImplementedError + + def broadcast(self, step: Broadcast[Any], stream: Any) -> Mapping[str, Any]: + raise NotImplementedError diff --git a/sentry_streams/tests/test_control.py b/sentry_streams/tests/test_control.py new file mode 100644 index 00000000..40bfd1b8 --- /dev/null +++ b/sentry_streams/tests/test_control.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import threading +import time + +import pytest + +from sentry_streams.adapters.stream_adapter import RuntimeState, RuntimeStatus +from sentry_streams.control import PipelineController, PipelineStateError +from tests.adapters.fake_adapter import FakeAdapter + + +def _wait_for_state( + controller: PipelineController, state: RuntimeState, timeout: float = 3.0 +) -> RuntimeStatus: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + snapshot = controller.snapshot + if snapshot.state is state: + return snapshot + time.sleep(0.01) + raise AssertionError(f"pipeline is {controller.snapshot.state}, expected {state}") + + +def _stop(controller: PipelineController) -> None: + controller.request_stop() + controller.wait_until_stopped(3.0) + + +def test_start_and_stop_are_non_blocking() -> None: + runtime = FakeAdapter() + controller = PipelineController(runtime) + try: + assert controller.request_start().state is RuntimeState.STARTING + assert controller.request_start().state in ( + RuntimeState.STARTING, + RuntimeState.CONSUMING, + ) + assert runtime.run_started.wait(3.0) + _wait_for_state(controller, RuntimeState.CONSUMING) + + assert controller.request_stop().state is RuntimeState.STOPPING + assert controller.request_stop().state in ( + RuntimeState.STOPPING, + RuntimeState.STOPPED, + ) + assert controller.wait_until_stopped(3.0).state is RuntimeState.STOPPED + assert runtime.run_calls == 1 + assert runtime.shutdown_calls == 1 + finally: + _stop(controller) + + +def test_controller_reports_adapter_runtime_status() -> None: + runtime = FakeAdapter(block_before_consuming=True) + controller = PipelineController(runtime) + try: + controller.request_start() + assert runtime.run_started.wait(3.0) + assert controller.snapshot.state is RuntimeState.STARTING + + runtime.allow_consume.set() + _wait_for_state(controller, RuntimeState.CONSUMING) + finally: + runtime.allow_consume.set() + _stop(controller) + + +def test_stop_during_start_is_not_lost() -> None: + runtime = FakeAdapter() + controller = PipelineController(runtime) + try: + controller.request_start() + assert controller.request_stop().state is RuntimeState.STOPPING + + assert controller.wait_until_stopped(3.0).state is RuntimeState.STOPPED + assert runtime.shutdown_calls == 1 + finally: + _stop(controller) + + +def test_stop_before_start_shuts_the_runtime_down() -> None: + runtime = FakeAdapter() + controller = PipelineController(runtime) + try: + assert controller.request_stop().state is RuntimeState.STOPPED + assert runtime.shutdown_calls == 1 + assert runtime.run_calls == 0 + finally: + _stop(controller) + + +def test_waiting_for_an_idle_pipeline_ends_when_it_is_stopped() -> None: + controller = PipelineController(FakeAdapter()) + try: + stopper = threading.Thread(target=controller.request_stop) + stopper.start() + assert controller.wait_until_finished(3.0).state is RuntimeState.STOPPED + stopper.join(3.0) + finally: + _stop(controller) + + +def test_waiting_for_a_running_pipeline_ends_when_its_run_loop_exits() -> None: + runtime = FakeAdapter() + controller = PipelineController(runtime) + try: + controller.request_start() + assert runtime.run_started.wait(3.0) + controller.request_stop() + + assert controller.wait_until_finished(3.0).state is RuntimeState.STOPPED + assert runtime.run_finished.is_set() + finally: + _stop(controller) + + +def test_stopping_a_failed_runtime_keeps_the_failure() -> None: + runtime = FakeAdapter(fail=True) + controller = PipelineController(runtime) + try: + assert controller.request_start().state is RuntimeState.STARTING + assert controller.wait_until_stopped(3.0).error == "runtime failed" + + controller.request_stop() + _stop(controller) + + snapshot = controller.snapshot + assert snapshot.state is RuntimeState.ERRORED + assert snapshot.error == "runtime failed" + finally: + _stop(controller) + + +def test_stopped_runtime_cannot_restart() -> None: + runtime = FakeAdapter() + controller = PipelineController(runtime) + try: + controller.request_start() + assert runtime.run_started.wait(3.0) + controller.request_stop() + controller.wait_until_stopped(3.0) + + with pytest.raises(PipelineStateError, match="cannot restart"): + controller.request_start() + assert runtime.run_calls == 1 + assert runtime.shutdown_calls == 1 + finally: + _stop(controller) + + +def test_shutdown_before_the_run_loop_never_starts() -> None: + runtime = FakeAdapter() + runtime.begin_start() + runtime.shutdown() + + runtime.run() + + assert runtime.status.state is RuntimeState.STOPPED + assert runtime.run_calls == 0 diff --git a/sentry_streams/tests/test_control_server.py b/sentry_streams/tests/test_control_server.py new file mode 100644 index 00000000..3f787fde --- /dev/null +++ b/sentry_streams/tests/test_control_server.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import json +import threading +import time +import urllib.error +import urllib.request +from typing import Any, Callable + +from sentry_streams.adapters.stream_adapter import RuntimeState +from sentry_streams.control import PipelineController +from sentry_streams.server.control_server import make_server +from tests.adapters.fake_adapter import FakeAdapter + + +def _wait_for(predicate: Callable[[], bool], timeout: float = 3.0) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(0.01) + return False + + +def _request(port: int, path: str, method: str) -> tuple[int, dict[str, Any]]: + data = b"" if method == "POST" else None + req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", method=method, data=data) + try: + with urllib.request.urlopen(req, timeout=3.0) as resp: + raw = resp.read() + return resp.status, (json.loads(raw) if raw else {}) + except urllib.error.HTTPError as exc: + raw = exc.read() + return exc.code, (json.loads(raw) if raw else {}) + + +def _stop(controller: PipelineController) -> None: + controller.request_stop() + controller.wait_until_stopped(3.0) + + +def test_control_server_endpoints() -> None: + runtime = FakeAdapter() + controller = PipelineController(runtime) + server = make_server(controller, "127.0.0.1", 0) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + code, body = _request(port, "/readyz", "GET") + assert code == 200 and body["state"] == RuntimeState.IDLE + + code, _ = _request(port, "/does-not-exist", "GET") + assert code == 404 + + code, _ = _request(port, "/start", "POST") + assert code == 202 + assert _wait_for( + lambda: _request(port, "/status", "GET")[1]["state"] == RuntimeState.CONSUMING + ) + + code, body = _request(port, "/stop", "POST") + assert code == 202 and body["state"] == RuntimeState.STOPPING + assert controller.wait_until_stopped(3.0).state is RuntimeState.STOPPED + assert _request(port, "/does-not-exist", "POST")[0] == 404 + assert runtime.shutdown_calls == 1 + finally: + server.shutdown() + server.server_close() + thread.join(timeout=3.0) + _stop(controller) + + +def test_readyz_reports_runtime_failure() -> None: + runtime = FakeAdapter(fail=True) + controller = PipelineController(runtime) + server = make_server(controller, "127.0.0.1", 0) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + assert _request(port, "/start", "POST")[0] == 202 + assert controller.wait_until_stopped(3.0).state is RuntimeState.ERRORED + code, body = _request(port, "/readyz", "GET") + assert code == 503 + assert body["state"] == RuntimeState.ERRORED + assert body["error"] == "runtime failed" + finally: + server.shutdown() + server.server_close() + thread.join(timeout=3.0) + _stop(controller) diff --git a/sentry_streams/tests/test_runner.py b/sentry_streams/tests/test_runner.py index 9f34e9fb..58182510 100644 --- a/sentry_streams/tests/test_runner.py +++ b/sentry_streams/tests/test_runner.py @@ -1,17 +1,30 @@ +import os +import signal +import threading from enum import Enum from typing import Any, cast import pytest from sentry_streams.adapters.loader import load_adapter -from sentry_streams.adapters.stream_adapter import PipelineConfig, RuntimeTranslator +from sentry_streams.adapters.stream_adapter import ( + PipelineConfig, + RuntimeState, + RuntimeTranslator, +) +from sentry_streams.control import PipelineController from sentry_streams.dummy.dummy_adapter import DummyAdapter from sentry_streams.pipeline import Map, PredicateFilter, branch, streaming_source from sentry_streams.pipeline.pipeline import ( DevNullSink, Pipeline, ) -from sentry_streams.runner import iterate_edges +from sentry_streams.runner import ( + _install_signal_handlers, + _run_pipeline, + iterate_edges, +) +from tests.adapters.fake_adapter import FakeAdapter class RouterBranch(Enum): @@ -94,3 +107,26 @@ def test_iterate_edges(create_pipeline: Pipeline[bytes]) -> None: "map4_segment", "map5_segment", ] + + +@pytest.mark.parametrize("signum", [signal.SIGINT, signal.SIGTERM]) +def test_run_pipeline_terminal_signals(signum: int) -> None: + runtime = FakeAdapter() + shutdown_requested = threading.Event() + _install_signal_handlers(shutdown_requested) + + def send_signal() -> None: + assert runtime.run_started.wait(3.0) + os.kill(os.getpid(), signum) + + signal_thread = threading.Thread(target=send_signal) + signal_thread.start() + try: + snapshot = _run_pipeline(PipelineController(runtime), shutdown_requested) + finally: + signal_thread.join(timeout=3.0) + signal.signal(signal.SIGINT, signal.default_int_handler) + signal.signal(signal.SIGTERM, signal.SIG_DFL) + + assert snapshot.state is RuntimeState.STOPPED + assert runtime.shutdown_calls == 1 From 60d871065b1305d03638ae42f0c5a386624151d4 Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Fri, 31 Jul 2026 11:37:05 -0700 Subject: [PATCH 2/6] ref(controller): exception semantics --- .../sentry_streams/adapters/stream_adapter.py | 71 ++++++++++++------- sentry_streams/sentry_streams/control.py | 23 +----- sentry_streams/sentry_streams/runner.py | 4 +- .../sentry_streams/server/control_server.py | 10 +-- .../tests/adapters/arroyo/test_adapter.py | 3 +- sentry_streams/tests/test_control.py | 23 +++--- sentry_streams/tests/test_runner.py | 11 +++ 7 files changed, 82 insertions(+), 63 deletions(-) diff --git a/sentry_streams/sentry_streams/adapters/stream_adapter.py b/sentry_streams/sentry_streams/adapters/stream_adapter.py index f0e83700..29954bdc 100644 --- a/sentry_streams/sentry_streams/adapters/stream_adapter.py +++ b/sentry_streams/sentry_streams/adapters/stream_adapter.py @@ -44,6 +44,12 @@ StreamSinkT = TypeVar("StreamSinkT") +class RuntimeStateError(RuntimeError): + """ + Raised when an operation conflicts with the runtime's lifecycle state. + """ + + class RuntimeState(StrEnum): IDLE = "idle" STARTING = "starting" @@ -56,32 +62,36 @@ class RuntimeState(StrEnum): def is_terminal(self) -> bool: return self in (RuntimeState.STOPPED, RuntimeState.ERRORED) + @property + def rank(self) -> int: + if self is RuntimeState.IDLE: + return 0 + if self is RuntimeState.STARTING: + return 1 + if self is RuntimeState.CONSUMING: + return 2 + if self is RuntimeState.STOPPING: + return 3 + return 4 -# To simplify lifecycle and state transitions we can use a forward-only ranking. -# A higher-ranked state should realistically never transition to a lower one. -# Would need to be updated if we allow restarting (but we don't for now). - -_STATE_RANKS: Mapping[RuntimeState, int] = { - RuntimeState.IDLE: 0, - RuntimeState.STARTING: 1, - RuntimeState.CONSUMING: 2, - RuntimeState.STOPPING: 3, - RuntimeState.STOPPED: 4, - RuntimeState.ERRORED: 4, -} + def can_transition_to(self, state: RuntimeState) -> bool: + return self.rank < state.rank @dataclass(frozen=True) class RuntimeStatus: state: RuntimeState - error: str | None = None + error: Exception | None = None @property def is_terminal(self) -> bool: return self.state.is_terminal def as_dict(self) -> dict[str, str | None]: - return {"state": self.state.value, "error": self.error} + return { + "state": self.state.value, + "error": str(self.error) if self.error is not None else None, + } class StreamAdapter(ABC, Generic[StreamT, StreamSinkT]): @@ -100,40 +110,49 @@ def status(self) -> RuntimeStatus: with self.__status_lock: return self.__status - def _set_status(self, state: RuntimeState, error: str | None = None) -> None: + def _set_status(self, state: RuntimeState, error: Exception | None = None) -> None: with self.__status_lock: - if _STATE_RANKS[state] <= _STATE_RANKS[self.__status.state]: + if not self.__status.state.can_transition_to(state): return self.__status = RuntimeStatus(state, error) def begin_start(self) -> RuntimeStatus: with self.__status_lock: - if self.__status.state is RuntimeState.IDLE: - self.__status = RuntimeStatus(RuntimeState.STARTING) + state = self.__status.state + if state.is_terminal: + raise RuntimeStateError(f"cannot restart runtime that is {state}") + if state is not RuntimeState.IDLE: + raise RuntimeStateError(f"cannot start runtime while it is {state}") + self.__status = RuntimeStatus(RuntimeState.STARTING) return self.__status def run(self) -> None: - if self.begin_start().state is not RuntimeState.STARTING: + state = self.status.state + if state is RuntimeState.STOPPING: self._set_status(RuntimeState.STOPPED) return + if state is not RuntimeState.STARTING: + raise RuntimeStateError(f"cannot run runtime while it is {state}") try: self._run() - except BaseException as exc: - self._set_status(RuntimeState.ERRORED, str(exc)) + except Exception as exc: + self._set_status(RuntimeState.ERRORED, exc) raise else: self._set_status(RuntimeState.STOPPED) - def shutdown(self) -> None: + def shutdown(self) -> RuntimeStatus: with self.__status_lock: state = self.__status.state - if not state.is_terminal: - self.__status = RuntimeStatus( - RuntimeState.STOPPED if state is RuntimeState.IDLE else RuntimeState.STOPPING - ) + if state is RuntimeState.STOPPING or state.is_terminal: + return self.__status + self.__status = RuntimeStatus( + RuntimeState.STOPPED if state is RuntimeState.IDLE else RuntimeState.STOPPING + ) self._shutdown() + return self.status @classmethod @abstractmethod diff --git a/sentry_streams/sentry_streams/control.py b/sentry_streams/sentry_streams/control.py index 5a7f9c43..4caf8c60 100644 --- a/sentry_streams/sentry_streams/control.py +++ b/sentry_streams/sentry_streams/control.py @@ -5,7 +5,6 @@ from typing import Any from sentry_streams.adapters.stream_adapter import ( - RuntimeState, RuntimeStatus, StreamAdapter, ) @@ -13,13 +12,6 @@ logger = logging.getLogger(__name__) -class PipelineStateError(RuntimeError): - """ - Raised when a request conflicts with the pipeline's state. - Used to differentiate between a fatal error and rejected request. - """ - - class PipelineController: """ Control panel for one streaming pipeline. Starts the adapter's blocking run loop @@ -44,12 +36,6 @@ def request_start(self) -> RuntimeStatus: Ask the pipeline to start (non-blocking). """ with self._lock: - state = self._runtime.status.state - if state is RuntimeState.STOPPING: - raise PipelineStateError("cannot start a pipeline that is stopping") - if state.is_terminal: - raise PipelineStateError(f"cannot restart a pipeline that is {state}") - status = self._runtime.begin_start() if self._thread is None: self._thread = threading.Thread( @@ -65,16 +51,11 @@ def request_stop(self) -> RuntimeStatus: Ask the pipeline to stop (non-blocking). """ with self._lock: - status = self._runtime.status - - if status.state is RuntimeState.STOPPING or status.is_terminal: - return status - - self._runtime.shutdown() + status = self._runtime.shutdown() if self._thread is None: self._finished.set() - return self._runtime.status + return status def wait_until_finished(self, timeout: float | None = None) -> RuntimeStatus: """ diff --git a/sentry_streams/sentry_streams/runner.py b/sentry_streams/sentry_streams/runner.py index 230a145e..0390b46f 100644 --- a/sentry_streams/sentry_streams/runner.py +++ b/sentry_streams/sentry_streams/runner.py @@ -52,7 +52,9 @@ def _handle_termination(signum: int, _frame: FrameType | None) -> None: def _raise_on_error(snapshot: RuntimeStatus) -> None: if snapshot.state is RuntimeState.ERRORED: - raise RuntimeError(snapshot.error or "pipeline run loop failed") + if snapshot.error is not None: + raise snapshot.error + raise RuntimeError("pipeline run loop failed") def _run_pipeline( diff --git a/sentry_streams/sentry_streams/server/control_server.py b/sentry_streams/sentry_streams/server/control_server.py index 72fe28af..15be1987 100644 --- a/sentry_streams/sentry_streams/server/control_server.py +++ b/sentry_streams/sentry_streams/server/control_server.py @@ -10,11 +10,11 @@ from typing import Any, cast from urllib.parse import urlparse -from sentry_streams.adapters.stream_adapter import RuntimeStatus -from sentry_streams.control import ( - PipelineController, - PipelineStateError, +from sentry_streams.adapters.stream_adapter import ( + RuntimeStateError, + RuntimeStatus, ) +from sentry_streams.control import PipelineController logger = logging.getLogger(__name__) @@ -66,7 +66,7 @@ def _respond_to_stop(self, snapshot: RuntimeStatus) -> None: self._respond(code, snapshot.as_dict()) def _respond_to_failure(self, exc: Exception) -> None: - if isinstance(exc, PipelineStateError): + if isinstance(exc, RuntimeStateError): logger.info("control-server rejected %s: %s", self.path, exc) self._respond(409, {"error": str(exc)}) else: diff --git a/sentry_streams/tests/adapters/arroyo/test_adapter.py b/sentry_streams/tests/adapters/arroyo/test_adapter.py index 28ca596c..a7e43a02 100644 --- a/sentry_streams/tests/adapters/arroyo/test_adapter.py +++ b/sentry_streams/tests/adapters/arroyo/test_adapter.py @@ -143,7 +143,8 @@ def test_shutdown_before_start_never_builds_a_processor() -> None: adapter.shutdown() assert adapter.status.state is RuntimeState.STOPPED - adapter.run() + with pytest.raises(RuntimeError, match="cannot run runtime while it is stopped"): + adapter.run() processor.run.assert_not_called() assert adapter.status.state is RuntimeState.STOPPED diff --git a/sentry_streams/tests/test_control.py b/sentry_streams/tests/test_control.py index 40bfd1b8..e4cb4890 100644 --- a/sentry_streams/tests/test_control.py +++ b/sentry_streams/tests/test_control.py @@ -5,8 +5,12 @@ import pytest -from sentry_streams.adapters.stream_adapter import RuntimeState, RuntimeStatus -from sentry_streams.control import PipelineController, PipelineStateError +from sentry_streams.adapters.stream_adapter import ( + RuntimeState, + RuntimeStateError, + RuntimeStatus, +) +from sentry_streams.control import PipelineController from tests.adapters.fake_adapter import FakeAdapter @@ -32,10 +36,8 @@ def test_start_and_stop_are_non_blocking() -> None: controller = PipelineController(runtime) try: assert controller.request_start().state is RuntimeState.STARTING - assert controller.request_start().state in ( - RuntimeState.STARTING, - RuntimeState.CONSUMING, - ) + with pytest.raises(RuntimeStateError, match="cannot start runtime"): + controller.request_start() assert runtime.run_started.wait(3.0) _wait_for_state(controller, RuntimeState.CONSUMING) @@ -120,14 +122,17 @@ def test_stopping_a_failed_runtime_keeps_the_failure() -> None: controller = PipelineController(runtime) try: assert controller.request_start().state is RuntimeState.STARTING - assert controller.wait_until_stopped(3.0).error == "runtime failed" + error = controller.wait_until_stopped(3.0).error + assert isinstance(error, RuntimeError) + assert str(error) == "runtime failed" + assert error.__traceback__ is not None controller.request_stop() _stop(controller) snapshot = controller.snapshot assert snapshot.state is RuntimeState.ERRORED - assert snapshot.error == "runtime failed" + assert snapshot.error is error finally: _stop(controller) @@ -141,7 +146,7 @@ def test_stopped_runtime_cannot_restart() -> None: controller.request_stop() controller.wait_until_stopped(3.0) - with pytest.raises(PipelineStateError, match="cannot restart"): + with pytest.raises(RuntimeStateError, match="cannot restart"): controller.request_start() assert runtime.run_calls == 1 assert runtime.shutdown_calls == 1 diff --git a/sentry_streams/tests/test_runner.py b/sentry_streams/tests/test_runner.py index 58182510..7fe72fbe 100644 --- a/sentry_streams/tests/test_runner.py +++ b/sentry_streams/tests/test_runner.py @@ -10,6 +10,7 @@ from sentry_streams.adapters.stream_adapter import ( PipelineConfig, RuntimeState, + RuntimeStatus, RuntimeTranslator, ) from sentry_streams.control import PipelineController @@ -21,6 +22,7 @@ ) from sentry_streams.runner import ( _install_signal_handlers, + _raise_on_error, _run_pipeline, iterate_edges, ) @@ -130,3 +132,12 @@ def send_signal() -> None: assert snapshot.state is RuntimeState.STOPPED assert runtime.shutdown_calls == 1 + + +def test_raise_on_error_reraises_the_stored_exception() -> None: + error = ValueError("runtime failed") + + with pytest.raises(ValueError) as exc_info: + _raise_on_error(RuntimeStatus(RuntimeState.ERRORED, error)) + + assert exc_info.value is error From 07c338fa468131ce072c481dc1be1922ccd8be0e Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Fri, 31 Jul 2026 11:42:38 -0700 Subject: [PATCH 3/6] ref(controller): use ThreadPoolExecutor + Future --- sentry_streams/sentry_streams/control.py | 47 +++++++++++------------- sentry_streams/tests/test_control.py | 12 ++++++ 2 files changed, 33 insertions(+), 26 deletions(-) diff --git a/sentry_streams/sentry_streams/control.py b/sentry_streams/sentry_streams/control.py index 4caf8c60..301e7b05 100644 --- a/sentry_streams/sentry_streams/control.py +++ b/sentry_streams/sentry_streams/control.py @@ -1,7 +1,7 @@ from __future__ import annotations -import logging import threading +from concurrent.futures import Future, ThreadPoolExecutor from typing import Any from sentry_streams.adapters.stream_adapter import ( @@ -9,8 +9,6 @@ StreamAdapter, ) -logger = logging.getLogger(__name__) - class PipelineController: """ @@ -24,7 +22,8 @@ class PipelineController: def __init__(self, runtime: StreamAdapter[Any, Any]) -> None: self._runtime = runtime self._lock = threading.Lock() - self._thread: threading.Thread | None = None + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pipeline-run") + self._run_future: Future[None] | None = None self._finished = threading.Event() @property @@ -37,13 +36,9 @@ def request_start(self) -> RuntimeStatus: """ with self._lock: status = self._runtime.begin_start() - if self._thread is None: - self._thread = threading.Thread( - target=self._run_runtime, - name="pipeline-run", - daemon=False, - ) - self._thread.start() + if self._run_future is None: + self._run_future = self._executor.submit(self._runtime.run) + self._run_future.add_done_callback(self._run_finished) return status def request_stop(self) -> RuntimeStatus: @@ -52,8 +47,9 @@ def request_stop(self) -> RuntimeStatus: """ with self._lock: status = self._runtime.shutdown() - if self._thread is None: + if self._run_future is None: self._finished.set() + self._executor.shutdown(wait=False) return status @@ -61,25 +57,24 @@ def wait_until_finished(self, timeout: float | None = None) -> RuntimeStatus: """ Wait until this controller is completely finished. """ - self._finished.wait(timeout) + if not self._finished.wait(timeout): + return self._runtime.status + + with self._lock: + run_future = self._run_future + + if run_future is not None: + run_future.result() + return self._runtime.status def wait_until_stopped(self, timeout: float | None = None) -> RuntimeStatus: """ Wait for a started pipeline's background thread to exit. """ - with self._lock: - thread = self._thread - - if thread is not None: - thread.join(timeout) - + self._finished.wait(timeout) return self._runtime.status - def _run_runtime(self) -> None: - try: - self._runtime.run() - except Exception: - logger.exception("pipeline run loop failed") - finally: - self._finished.set() + def _run_finished(self, _future: Future[None]) -> None: + self._finished.set() + self._executor.shutdown(wait=False) diff --git a/sentry_streams/tests/test_control.py b/sentry_streams/tests/test_control.py index e4cb4890..d393564c 100644 --- a/sentry_streams/tests/test_control.py +++ b/sentry_streams/tests/test_control.py @@ -117,6 +117,18 @@ def test_waiting_for_a_running_pipeline_ends_when_its_run_loop_exits() -> None: _stop(controller) +def test_wait_until_finished_reraises_the_run_loop_exception() -> None: + runtime = FakeAdapter(fail=True) + controller = PipelineController(runtime) + controller.request_start() + + with pytest.raises(RuntimeError, match="runtime failed") as exc_info: + controller.wait_until_finished(3.0) + + assert exc_info.value is runtime.status.error + assert exc_info.value.__traceback__ is not None + + def test_stopping_a_failed_runtime_keeps_the_failure() -> None: runtime = FakeAdapter(fail=True) controller = PipelineController(runtime) From d9dd62ff35c89a397e103cd398935dc40175748d Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Fri, 31 Jul 2026 11:44:28 -0700 Subject: [PATCH 4/6] ref(controller): redundant request_stop --- sentry_streams/sentry_streams/runner.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sentry_streams/sentry_streams/runner.py b/sentry_streams/sentry_streams/runner.py index 0390b46f..cee7b19a 100644 --- a/sentry_streams/sentry_streams/runner.py +++ b/sentry_streams/sentry_streams/runner.py @@ -102,7 +102,6 @@ def _serve_until_shutdown() -> None: finally: shutdown_requested.set() shutdown_thread.join(SHUTDOWN_TIMEOUT_SEC) - controller.request_stop() snapshot = controller.wait_until_stopped(SHUTDOWN_TIMEOUT_SEC) if not snapshot.is_terminal: logger.warning("pipeline did not stop within %ss, exiting anyway", SHUTDOWN_TIMEOUT_SEC) From 774e218de5fdb0fef5251e091b6941c118baafb6 Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Fri, 31 Jul 2026 12:02:18 -0700 Subject: [PATCH 5/6] chore(server): better docstring --- .../sentry_streams/server/control_server.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/sentry_streams/sentry_streams/server/control_server.py b/sentry_streams/sentry_streams/server/control_server.py index 15be1987..1f764918 100644 --- a/sentry_streams/sentry_streams/server/control_server.py +++ b/sentry_streams/sentry_streams/server/control_server.py @@ -1,7 +1,3 @@ -""" -Operator control server for a consumer process. -""" - from __future__ import annotations import json @@ -75,6 +71,22 @@ def _respond_to_failure(self, exc: Exception) -> None: class ControlServer(ThreadingHTTPServer): + """ + Expose lifecycle control for a streaming consumer process over HTTP. + + In operator-controlled mode, the consumer process loads its pipeline but leaves + it idle. This server wraps its PipelineController so an external manager + (like the operator) can manually start consuming and observe/change state. + + - POST /start: Start consuming. + - POST /stop: Stop consuming. + + - GET /status: Get the current state. + - GET /readyz: Get a readiness response. + + Invalid state transitions are rejected. + """ + def __init__(self, address: tuple[str, int], controller: PipelineController) -> None: self.controller = controller super().__init__(address, ControlHandler) From 4937081cb63d772efce3674916ec439cde60f5e7 Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Fri, 31 Jul 2026 12:06:07 -0700 Subject: [PATCH 6/6] chore(controller): better docstring --- sentry_streams/sentry_streams/control.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sentry_streams/sentry_streams/control.py b/sentry_streams/sentry_streams/control.py index 301e7b05..458fe9bf 100644 --- a/sentry_streams/sentry_streams/control.py +++ b/sentry_streams/sentry_streams/control.py @@ -17,6 +17,9 @@ class PipelineController: A pipeline is single-use: after it stops or fails, it cannot be restarted. A replacement deployment must create a new controller and adapter. + + Start and stop requests can arrive concurrently from HTTP handler and process + shutdown threads. The lock prevents state transition race conditions. """ def __init__(self, runtime: StreamAdapter[Any, Any]) -> None: