diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0d425e61..d7cc2cac 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/sentry_streams_k8s/chart/streaming-operator/templates/_helpers.tpl b/sentry_streams_k8s/chart/streaming-operator/templates/_helpers.tpl index 22e4be70..e5357392 100644 --- a/sentry_streams_k8s/chart/streaming-operator/templates/_helpers.tpl +++ b/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 -}} diff --git a/sentry_streams_k8s/chart/streaming-operator/templates/deployment.yaml b/sentry_streams_k8s/chart/streaming-operator/templates/deployment.yaml index a419c550..6dc4f560 100644 --- a/sentry_streams_k8s/chart/streaming-operator/templates/deployment.yaml +++ b/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 }} diff --git a/sentry_streams_k8s/chart/streaming-operator/values.yaml b/sentry_streams_k8s/chart/streaming-operator/values.yaml index f422782e..17e089eb 100644 --- a/sentry_streams_k8s/chart/streaming-operator/values.yaml +++ b/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: "" diff --git a/sentry_streams_k8s/sentry_streams_k8s/consumer_builder.py b/sentry_streams_k8s/sentry_streams_k8s/consumer_builder.py index 81d916c0..22c25845 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/consumer_builder.py +++ b/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": { diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/constants.py b/sentry_streams_k8s/sentry_streams_k8s/operator/constants.py index 778c8a70..5dc8a9de 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/constants.py +++ b/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 diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/control_client.py b/sentry_streams_k8s/sentry_streams_k8s/operator/control_client.py new file mode 100644 index 00000000..efac2f28 --- /dev/null +++ b/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 diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/operator.py b/sentry_streams_k8s/sentry_streams_k8s/operator/operator.py index 851ac554..84298a61 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/operator.py +++ b/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) 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 4b3873fe..0fad4b39 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/pod_resources.py +++ b/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}" diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/reconcile.py b/sentry_streams_k8s/sentry_streams_k8s/operator/reconcile.py index db3c08e7..667f724f 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/reconcile.py +++ b/sentry_streams_k8s/sentry_streams_k8s/operator/reconcile.py @@ -1,6 +1,8 @@ from __future__ import annotations +import time from collections.abc import Mapping +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, TypedDict, cast @@ -18,6 +20,8 @@ from sentry_streams_k8s.operator.constants import ( APPLY_PATCH_CONTENT_TYPE, FIELD_MANAGER, + HANDOFF_DRAIN_TIMEOUT_SECONDS, + HANDOFF_POLL_INTERVAL_SECONDS, MAX_BASE_NAME_LENGTH, MAX_GENERATION, MAX_REPLICAS, @@ -26,6 +30,11 @@ OWNER_UID_LABEL, Logger, ) +from sentry_streams_k8s.operator.control_client import ( + ControlClient, + ControlError, + RuntimeState, +) from sentry_streams_k8s.operator.pod_health import ( PodHealth, is_deleting, @@ -37,8 +46,10 @@ build_pipeline_pod, consumer_pod_name, delete_pod, + group_instance_id, list_owned_pods, pod_generation, + pod_ip, pod_keep_key, pod_name, pod_ordinal, @@ -233,6 +244,152 @@ def _allocate_generation(generations: dict[int, int], ordinal: int, pods: list[V return generation +@dataclass(frozen=True) +class _HandoffTarget: + """A replica's incoming Pod and the outdated Pods it has to take over from.""" + + ordinal: int + keep: V1Pod + outdated: list[V1Pod] + + +def _is_drained(pod: V1Pod, control: ControlClient) -> bool: + """True once a Pod has released its partitions.""" + + ip = pod_ip(pod) + if ip is None: + return True + + state = control.status(ip) + return state is None or state.is_terminal + + +def _request_stops( + *, + core: client.CoreV1Api, + workload_namespace: str, + targets: list[_HandoffTarget], + health_by_name: dict[str, PodHealth], + control: ControlClient, + logger: Logger, +) -> tuple[dict[int, list[V1Pod]], set[int]]: + """Ask every outdated Pod to stop, for all replicas at once.""" + + awaiting: dict[int, list[V1Pod]] = {} + blocked: set[int] = set() + + for target in targets: + for pod in target.outdated: + if is_deleting(pod): + awaiting.setdefault(target.ordinal, []).append(pod) + continue + + ip = pod_ip(pod) + if ip is not None and control.stop(ip): + awaiting.setdefault(target.ordinal, []).append(pod) + continue + + # Fall back to deleting the Pod, but wait until the next pass to start its replacement. + _delete_current_pod( + core, pod, workload_namespace, logger, health_by_name[pod_name(pod)], "Outdated" + ) + blocked.add(target.ordinal) + + return awaiting, blocked + + +def _wait_for_drain( + awaiting: dict[int, list[V1Pod]], + control: ControlClient, + logger: Logger, +) -> set[int]: + """ + Wait for the outdated Pods to finish committing + Returns a set of ordinals that are done. + """ + + pending = {ordinal: list(pods) for ordinal, pods in awaiting.items()} + deadline = time.monotonic() + HANDOFF_DRAIN_TIMEOUT_SECONDS + + while pending: + for ordinal, pods in list(pending.items()): + remaining = [pod for pod in pods if not _is_drained(pod, control)] + if remaining: + pending[ordinal] = remaining + else: + del pending[ordinal] + + if not pending: + break + + if time.monotonic() >= deadline: + logger.info( + "replicas %s are still draining; continuing the handoff on the next reconcile", + sorted(pending), + ) + break + + time.sleep(HANDOFF_POLL_INTERVAL_SECONDS) + + return set(awaiting) - set(pending) + + +def _handoff_pods( + *, + core: client.CoreV1Api, + workload_namespace: str, + base_name: str, + targets: list[_HandoffTarget], + health_by_name: dict[str, PodHealth], + control: ControlClient, + logger: Logger, +) -> None: + if not targets: + return + + awaiting, blocked = _request_stops( + core=core, + workload_namespace=workload_namespace, + targets=targets, + health_by_name=health_by_name, + control=control, + logger=logger, + ) + drained = _wait_for_drain(awaiting, control, logger) + + for target in targets: + ordinal = target.ordinal + if ordinal in blocked or (ordinal in awaiting and ordinal not in drained): + continue + + ip = pod_ip(target.keep) + if ip is None: + continue + + state = control.status(ip) + if state is None: + logger.info("replica %d is not reachable yet; retrying later", ordinal) + continue + + if state is RuntimeState.IDLE: + try: + control.start(ip, group_instance_id(base_name, ordinal)) + except ControlError as error: + logger.warning("could not start replica %d: %s", ordinal, error) + continue + logger.info( + "started pipeline Pod %s/%s as %s", + workload_namespace, + pod_name(target.keep), + group_instance_id(base_name, ordinal), + ) + + for pod in awaiting.get(ordinal, []): + _delete_current_pod( + core, pod, workload_namespace, logger, health_by_name[pod_name(pod)], "Outdated" + ) + + def reconcile_pipeline_pods( *, core: client.CoreV1Api, @@ -247,6 +404,7 @@ def reconcile_pipeline_pods( generations: dict[int, int], logger: Logger, workload_set: str, + control: ControlClient, ) -> PodSetResult: desired_ordinals = set(range(max(replicas, 0))) current = list_owned_pods( @@ -261,6 +419,7 @@ def reconcile_pipeline_pods( health_by_name: dict[str, PodHealth] = {} reported_statuses: list[ReportedPodStatus] = [] active_pod_names: list[str] = [] + handoff_targets: list[_HandoffTarget] = [] def _build(ordinal: int, generation: int) -> V1PodDict: return build_pipeline_pod( @@ -319,19 +478,35 @@ def _build(ordinal: int, generation: int) -> V1PodDict: keep_name, ) + # Keep outdated Pods running until their replacement is ready for the handoff. + + outdated: list[V1Pod] = [] for pod in pods: if pod is keep: continue health = health_by_name[pod_name(pod)] - if pod_spec_changed(pod, desired_template): - _delete_current_pod(core, pod, workload_namespace, logger, health, "Outdated") - elif health.delete: + if health.delete: _delete_current_pod( core, pod, workload_namespace, logger, health, health.reason or "Unhealthy" ) + elif pod_spec_changed(pod, desired_template): + outdated.append(pod) elif pod in candidates: _delete_current_pod(core, pod, workload_namespace, logger, health, "Duplicate") + if keep is not None and health_by_name[pod_name(keep)].ready: + handoff_targets.append(_HandoffTarget(ordinal, keep, outdated)) + + _handoff_pods( + core=core, + workload_namespace=workload_namespace, + base_name=base_name, + targets=handoff_targets, + health_by_name=health_by_name, + control=control, + logger=logger, + ) + active = set(active_pod_names) ready_ordinals = { ordinal @@ -367,6 +542,7 @@ def _reconcile_pod_set( owner_namespace: str, logger: Logger, previous_generations: dict[int, int], + control: ControlClient, ) -> tuple[PodSetResult, dict[int, int]]: base_name = workload.name @@ -397,6 +573,7 @@ def _reconcile_pod_set( generations=generations, logger=logger, workload_set=workload_set, + control=control, ) return pod_result, generations @@ -441,7 +618,11 @@ def reconcile_pipeline( status: PipelineStatusPatch | None = None, previous_conditions: list[V1ConditionDict] | None = None, previous_generations: object = None, + control_host: str, + control_port: int, + control: ControlClient | None = None, ) -> CombinedPodResult: + control = control if control is not None else ControlClient(control_port) status_patch = ( status if status is not None @@ -456,7 +637,7 @@ def reconcile_pipeline( consumer = from_crd_spec(dict(spec), name=name) try: validate(consumer) - result = render_pods(consumer) + result = render_pods(consumer, control_host, control_port) except Exception as e: if status_patch is not None: failed = [ @@ -500,6 +681,7 @@ def reconcile_pipeline( owner_namespace=namespace, logger=logger, previous_generations=_parse_generations(ledger.get(workload_set)), + control=control, ) pod_set_results[workload_set] = set_result generations_by_set[workload_set] = generations diff --git a/sentry_streams_k8s/sentry_streams_k8s/operator/streaming_pipeline.py b/sentry_streams_k8s/sentry_streams_k8s/operator/streaming_pipeline.py index b6faa3fc..1ac5f4ba 100644 --- a/sentry_streams_k8s/sentry_streams_k8s/operator/streaming_pipeline.py +++ b/sentry_streams_k8s/sentry_streams_k8s/operator/streaming_pipeline.py @@ -82,8 +82,8 @@ def render_deployments(spec: StreamingPipelineSpec) -> RenderedDeployments: return builder.build_deployments(consumer, spec["pipeline_config"]) -def render_pods(spec: StreamingPipelineSpec) -> RenderedPods: +def render_pods(spec: StreamingPipelineSpec, control_host: str, control_port: int) -> RenderedPods: builder = ConsumerBuilder(spec["deployment_template"], spec["container_template"]) consumer = to_consumer_spec(spec) builder.validate(consumer, spec["pipeline_config"]) - return builder.build_pods(consumer, spec["pipeline_config"]) + return builder.build_pods(consumer, spec["pipeline_config"], control_host, control_port) diff --git a/sentry_streams_k8s/tests/k8s_fixtures.py b/sentry_streams_k8s/tests/k8s_fixtures.py index 2cc9912d..1ad1658f 100644 --- a/sentry_streams_k8s/tests/k8s_fixtures.py +++ b/sentry_streams_k8s/tests/k8s_fixtures.py @@ -21,6 +21,7 @@ ) from sentry_streams_k8s.k8s_types import V1ConfigMapDict, V1PodDict +from sentry_streams_k8s.operator.control_client import ControlError, RuntimeState def _matches_selector(labels: Mapping[str, str] | None, selector: str | None) -> bool: @@ -193,6 +194,7 @@ def make_pod( init_container_statuses: list[V1ContainerStatus] | None = None, start_time: datetime | None = None, reason: str | None = None, + pod_ip: str | None = None, ) -> V1Pod: return V1Pod( metadata=V1ObjectMeta( @@ -209,5 +211,38 @@ def make_pod( init_container_statuses=init_container_statuses, start_time=start_time, reason=reason, + pod_ip=pod_ip, ), ) + + +@dataclass +class FakeControlClient: + """In-memory control client used by operator tests.""" + + states: dict[str, RuntimeState] = field(default_factory=dict) + unreachable: set[str] = field(default_factory=set) + + started: list[tuple[str, str]] = field(default_factory=list, init=False) + stopped: list[str] = field(default_factory=list, init=False) + + def status(self, ip: str) -> RuntimeState | None: + if ip in self.unreachable: + return None + return self.states.get(ip) + + def readyz(self, ip: str) -> bool: + return ip not in self.unreachable + + def start(self, ip: str, group_instance_id: str) -> None: + if ip in self.unreachable: + raise ControlError(f"cannot reach {ip}") + self.started.append((ip, group_instance_id)) + self.states[ip] = RuntimeState.CONSUMING + + def stop(self, ip: str) -> bool: + if ip in self.unreachable: + return False + self.stopped.append(ip) + self.states[ip] = RuntimeState.STOPPED + return True diff --git a/sentry_streams_k8s/tests/test_operator.py b/sentry_streams_k8s/tests/test_operator.py index 35f77cbd..db739a73 100644 --- a/sentry_streams_k8s/tests/test_operator.py +++ b/sentry_streams_k8s/tests/test_operator.py @@ -45,6 +45,9 @@ WORKLOAD_NAMESPACE = "test-streaming-pipelines" +CONTROL_HOST = "127.0.0.5" +CONTROL_PORT = 9137 + def test_prepare_manifest_routes_workload_and_records_source_cr() -> None: manifest = { @@ -205,7 +208,7 @@ def _stub_render(monkeypatch: pytest.MonkeyPatch) -> MagicMock: monkeypatch.setattr("sentry_streams_k8s.operator.reconcile.validate", lambda _consumer: None) monkeypatch.setattr( "sentry_streams_k8s.operator.reconcile.render_pods", - lambda _consumer: { + lambda _consumer, _host, _port: { "configmap": configmap, "sets": { PRIMARY_WORKLOAD_SET: _workload("consumer", 1), @@ -232,6 +235,8 @@ def test_reconcile_applies_pods_and_reports_status_through_a_plain_dict( namespace="source", uid="owner-uid", workload_namespace=WORKLOAD_NAMESPACE, + control_host=CONTROL_HOST, + control_port=CONTROL_PORT, logger=MagicMock(), status=status, ) @@ -268,7 +273,7 @@ def test_reconcile_nulls_out_a_workload_set_that_is_no_longer_rendered( core = _stub_render(monkeypatch) monkeypatch.setattr( "sentry_streams_k8s.operator.reconcile.render_pods", - lambda _consumer: { + lambda _consumer, _host, _port: { "configmap": { "apiVersion": "v1", "kind": "ConfigMap", @@ -285,6 +290,8 @@ def test_reconcile_nulls_out_a_workload_set_that_is_no_longer_rendered( namespace="source", uid="owner-uid", workload_namespace=WORKLOAD_NAMESPACE, + control_host=CONTROL_HOST, + control_port=CONTROL_PORT, logger=MagicMock(), status=status, previous_generations={PRIMARY_WORKLOAD_SET: {"0": 4}, CANARY_WORKLOAD_SET: {"0": 2}}, @@ -314,6 +321,8 @@ def test_reconcile_records_render_failure_in_the_status_dict( namespace="source", uid="owner-uid", workload_namespace=WORKLOAD_NAMESPACE, + control_host=CONTROL_HOST, + control_port=CONTROL_PORT, logger=MagicMock(), status=status, ) @@ -418,6 +427,8 @@ def test_reconcile_once_publishes_status_and_schedules_the_health_scan( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("WORKLOAD_NAMESPACE", WORKLOAD_NAMESPACE) + monkeypatch.setenv("CONTROL_HOST", CONTROL_HOST) + monkeypatch.setenv("CONTROL_PORT", str(CONTROL_PORT)) patch_status = MagicMock() monkeypatch.setattr(operator_module, "_patch_pipeline_status", patch_status) @@ -438,6 +449,8 @@ def test_reconcile_once_reports_a_permanent_error_without_scheduling_a_retry( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("WORKLOAD_NAMESPACE", WORKLOAD_NAMESPACE) + monkeypatch.setenv("CONTROL_HOST", CONTROL_HOST) + monkeypatch.setenv("CONTROL_PORT", str(CONTROL_PORT)) patch_status = MagicMock() monkeypatch.setattr(operator_module, "_patch_pipeline_status", patch_status) @@ -456,6 +469,8 @@ def test_reconcile_once_retries_soon_after_an_unexpected_error( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("WORKLOAD_NAMESPACE", WORKLOAD_NAMESPACE) + monkeypatch.setenv("CONTROL_HOST", CONTROL_HOST) + monkeypatch.setenv("CONTROL_PORT", str(CONTROL_PORT)) monkeypatch.setattr(operator_module, "_patch_pipeline_status", MagicMock()) monkeypatch.setattr( operator_module, "reconcile_pipeline", MagicMock(side_effect=RuntimeError("boom")) @@ -468,6 +483,8 @@ def test_reconcile_once_skips_the_pass_when_already_stopped( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("WORKLOAD_NAMESPACE", WORKLOAD_NAMESPACE) + monkeypatch.setenv("CONTROL_HOST", CONTROL_HOST) + monkeypatch.setenv("CONTROL_PORT", str(CONTROL_PORT)) reconcile = MagicMock() monkeypatch.setattr(operator_module, "reconcile_pipeline", reconcile) stopped = FakeStopped() @@ -574,6 +591,8 @@ def test_cleanup_deletes_owned_pods_and_prunes_configmaps( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("WORKLOAD_NAMESPACE", WORKLOAD_NAMESPACE) + monkeypatch.setenv("CONTROL_HOST", CONTROL_HOST) + monkeypatch.setenv("CONTROL_PORT", str(CONTROL_PORT)) core = MagicMock() delete_pods = MagicMock() prune = MagicMock() @@ -631,3 +650,32 @@ async def scenario() -> tuple[bool, bool]: return healthy_woke, event.is_set() assert asyncio.run(scenario()) == (False, True) + + +def test_pod_event_wakes_when_a_pod_becomes_ready() -> None: + scheduler = ReconcileScheduler() + memo = SimpleNamespace(reconcile_scheduler=scheduler) + + async def scenario() -> bool: + event = scheduler.register("uid") + await handle_pipeline_pod_event( + type="MODIFIED", + body=kopf.Body( + { + "metadata": {"name": "consumer-0-0"}, + "status": { + "phase": "Running", + "conditions": [{"type": "Ready", "status": "True"}], + }, + } + ), + meta=kopf.Meta({}), + labels={OWNER_UID_LABEL: "uid"}, + name="consumer-0-0", + namespace=WORKLOAD_NAMESPACE, + memo=memo, + logger=MagicMock(), + ) + return event.is_set() + + assert asyncio.run(scenario()) is True diff --git a/sentry_streams_k8s/tests/test_pipeline_step.py b/sentry_streams_k8s/tests/test_pipeline_step.py index 2e4b519f..56669cfb 100644 --- a/sentry_streams_k8s/tests/test_pipeline_step.py +++ b/sentry_streams_k8s/tests/test_pipeline_step.py @@ -140,6 +140,8 @@ def test_build_container() -> None: cpu_per_process=1000, memory_per_process=512, segment_id=0, + control_host=None, + control_port=None, ) assert container == { @@ -208,6 +210,8 @@ def test_build_container_custom_name() -> None: memory_per_process=512, segment_id=0, container_name="my-custom-container", + control_host=None, + control_port=None, ) assert container["name"] == "my-custom-container" @@ -224,6 +228,8 @@ def test_build_container_with_log_level() -> None: memory_per_process=512, segment_id=0, log_level="ERROR", + control_host=None, + control_port=None, ) assert container["args"] == [ @@ -891,6 +897,8 @@ def test_build_container_with_multiprocessing() -> None: memory_per_process=512, segment_id=0, process_count=4, + control_host=None, + control_port=None, ) # Check resources are multiplied by process count @@ -925,6 +933,8 @@ def test_build_container_without_multiprocessing() -> None: memory_per_process=512, segment_id=0, process_count=None, + control_host=None, + control_port=None, ) # Check resources are NOT multiplied diff --git a/sentry_streams_k8s/tests/test_pods.py b/sentry_streams_k8s/tests/test_pods.py index aa96494e..185b894f 100644 --- a/sentry_streams_k8s/tests/test_pods.py +++ b/sentry_streams_k8s/tests/test_pods.py @@ -3,7 +3,7 @@ import copy import logging from datetime import datetime, timezone -from typing import Any +from typing import Any, cast import pytest from kubernetes.client import V1Pod @@ -21,6 +21,7 @@ SPEC_HASH_ANNOTATION, WORKLOAD_SET_LABEL, ) +from sentry_streams_k8s.operator.control_client import RuntimeState from sentry_streams_k8s.operator.pod_health import PodHealth from sentry_streams_k8s.operator.pod_resources import ( build_pipeline_pod, @@ -39,6 +40,7 @@ reconcile_pipeline_pods, ) from tests.k8s_fixtures import ( + FakeControlClient, FakeCoreV1Api, make_condition, make_pod, @@ -268,6 +270,7 @@ def _observed( container_statuses: list[Any] | None = None, start_time: datetime | None = None, deletion_timestamp: datetime | None = None, + pod_ip: str | None = None, ) -> V1Pod: metadata = manifest["metadata"] return make_pod( @@ -279,6 +282,7 @@ def _observed( container_statuses=container_statuses, start_time=start_time, deletion_timestamp=deletion_timestamp, + pod_ip=pod_ip, ) @@ -298,6 +302,7 @@ def _reconcile( replicas: int = 1, generations: dict[int, int] | None = None, workload_set: str = PRIMARY_WORKLOAD_SET, + control: FakeControlClient | None = None, ) -> tuple[list[V1PodDict], list[tuple[str, bool]], PodSetResult, dict[int, int]]: core = FakeCoreV1Api(pods=pods) metadata, spec = _template() @@ -315,6 +320,7 @@ def _reconcile( generations=ledger, logger=LOGGER, workload_set=workload_set, + control=cast(Any, control if control is not None else FakeControlClient()), ) return core.applied_pods, core.deleted_pods, result, ledger @@ -369,12 +375,13 @@ def test_reconcile_replaces_a_failed_pod_with_the_next_generation() -> None: ] -def test_reconcile_applies_the_replacement_before_deleting_outdated() -> None: - outdated = _observed(_manifest(0, 1), phase="Running", ready=True) +def test_reconcile_keeps_outdated_consuming_until_the_replacement_exists() -> None: + outdated = _observed(_manifest(0, 1), phase="Running", ready=True, pod_ip="10.0.0.1") assert outdated.metadata is not None outdated.metadata.annotations[SPEC_HASH_ANNOTATION] = "old" core = FakeCoreV1Api(pods=[outdated]) metadata, spec = _template() + control = FakeControlClient(states={"10.0.0.1": RuntimeState.CONSUMING}) reconcile_pipeline_pods( core=core, @@ -389,9 +396,12 @@ def test_reconcile_applies_the_replacement_before_deleting_outdated() -> None: generations={}, logger=LOGGER, workload_set=PRIMARY_WORKLOAD_SET, + control=cast(Any, control), ) - assert core.operations == ["apply:consumer-0-2", "delete:consumer-0-1"] + assert core.operations == ["apply:consumer-0-2"] + assert control.stopped == [] + assert control.started == [] def test_reconcile_keeps_the_best_duplicate_and_deletes_the_rest() -> None: @@ -409,6 +419,109 @@ def test_reconcile_keeps_the_best_duplicate_and_deletes_the_rest() -> None: assert ledger == {0: 2} +def test_reconcile_starts_a_ready_pod_that_has_no_predecessor() -> None: + pod = _observed(_manifest(0, 0), phase="Running", ready=True, pod_ip="10.0.0.9") + control = FakeControlClient(states={"10.0.0.9": RuntimeState.IDLE}) + + applied, deleted, _result, _ledger = _reconcile([pod], control=control) + + assert applied == [] + assert deleted == [] + assert control.started == [("10.0.0.9", "consumer-0")] + + +def test_reconcile_leaves_an_already_consuming_pod_alone() -> None: + pod = _observed(_manifest(0, 0), phase="Running", ready=True, pod_ip="10.0.0.9") + control = FakeControlClient(states={"10.0.0.9": RuntimeState.CONSUMING}) + + applied, deleted, _result, _ledger = _reconcile([pod], control=control) + + assert applied == [] and deleted == [] + assert control.started == [] and control.stopped == [] + + +def test_reconcile_hands_partitions_over_before_deleting_the_outdated_pod() -> None: + outdated = _observed(_manifest(0, 1), phase="Running", ready=True, pod_ip="10.0.0.1") + assert outdated.metadata is not None + outdated.metadata.annotations[SPEC_HASH_ANNOTATION] = "old" + replacement = _observed(_manifest(0, 2), phase="Running", ready=True, pod_ip="10.0.0.2") + core = FakeCoreV1Api(pods=[outdated, replacement]) + metadata, spec = _template() + control = FakeControlClient( + states={"10.0.0.1": RuntimeState.CONSUMING, "10.0.0.2": RuntimeState.IDLE} + ) + + reconcile_pipeline_pods( + core=core, + workload_namespace=NAMESPACE, + owner_uid=OWNER_UID, + owner_name="pipeline", + owner_namespace="source", + base_name="consumer", + template_metadata=metadata, + template_spec=spec, + replicas=1, + generations={}, + logger=LOGGER, + workload_set=PRIMARY_WORKLOAD_SET, + control=cast(Any, control), + ) + + assert control.stopped == ["10.0.0.1"] + assert control.started == [("10.0.0.2", "consumer-0")] + assert core.operations == ["delete:consumer-0-1"] + + +def test_reconcile_does_not_hand_over_to_an_unready_replacement() -> None: + outdated = _observed(_manifest(0, 1), phase="Running", ready=True, pod_ip="10.0.0.1") + assert outdated.metadata is not None + outdated.metadata.annotations[SPEC_HASH_ANNOTATION] = "old" + replacement = _observed(_manifest(0, 2), phase="Pending", ready=False, pod_ip="10.0.0.2") + control = FakeControlClient( + states={"10.0.0.1": RuntimeState.CONSUMING, "10.0.0.2": RuntimeState.IDLE} + ) + + applied, deleted, _result, _ledger = _reconcile([outdated, replacement], control=control) + + assert applied == [] and deleted == [] + assert control.stopped == [] and control.started == [] + + +def test_reconcile_deletes_an_unreachable_predecessor_and_waits() -> None: + outdated = _observed(_manifest(0, 1), phase="Running", ready=True, pod_ip="10.0.0.1") + assert outdated.metadata is not None + outdated.metadata.annotations[SPEC_HASH_ANNOTATION] = "old" + replacement = _observed(_manifest(0, 2), phase="Running", ready=True, pod_ip="10.0.0.2") + control = FakeControlClient(states={"10.0.0.2": RuntimeState.IDLE}, unreachable={"10.0.0.1"}) + + applied, deleted, _result, _ledger = _reconcile([outdated, replacement], control=control) + + assert applied == [] + assert deleted == [("consumer-0-1", False)] + assert control.started == [] + + +def test_reconcile_waits_for_a_terminating_predecessor_to_finish_committing() -> None: + outdated = _observed( + _manifest(0, 1), + phase="Running", + ready=True, + pod_ip="10.0.0.1", + deletion_timestamp=datetime.now(timezone.utc), + ) + assert outdated.metadata is not None + outdated.metadata.annotations[SPEC_HASH_ANNOTATION] = "old" + replacement = _observed(_manifest(0, 2), phase="Running", ready=True, pod_ip="10.0.0.2") + control = FakeControlClient( + states={"10.0.0.1": RuntimeState.STOPPING, "10.0.0.2": RuntimeState.IDLE} + ) + + applied, deleted, _result, _ledger = _reconcile([outdated, replacement], control=control) + + assert applied == [] and deleted == [] + assert control.started == [] + + def test_delete_obsolete_pod_sets_removes_canary_when_disabled() -> None: primary = _pod(0, 1, ready=True, phase="Running") canary = _pod(0, 2, ready=True, phase="Running")