feat(operator): static membership

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
         """

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()
 

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
         """

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
 

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:

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

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:

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<Option<String>>,
+
     /// The ProcessorHandle allows the main thread to stop the StreamingProcessor
     /// from a different thread.
     handle: Mutex<Option<ProcessorHandle>>,
@@ -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());

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<PyKafkaConsumerConfig> for KafkaConfig {

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")

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)