diff --git a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py b/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py index 94d4e67b..0fc81304 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py +++ b/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, diff --git a/sentry_streams/tests/adapters/arroyo/test_rust_arroyo.py b/sentry_streams/tests/adapters/arroyo/test_rust_arroyo.py index ef985b7a..75688bc5 100644 --- a/sentry_streams/tests/adapters/arroyo/test_rust_arroyo.py +++ b/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 == {} diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/constants.py b/sentry_streams_k8s/sentry_streams_k8s/operator/constants.py index b5199180..17148580 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/constants.py +++ b/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"}) diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/pod_resources.py b/sentry_streams_k8s/sentry_streams_k8s/operator/pod_resources.py index 9f0d28f4..741f0451 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/pod_resources.py +++ b/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. diff --git a/sentry_streams_k8s/tests/test_pods.py b/sentry_streams_k8s/tests/test_pods.py index 4efddd76..48ffade0 100644 --- a/sentry_streams_k8s/tests/test_pods.py +++ b/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 = []