feat(operator): pod handoff

.pre-commit-config.yaml

@@ -8,6 +8,8 @@ repos:
       - id: check-symlinks
       - id: check-xml
       - id: check-yaml
+        # Helm templates have to be rendered before they are valid Yaml.
+        exclude: "^sentry_streams_k8s/chart/[^/]+/templates/"
       - id: detect-private-key
       - id: end-of-file-fixer
         # Exclude CHANGELOG.md from this hook as Craft generates updates

sentry_streams_k8s/chart/streaming-operator/templates/_helpers.tpl

@@ -28,3 +28,11 @@ helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version }}
 {{- define "streaming-operator.workloadNamespace" -}}
 {{- required "workloadNamespace must be set" .Values.workloadNamespace -}}
 {{- end -}}
+
+{{- define "streaming-operator.controlHost" -}}
+{{- required "control.host must be set" .Values.control.host -}}
+{{- end -}}
+
+{{- define "streaming-operator.controlPort" -}}
+{{- required "control.port must be set" .Values.control.port -}}
+{{- end -}}

sentry_streams_k8s/chart/streaming-operator/templates/deployment.yaml

@@ -38,6 +38,10 @@ spec:
           env:
             - name: WORKLOAD_NAMESPACE
               value: {{ include "streaming-operator.workloadNamespace" . | quote }}
+            - name: CONTROL_HOST
+              value: {{ include "streaming-operator.controlHost" . | quote }}
+            - name: CONTROL_PORT
+              value: {{ include "streaming-operator.controlPort" . | quote }}
             {{- with .Values.env }}
             {{- toYaml . | nindent 12 }}
             {{- end }}

sentry_streams_k8s/chart/streaming-operator/values.yaml

@@ -4,6 +4,10 @@
 
 workloadNamespace: ""
 
+control:
+  host: ""
+  port: ""
+
 image:
   repository: us-central1-docker.pkg.dev/sentryio/streaming-operator/image
   tag: ""

sentry_streams_k8s/sentry_streams_k8s/consumer_builder.py

@@ -9,7 +9,10 @@
 
 import yaml
 
-from sentry_streams_k8s.constants import CANARY_WORKLOAD_SET, PRIMARY_WORKLOAD_SET
+from sentry_streams_k8s.constants import (
+    CANARY_WORKLOAD_SET,
+    PRIMARY_WORKLOAD_SET,
+)
 from sentry_streams_k8s.k8s_types import V1ConfigMapDict, V1DeploymentDict
 from sentry_streams_k8s.merge import ScalarOverwriteError, deepmerge
 from sentry_streams_k8s.validation import validate_pipeline_config
@@ -138,6 +141,9 @@ def build_container(
     enable_liveness_probe: bool = True,
     multiprocess_enabled: bool | None = None,
     container_name: str = "pipeline-consumer",
+    *,
+    control_host: str | None,
+    control_port: int | None,
 ) -> dict[str, Any]:
     """
     Build a complete container specification for the pipeline step.
@@ -187,23 +193,29 @@ def build_container(
             }
         )
 
+    args = [
+        "-n",
+        pipeline_name,
+        "--log-level",
+        log_level,
+        "--adapter",
+        "rust_arroyo",
+        "--segment-id",
+        str(segment_id),
+        "--config",
+        "/etc/pipeline-config/pipeline_config.yaml",
+    ]
+
+    if control_host is not None and control_port is not None:
+        args += ["--control-host", control_host, "--control-port", str(control_port)]
+
+    args.append(pipeline_module)
+
     pipeline_additions: dict[str, Any] = {
         "name": container_name,
         "image": image_name,
         "command": ["python", "-m", "sentry_streams.runner"],
-        "args": [
-            "-n",
-            pipeline_name,
-            "--log-level",
-            log_level,
-            "--adapter",
-            "rust_arroyo",
-            "--segment-id",
-            str(segment_id),
-            "--config",
-            "/etc/pipeline-config/pipeline_config.yaml",
-            pipeline_module,
-        ],
+        "args": args,
         "resources": {
             "requests": {
                 "cpu": f"{cpu_total}m",
@@ -225,6 +237,13 @@ def build_container(
             "periodSeconds": 10,
         }
 
+    if control_port is not None:
+        pipeline_additions["readinessProbe"] = {
+            "httpGet": {"path": "/readyz", "port": control_port},
+            "periodSeconds": 5,
+            "failureThreshold": 3,
+        }
+
     return deepmerge(container, pipeline_additions)
 
 
@@ -356,7 +375,11 @@ def validate(self, spec: ConsumerSpec, pipeline_config: Mapping[str, Any]) -> No
             )
 
     def _workload_deployments(
-        self, spec: ConsumerSpec, config: dict[str, Any]
+        self,
+        spec: ConsumerSpec,
+        config: dict[str, Any],
+        control_host: str | None,
+        control_port: int | None,
     ) -> dict[str, V1DeploymentDict]:
         """
         Generate the Kubernetes Deployments for the pipeline.
@@ -385,6 +408,8 @@ def _workload_deployments(
             spec.enable_liveness_probe,
             multiprocess_enabled,
             spec.container_name,
+            control_host=control_host,
+            control_port=control_port,
         )
 
         base_deployment = load_base_template("deployment")
@@ -488,7 +513,7 @@ def build_deployments(
         """Render the pipeline as Deployments and a ConfigMap."""
 
         config = dict(pipeline_config)
-        deployments = self._workload_deployments(spec, config)
+        deployments = self._workload_deployments(spec, config, control_host=None, control_port=None)
         primary = deployments[PRIMARY_WORKLOAD_SET]
 
         result: RenderedDeployments = {
@@ -501,11 +526,19 @@ def build_deployments(
 
         return result
 
-    def build_pods(self, spec: ConsumerSpec, pipeline_config: Mapping[str, Any]) -> RenderedPods:
+    def build_pods(
+        self,
+        spec: ConsumerSpec,
+        pipeline_config: Mapping[str, Any],
+        control_host: str,
+        control_port: int,
+    ) -> RenderedPods:
         """Render the pipeline as Pods and a ConfigMap."""
 
         config = dict(pipeline_config)
-        deployments = self._workload_deployments(spec, config)
+        deployments = self._workload_deployments(
+            spec, config, control_host=control_host, control_port=control_port
+        )
 
         return {
             "sets": {

sentry_streams_k8s/sentry_streams_k8s/operator/constants.py

@@ -69,6 +69,20 @@
 POD_WAITING_GRACE_SECONDS = 300
 POD_TERMINATING_GRACE_SECONDS = 600
 
+# Environment variables for the control server address:
+
+CONTROL_HOST_ENV = "CONTROL_HOST"
+CONTROL_PORT_ENV = "CONTROL_PORT"
+
+# Handoff timeout and poll interval:
+
+HANDOFF_DRAIN_TIMEOUT_SECONDS = 30.0
+HANDOFF_POLL_INTERVAL_SECONDS = 0.5
+
+# Request timeout to a Pod's control server:
+
+CONTROL_REQUEST_TIMEOUT_SECONDS = 5.0
+
 # How often the daemon re-runs a full reconcile as a safety net behind the watch:
 
 HEALTH_SCAN_INTERVAL_SECONDS = 60

sentry_streams_k8s/sentry_streams_k8s/operator/control_client.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import json
+import logging
+import urllib.error
+import urllib.request
+from enum import StrEnum
+from typing import Any, cast
+
+from sentry_streams_k8s.operator.constants import CONTROL_REQUEST_TIMEOUT_SECONDS
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeState(StrEnum):
+    # Mirrors sentry_streams.adapters.stream_adapter.RuntimeState.
+    # TODO: Remove  once sentry_streams is updated to expose RuntimeState.
+
+    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)
+
+
+class ControlError(Exception):
+    """Raised when a control request could not be completed."""
+
+
+class ControlClient:
+    def __init__(
+        self,
+        port: int,
+        timeout: float = CONTROL_REQUEST_TIMEOUT_SECONDS,
+    ) -> None:
+        self._port = port
+        self._timeout = timeout
+
+    def _request(self, ip: str, path: str, method: str, body: dict[str, Any] | None = None) -> Any:
+        data = json.dumps(body).encode() if body is not None else b""
+        request = urllib.request.Request(
+            f"http://{ip}:{self._port}{path}",
+            method=method,
+            data=data,
+            headers={"Content-Type": "application/json"} if body is not None else {},
+        )
+        with urllib.request.urlopen(request, timeout=self._timeout) as response:
+            raw = response.read()
+        return json.loads(raw) if raw else {}
+
+    def status(self, ip: str) -> RuntimeState | None:
+        try:
+            payload = self._request(ip, "/status", "GET")
+        except (urllib.error.URLError, OSError, ValueError) as error:
+            logger.debug("control server at %s is unreachable: %s", ip, error)
+            return None
+
+        state = cast(dict[str, Any], payload).get("state")
+        try:
+            return RuntimeState(state) if isinstance(state, str) else None
+        except ValueError:
+            logger.warning("control server at %s reported unknown state %r", ip, state)
+            return None
+
+    def readyz(self, ip: str) -> bool:
+        try:
+            self._request(ip, "/readyz", "GET")
+            return True
+        except (urllib.error.URLError, OSError, ValueError):
+            return False
+
+    def start(self, ip: str, group_instance_id: str) -> None:
+        try:
+            self._request(ip, "/start", "POST", {"group_instance_id": group_instance_id})
+        except urllib.error.HTTPError as error:
+            if error.code == 409:
+                logger.info("consumer at %s was not idle when started: %s", ip, error)
+                return
+            raise ControlError(f"failed to start consumer at {ip}: {error}") from error
+        except (urllib.error.URLError, OSError, ValueError) as error:
+            raise ControlError(f"failed to start consumer at {ip}: {error}") from error
+
+    def stop(self, ip: str) -> bool:
+        try:
+            self._request(ip, "/stop", "POST")
+            return True
+        except (urllib.error.HTTPError, urllib.error.URLError, OSError, ValueError) as error:
+            logger.info("could not stop consumer at %s, treating it as stopped: %s", ip, error)
+            return False

sentry_streams_k8s/sentry_streams_k8s/operator/operator.py

@@ -16,6 +16,8 @@
 
 from sentry_streams_k8s.k8s_types import V1ConditionDict
 from sentry_streams_k8s.operator.constants import (
+    CONTROL_HOST_ENV,
+    CONTROL_PORT_ENV,
     FIELD_MANAGER,
     GROUP,
     HEALTH_SCAN_INTERVAL_SECONDS,
@@ -95,6 +97,23 @@ def _workload_namespace() -> str:
     return namespace
 
 
+def _control_host() -> str:
+    host = os.environ.get(CONTROL_HOST_ENV, "").strip()
+    if not host:
+        raise RuntimeError(f"{CONTROL_HOST_ENV} must be set.")
+    return host
+
+
+def _control_port() -> int:
+    port = os.environ.get(CONTROL_PORT_ENV, "").strip()
+    if not port:
+        raise RuntimeError(f"{CONTROL_PORT_ENV} must be set.")
+    try:
+        return int(port)
+    except ValueError:
+        raise RuntimeError(f"{CONTROL_PORT_ENV} must be an integer, got {port!r}.")
+
+
 def _published_conditions(status: Mapping[str, object]) -> list[V1ConditionDict] | None:
     conditions = status.get("conditions")
     if not isinstance(conditions, list):
@@ -196,6 +215,8 @@ async def _reconcile_once(
                         namespace=namespace,
                         uid=uid,
                         workload_namespace=_workload_namespace(),
+                        control_host=_control_host(),
+                        control_port=_control_port(),
                         logger=logger,
                         status=status_patch,
                         previous_conditions=_published_conditions(published),
@@ -270,7 +291,7 @@ async def handle_pipeline_pod_event(
 
     if type == "MODIFIED" and meta.deletion_timestamp is None:
         health = pod_health(_deserialize_pod(body), datetime.now(timezone.utc))
-        if not health.delete:
+        if not (health.delete or health.ready):
             return
 
     owner_uid = labels.get(OWNER_UID_LABEL)

sentry_streams_k8s/sentry_streams_k8s/operator/pod_resources.py

@@ -63,11 +63,20 @@ def delete_pod(core: client.CoreV1Api, name: str, namespace: str, force: bool =
         core.delete_namespaced_pod(name=name, namespace=namespace)
 
 
+def pod_ip(pod: V1Pod) -> str | None:
+    status = pod.status
+    return status.pod_ip if status and status.pod_ip else None
+
+
 def pod_name(pod: V1Pod) -> str:
     metadata = pod.metadata
     return metadata.name if metadata and metadata.name else ""
 
 
+def group_instance_id(base_name: str, ordinal: int) -> str:
+    return f"{base_name}-{ordinal}"
+
+
 def consumer_pod_name(base_name: str, ordinal: int, generation: int) -> str:
     return f"{base_name}-{ordinal}-{generation}"