From 5d66f0eadf091516ac8785014efb9dd87cb55bed Mon Sep 17 00:00:00 2001 From: bmcquilkin Date: Thu, 30 Jul 2026 15:51:14 -0700 Subject: [PATCH] feat(operator): static membership v2 --- .../sentry_streams/adapters/arroyo/adapter.py | 3 +- .../adapters/arroyo/rust_arroyo.py | 5 ++- .../sentry_streams/adapters/stream_adapter.py | 11 ++++-- sentry_streams/sentry_streams/control.py | 5 ++- .../sentry_streams/dummy/dummy_adapter.py | 8 +++- .../sentry_streams/rust_streams.pyi | 1 + .../sentry_streams/server/control_server.py | 39 ++++++++++++++++++- sentry_streams/src/consumer.rs | 17 +++++++- sentry_streams/src/kafka_config.rs | 7 ++++ sentry_streams/tests/adapters/fake_adapter.py | 5 ++- sentry_streams/tests/test_control_server.py | 27 +++++++++++-- 11 files changed, 112 insertions(+), 16 deletions(-) diff --git a/sentry_streams/sentry_streams/adapters/arroyo/adapter.py b/sentry_streams/sentry_streams/adapters/arroyo/adapter.py index a3f829bc..d7a53f59 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/adapter.py +++ b/sentry_streams/sentry_streams/adapters/arroyo/adapter.py @@ -36,6 +36,7 @@ from sentry_streams.adapters.stream_adapter import ( PipelineConfig, RuntimeState, + StartOptions, StreamAdapter, ) from sentry_streams.config_types import ( @@ -326,7 +327,7 @@ def create_processors(self) -> None: for source, consumer in self.__consumers.items() } - def _run(self) -> None: + def _run(self, options: StartOptions) -> None: """ Starts the pipeline """ diff --git a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py b/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py index 9c0ff705..ce816430 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py +++ b/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py @@ -30,6 +30,7 @@ from sentry_streams.adapters.stream_adapter import ( PipelineConfig, RuntimeState, + StartOptions, StreamAdapter, ) from sentry_streams.config_types import ( @@ -591,13 +592,15 @@ def routing_function(msg: Message[Any]) -> str: ) return build_branches(stream, step.routing_table.values()) - def _run(self) -> None: + def _run(self, options: StartOptions) -> None: """ Starts the pipeline """ # TODO: Support multiple consumers assert len(self.__consumers) == 1, "Multiple consumers not supported yet" consumer = next(iter(self.__consumers.values())) + if options.group_instance_id is not None: + consumer.set_group_instance_id(options.group_instance_id) self._set_status(RuntimeState.CONSUMING) consumer.run() diff --git a/sentry_streams/sentry_streams/adapters/stream_adapter.py b/sentry_streams/sentry_streams/adapters/stream_adapter.py index 29954bdc..4845e4cf 100644 --- a/sentry_streams/sentry_streams/adapters/stream_adapter.py +++ b/sentry_streams/sentry_streams/adapters/stream_adapter.py @@ -78,6 +78,11 @@ def can_transition_to(self, state: RuntimeState) -> bool: return self.rank < state.rank +@dataclass(frozen=True) +class StartOptions: + group_instance_id: str | None = None + + @dataclass(frozen=True) class RuntimeStatus: state: RuntimeState @@ -126,7 +131,7 @@ def begin_start(self) -> RuntimeStatus: self.__status = RuntimeStatus(RuntimeState.STARTING) return self.__status - def run(self) -> None: + def run(self, options: StartOptions | None = None) -> None: state = self.status.state if state is RuntimeState.STOPPING: self._set_status(RuntimeState.STOPPED) @@ -135,7 +140,7 @@ def run(self) -> None: raise RuntimeStateError(f"cannot run runtime while it is {state}") try: - self._run() + self._run(options if options is not None else StartOptions()) except Exception as exc: self._set_status(RuntimeState.ERRORED, exc) raise @@ -251,7 +256,7 @@ def broadcast( raise NotImplementedError @abstractmethod - def _run(self) -> None: + def _run(self, options: StartOptions) -> None: """ Starts the pipeline """ diff --git a/sentry_streams/sentry_streams/control.py b/sentry_streams/sentry_streams/control.py index 458fe9bf..ca16d54a 100644 --- a/sentry_streams/sentry_streams/control.py +++ b/sentry_streams/sentry_streams/control.py @@ -6,6 +6,7 @@ from sentry_streams.adapters.stream_adapter import ( RuntimeStatus, + StartOptions, StreamAdapter, ) @@ -33,14 +34,14 @@ def __init__(self, runtime: StreamAdapter[Any, Any]) -> None: def snapshot(self) -> RuntimeStatus: return self._runtime.status - def request_start(self) -> RuntimeStatus: + def request_start(self, options: StartOptions | None = None) -> RuntimeStatus: """ Ask the pipeline to start (non-blocking). """ with self._lock: status = self._runtime.begin_start() if self._run_future is None: - self._run_future = self._executor.submit(self._runtime.run) + self._run_future = self._executor.submit(self._runtime.run, options) self._run_future.add_done_callback(self._run_finished) return status diff --git a/sentry_streams/sentry_streams/dummy/dummy_adapter.py b/sentry_streams/sentry_streams/dummy/dummy_adapter.py index b7651e3e..89e34489 100644 --- a/sentry_streams/sentry_streams/dummy/dummy_adapter.py +++ b/sentry_streams/sentry_streams/dummy/dummy_adapter.py @@ -1,6 +1,10 @@ from typing import Any, Callable, Optional, Self, Sequence, Type, TypeVar, cast -from sentry_streams.adapters.stream_adapter import PipelineConfig, StreamAdapter +from sentry_streams.adapters.stream_adapter import ( + PipelineConfig, + StartOptions, + StreamAdapter, +) from sentry_streams.pipeline.function_template import ( InputType, OutputType, @@ -95,7 +99,7 @@ def router(self, step: Router[RoutingFuncReturnType, Any], stream: Any) -> Any: ret[branch.root.name] = branch return ret - def _run(self) -> None: + def _run(self, options: StartOptions) -> None: pass def _shutdown(self) -> None: diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index e4c359cf..27b2c567 100644 --- a/sentry_streams/sentry_streams/rust_streams.pyi +++ b/sentry_streams/sentry_streams/rust_streams.pyi @@ -159,6 +159,7 @@ class ArroyoConsumer: ) -> None: ... def add_threadpool(self, step_name: str, thread_count: int) -> None: ... def add_step(self, step: RuntimeOperator) -> None: ... + def set_group_instance_id(self, group_instance_id: str) -> None: ... def run(self) -> None: ... def shutdown(self) -> None: ... @property diff --git a/sentry_streams/sentry_streams/server/control_server.py b/sentry_streams/sentry_streams/server/control_server.py index 1f764918..6fe5637b 100644 --- a/sentry_streams/sentry_streams/server/control_server.py +++ b/sentry_streams/sentry_streams/server/control_server.py @@ -9,17 +9,48 @@ from sentry_streams.adapters.stream_adapter import ( RuntimeStateError, RuntimeStatus, + StartOptions, ) from sentry_streams.control import PipelineController logger = logging.getLogger(__name__) +class InvalidRequest(Exception): + """Raised when a request body cannot be read.""" + + class ControlHandler(BaseHTTPRequestHandler): @property def _controller(self) -> PipelineController: return cast(ControlServer, self.server).controller + def _read_start_options(self) -> StartOptions | None: + length = int(self.headers.get("Content-Length") or 0) + raw = self.rfile.read(length) if length > 0 else b"" + if not raw.strip(): + return None + + try: + payload = json.loads(raw) + except json.JSONDecodeError as exc: + raise InvalidRequest(f"body is not valid JSON: {exc}") from exc + + if not isinstance(payload, dict): + raise InvalidRequest("body must be a JSON object") + + unknown = set(payload) - {"group_instance_id"} + if unknown: + raise InvalidRequest(f"unknown fields: {', '.join(sorted(unknown))}") + + group_instance_id = payload.get("group_instance_id") + if group_instance_id is not None and ( + not isinstance(group_instance_id, str) or not group_instance_id + ): + raise InvalidRequest("group_instance_id must be a non-empty string") + + return StartOptions(group_instance_id=group_instance_id) + def log_message(self, format: str, *args: Any) -> None: logger.debug("control-server %s - %s", self.address_string(), format % args) @@ -49,7 +80,8 @@ def do_POST(self) -> None: try: path = urlparse(self.path).path if path == "/start": - self._respond(202, self._controller.request_start().as_dict()) + options = self._read_start_options() + self._respond(202, self._controller.request_start(options).as_dict()) elif path == "/stop": self._respond_to_stop(self._controller.request_stop()) else: @@ -62,7 +94,10 @@ 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, RuntimeStateError): + if isinstance(exc, InvalidRequest): + logger.info("control-server could not read %s: %s", self.path, exc) + self._respond(400, {"error": str(exc)}) + elif 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/src/consumer.rs b/sentry_streams/src/consumer.rs index 5ef2aa52..82cf299e 100644 --- a/sentry_streams/src/consumer.rs +++ b/sentry_streams/src/consumer.rs @@ -91,6 +91,9 @@ pub struct ArroyoConsumer { /// processor is not lost. `run()` checks it before entering the run loop. shutdown_requested: AtomicBool, + /// Static membership id provided by controller. + group_instance_id: Mutex>, + /// The ProcessorHandle allows the main thread to stop the StreamingProcessor /// from a different thread. handle: Mutex>, @@ -140,6 +143,7 @@ impl ArroyoConsumer { schema, steps: Vec::new(), shutdown_requested: AtomicBool::new(false), + group_instance_id: Mutex::new(None), handle: Mutex::new(None), concurrency_config: Arc::new(ConcurrencyConfig::new(1)), step_concurrency_configs: HashMap::new(), @@ -167,6 +171,10 @@ impl ArroyoConsumer { self.steps.push(step); } + fn set_group_instance_id(&self, group_instance_id: String) { + *self.group_instance_id.lock().unwrap() = Some(group_instance_id); + } + /// Runs the consumer. /// This method is blocking and will run until the consumer /// is stopped via shutdown(). Signals are handled by the Python @@ -197,7 +205,14 @@ impl ArroyoConsumer { self.topic.clone(), self.consumer_config.group_id().to_string(), ); - let config = self.consumer_config.clone().into(); + + let mut consumer_config = self.consumer_config.clone(); + + if let Some(id) = self.group_instance_id.lock().unwrap().as_deref() { + consumer_config.set_group_instance_id(id); + } + + let config = consumer_config.into(); // Build DLQ policy if configured let dlq_policy = build_dlq_policy(&self.dlq_config, self.concurrency_config.handle()); diff --git a/sentry_streams/src/kafka_config.rs b/sentry_streams/src/kafka_config.rs index 04d2dc02..7dbf3b49 100644 --- a/sentry_streams/src/kafka_config.rs +++ b/sentry_streams/src/kafka_config.rs @@ -98,6 +98,13 @@ impl PyKafkaConsumerConfig { pub fn group_id(&self) -> &str { &self.group_id } + + pub fn set_group_instance_id(&mut self, group_instance_id: &str) { + self.override_params.get_or_insert_default().insert( + "group.instance.id".to_string(), + group_instance_id.to_string(), + ); + } } impl From for KafkaConfig { diff --git a/sentry_streams/tests/adapters/fake_adapter.py b/sentry_streams/tests/adapters/fake_adapter.py index 7f6328d8..3f918418 100644 --- a/sentry_streams/tests/adapters/fake_adapter.py +++ b/sentry_streams/tests/adapters/fake_adapter.py @@ -6,6 +6,7 @@ from sentry_streams.adapters.stream_adapter import ( PipelineConfig, RuntimeState, + StartOptions, StreamAdapter, ) from sentry_streams.pipeline.function_template import InputType, OutputType @@ -35,9 +36,11 @@ def __init__(self, fail: bool = False, block_before_consuming: bool = False) -> self.run_finished = threading.Event() self.run_calls = 0 self.shutdown_calls = 0 + self.start_options: StartOptions | None = None - def _run(self) -> None: + def _run(self, options: StartOptions) -> None: self.run_calls += 1 + self.start_options = options self.run_started.set() if self._fail: raise RuntimeError("runtime failed") diff --git a/sentry_streams/tests/test_control_server.py b/sentry_streams/tests/test_control_server.py index 3f787fde..a5f66798 100644 --- a/sentry_streams/tests/test_control_server.py +++ b/sentry_streams/tests/test_control_server.py @@ -7,7 +7,7 @@ import urllib.request from typing import Any, Callable -from sentry_streams.adapters.stream_adapter import RuntimeState +from sentry_streams.adapters.stream_adapter import RuntimeState, StartOptions from sentry_streams.control import PipelineController from sentry_streams.server.control_server import make_server from tests.adapters.fake_adapter import FakeAdapter @@ -22,8 +22,10 @@ def _wait_for(predicate: Callable[[], bool], timeout: float = 3.0) -> bool: return False -def _request(port: int, path: str, method: str) -> tuple[int, dict[str, Any]]: - data = b"" if method == "POST" else None +def _request( + port: int, path: str, method: str, body: bytes | None = None +) -> tuple[int, dict[str, Any]]: + data = body if body is not None else (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: @@ -90,3 +92,22 @@ def test_readyz_reports_runtime_failure() -> None: server.server_close() thread.join(timeout=3.0) _stop(controller) + + +def test_start_passes_a_group_instance_id_to_the_runtime() -> 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: + body = json.dumps({"group_instance_id": "consumer-2"}).encode() + assert _request(port, "/start", "POST", body)[0] == 202 + assert _wait_for(lambda: runtime.start_options is not None) + assert runtime.start_options == StartOptions(group_instance_id="consumer-2") + finally: + server.shutdown() + server.server_close() + thread.join(timeout=3.0) + _stop(controller)