feat(operator): static membership

sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py

@@ -2,6 +2,7 @@
 
 import functools
 import logging
+import os
 import time
 from dataclasses import replace
 from typing import (
@@ -144,20 +145,29 @@ def build_initial_offset(offset_reset: str) -> InitialOffset:
 
 def build_kafka_consumer_config(
     source: str,
-    source_config: StepConfig,
+    source_config: KafkaConsumerConfig,
     consumer_group_override: str | None = None,
 ) -> PyKafkaConsumerConfig:
     """
     Build the Kafka consumer configuration for the source.
     """
-    consumer_config = cast(KafkaConsumerConfig, source_config)
+    consumer_config = source_config
     bootstrap_servers = consumer_config["bootstrap_servers"]
     group_id = (
         consumer_group_override or consumer_config.get("consumer_group") or f"pipeline-{source}"
     )
     auto_offset_reset = build_initial_offset(consumer_config.get("auto_offset_reset", "latest"))
     strict_offset_reset = bool(consumer_config.get("strict_offset_reset", False))
-    override_params = consumer_config.get("override_params", {})
+    override_params = dict(consumer_config.get("override_params", {}))
+
+    group_instance_id = os.environ.get("STREAMS_KAFKA_GROUP_INSTANCE_ID")
+    if group_instance_id and "group.instance.id" not in override_params:
+        override_params["group.instance.id"] = group_instance_id
+        logger.info(
+            "Enabling Kafka static membership for source %s with group.instance.id=%s",
+            source,
+            group_instance_id,
+        )
 
     return PyKafkaConsumerConfig(
         bootstrap_servers=bootstrap_servers,
@@ -312,7 +322,7 @@ def source(self, step: Source[Any]) -> Route:
         self.__consumers[source_name] = ArroyoConsumer(
             source=source_name,
             kafka_config=build_kafka_consumer_config(
-                source_name, source_config, step.consumer_group
+                source_name, cast(KafkaConsumerConfig, source_config), step.consumer_group
             ),
             topic=step.stream_name,
             schema=schema_name,

sentry_streams/tests/adapters/arroyo/test_rust_arroyo.py

@@ -8,7 +8,7 @@
     build_kafka_producer_config,
 )
 from sentry_streams.adapters.stream_adapter import RuntimeTranslator
-from sentry_streams.config_types import StepConfig
+from sentry_streams.config_types import KafkaConsumerConfig, StepConfig
 from sentry_streams.pipeline.pipeline import Pipeline
 from sentry_streams.runner import iterate_edges
 from sentry_streams.rust_streams import InitialOffset
@@ -46,6 +46,7 @@ def test_rust_arroyo_adapter(
     [
         pytest.param(
             {
+                "starts_segment": None,
                 "bootstrap_servers": ["localhost:9092"],
                 "auto_offset_reset": "earliest",
             },
@@ -59,6 +60,7 @@ def test_rust_arroyo_adapter(
         ),
         pytest.param(
             {
+                "starts_segment": None,
                 "bootstrap_servers": ["localhost:9092"],
                 "auto_offset_reset": "latest",
                 "consumer_group": "my-group",
@@ -75,6 +77,7 @@ def test_rust_arroyo_adapter(
         ),
         pytest.param(
             {
+                "starts_segment": None,
                 "bootstrap_servers": ["broker1:9092", "broker2:9092"],
                 "auto_offset_reset": "earliest",
                 "consumer_group": "config-group",
@@ -94,7 +97,7 @@ def test_rust_arroyo_adapter(
     ],
 )
 def test_build_kafka_consumer_config(
-    source_config: StepConfig,
+    source_config: KafkaConsumerConfig,
     consumer_group_override: str | None,
     expected_bootstrap_servers: Sequence[str],
     expected_group_id: str,
@@ -157,3 +160,54 @@ def test_build_kafka_producer_config(
     assert result is not None
     assert list(result.bootstrap_servers) == expected_bootstrap_servers
     assert result.override_params == expected_override_params
+
+
+def test_build_kafka_consumer_config_injects_static_membership(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("STREAMS_KAFKA_GROUP_INSTANCE_ID", "consumer-2")
+    result = build_kafka_consumer_config(
+        source="test_source",
+        source_config={
+            "starts_segment": None,
+            "bootstrap_servers": ["localhost:9092"],
+            "auto_offset_reset": "earliest",
+        },
+    )
+    assert result.override_params is not None
+    assert result.override_params["group.instance.id"] == "consumer-2"
+
+
+def test_build_kafka_consumer_config_respects_explicit_instance_id(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("STREAMS_KAFKA_GROUP_INSTANCE_ID", "consumer-2")
+    result = build_kafka_consumer_config(
+        source="test_source",
+        source_config={
+            "starts_segment": None,
+            "bootstrap_servers": ["localhost:9092"],
+            "auto_offset_reset": "earliest",
+            "override_params": {
+                "group.instance.id": "explicit",
+                "session.timeout.ms": "10000",
+            },
+        },
+    )
+    assert result.override_params is not None
+    assert result.override_params["group.instance.id"] == "explicit"
+
+
+def test_build_kafka_consumer_config_no_static_membership_without_env(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("STREAMS_KAFKA_GROUP_INSTANCE_ID", raising=False)
+    result = build_kafka_consumer_config(
+        source="test_source",
+        source_config={
+            "starts_segment": None,
+            "bootstrap_servers": ["localhost:9092"],
+            "auto_offset_reset": "earliest",
+        },
+    )
+    assert result.override_params == {}

sentry_streams_k8s/sentry_streams_k8s/operator/constants.py

@@ -52,6 +52,11 @@
 
 SPEC_HASH_ANNOTATION = "streams.sentry.io/spec-hash"
 
+# Environment variable passed to each consumer Pod giving it a stable identity
+# (group.instance.id). Replacement Pods rejoin as the same static member.
+
+GROUP_INSTANCE_ID_ENV = "STREAMS_KAFKA_GROUP_INSTANCE_ID"
+
 # Waiting states that may recover when the operator creates a new Pod:
 
 UNHEALTHY_WAITING_REASONS = frozenset({"ErrImagePull", "ImagePullBackOff"})

sentry_streams_k8s/sentry_streams_k8s/operator/pod_resources.py

@@ -12,6 +12,7 @@
 from sentry_streams_k8s.operator.constants import (
     FIELD_MANAGER,
     GENERATION_LABEL,
+    GROUP_INSTANCE_ID_ENV,
     MANAGED_BY_LABEL,
     ORDINAL_LABEL,
     OWNER_NAME_ANNOTATION,
@@ -79,7 +80,11 @@ def pod_name(pod: Mapping[str, Any]) -> str:
 
 
 def consumer_pod_name(base_name: str, ordinal: int, generation: int) -> str:
-    return f"{base_name}-{ordinal}-{generation}"
+    return f"{group_instance_id(base_name, ordinal)}-{generation}"
+
+
+def group_instance_id(base_name: str, ordinal: int) -> str:
+    return f"{base_name}-{ordinal}"
 
 
 def pod_ordinal(pod: Mapping[str, Any]) -> int | None:
@@ -157,6 +162,15 @@ def build_pipeline_pod(
 
     pod_spec["restartPolicy"] = "Never"
 
+    # Give the replica a stable Kafka identity so replacement
+    # Pods rejoin without triggering a rebalance:
+
+    for container in pod_spec.get("containers", []) or []:
+        env = container.setdefault("env", [])
+        if any(item.get("name") == GROUP_INSTANCE_ID_ENV for item in env):
+            continue
+        env.append({"name": GROUP_INSTANCE_ID_ENV, "value": group_instance_id(base_name, ordinal)})
+
     # Pods can live in a different namespace from their StreamingPipeline, so they
     # cannot use owner references. The operator finds them by owner label instead.
 

sentry_streams_k8s/tests/test_pods.py

@@ -10,6 +10,7 @@
 from sentry_streams_k8s.operator.constants import (
     CANARY_WORKLOAD_SET,
     GENERATION_LABEL,
+    GROUP_INSTANCE_ID_ENV,
     MAX_BASE_NAME_LENGTH,
     MAX_GENERATION,
     MAX_REPLICAS,
@@ -21,6 +22,7 @@
 from sentry_streams_k8s.operator.pod_resources import (
     build_pipeline_pod,
     delete_owned_pods,
+    group_instance_id,
     list_owned_pods,
     pod_workload_set,
 )
@@ -31,6 +33,12 @@
     reconcile_pipeline_pods,
 )
 
+
+def _container_env(pod: dict[str, Any], container_index: int = 0) -> dict[str, str]:
+    env = pod["spec"]["containers"][container_index].get("env", [])
+    return {item["name"]: item["value"] for item in env}
+
+
 NAMESPACE = "workloads"
 OWNER_UID = "owner-uid"
 
@@ -145,6 +153,64 @@ def test_build_pipeline_pod_stamps_identity_without_mutating_template() -> None:
     assert "restartPolicy" not in spec
 
 
+def test_group_instance_id_is_stable_across_generations() -> None:
+    assert group_instance_id("consumer", 2) == "consumer-2"
+    assert group_instance_id("consumer-canary", 2) == "consumer-canary-2"
+
+
+def test_build_pipeline_pod_injects_generation_stable_instance_id() -> None:
+    metadata, spec = _template()
+
+    def _build(generation: int) -> dict[str, Any]:
+        return build_pipeline_pod(
+            base_name="consumer",
+            template_metadata=metadata,
+            template_spec=spec,
+            ordinal=2,
+            generation=generation,
+            owner_uid=OWNER_UID,
+            owner_name="pipeline",
+            owner_namespace="source",
+            workload_set=PRIMARY_WORKLOAD_SET,
+        )
+
+    gen4 = _build(4)
+    gen7 = _build(7)
+
+    assert _container_env(gen4)[GROUP_INSTANCE_ID_ENV] == "consumer-2"
+    assert _container_env(gen7)[GROUP_INSTANCE_ID_ENV] == "consumer-2"
+
+    assert (
+        gen4["metadata"]["annotations"][SPEC_HASH_ANNOTATION]
+        == gen7["metadata"]["annotations"][SPEC_HASH_ANNOTATION]
+    )
+
+
+def test_build_pipeline_pod_respects_explicit_instance_id_env() -> None:
+    metadata, _ = _template()
+    spec = {
+        "containers": [
+            {
+                "name": "consumer",
+                "image": "example/consumer:v1",
+                "env": [{"name": GROUP_INSTANCE_ID_ENV, "value": "explicit"}],
+            }
+        ]
+    }
+    pod = build_pipeline_pod(
+        base_name="consumer",
+        template_metadata=metadata,
+        template_spec=spec,
+        ordinal=2,
+        generation=0,
+        owner_uid=OWNER_UID,
+        owner_name="pipeline",
+        owner_namespace="source",
+        workload_set=PRIMARY_WORKLOAD_SET,
+    )
+    assert _container_env(pod)[GROUP_INSTANCE_ID_ENV] == "explicit"
+
+
 def test_list_owned_pods_selects_owner_and_workload_set() -> None:
     resource = MagicMock()
     resource.get.return_value.items = []