From d129fa7d61dcaf57f2bb3557577a404144dbc19d Mon Sep 17 00:00:00 2001 From: Alex Hunt Date: Tue, 28 Jul 2026 17:42:02 +0100 Subject: [PATCH 1/2] orchestratord: use leader election and run multiple operator replicas Run orchestratord's controllers (materialize, balancer, console) under k8s-controller's lease-based leader election, with a single lease guarding all of them, so that multiple replicas of orchestratord can run while only one replica reconciles at a time. The other replicas keep serving the conversion webhook, which avoids webhook downtime during operator rollouts and node drains. The leader election identity must be unique per replica, otherwise replicas mistake each other's lease renewals for their own and all act as leader at once. It comes from --leader-election-identity, which the chart sets from the pod name via the downward API, and falls back to $HOSTNAME so that a replica is still uniquely identified without the chart's help. Losing the lease drops the controllers, cancelling their in-flight reconciliations, and the replica then rejoins the election in process rather than restarting. Reconciliation is polled inline rather than spawned, so none of it outlives the lease, and staying up keeps the replica serving the conversion webhook. The GCP node upgrade watcher moves under the lease as well. It triggers rollouts by writing the force-rollout annotation, so running it on every replica would let a single node pool upgrade trigger several rollouts of the same instance, each replica writing its own annotation value. Its in-memory dedup map does not coordinate across replicas, and the "rollout already in progress" check it makes is a read followed by a write, so it does not serialize them either. Its Pub/Sub subscriber is polled inline rather than spawned for the same reason the controllers are: a subscriber that outlived the lease would keep pulling and acking notifications out from under the next holder. Its configuration is still validated at startup, so a misconfiguration fails immediately rather than only once the replica wins the election. On SIGTERM the leader stops reconciling, releases the lease so that a standby takes over immediately instead of waiting for the lease to expire, drains the conversion webhook server, and exits. The webhook keeps accepting requests for a few seconds before draining, since Kubernetes removes the pod from the webhook service's endpoints asynchronously. Every step of that sequence is bounded, the lease release included, because its API call could otherwise consume the whole termination grace period and leave no time to drain. The chart sets terminationGracePeriodSeconds explicitly so that the budget the binary assumes cannot drift away from what Kubernetes grants it. The SIGTERM handler is installed before the rest of startup, rather than alongside the controllers that act on it. The webhook server starts early and the readiness probe only checks that server, so the pod can be in the webhook service's endpoints while the rest of startup is still running. With no handler installed yet, the default action for SIGTERM would kill the process there and drop the conversion requests in flight. Tokio records a signal received before the first recv(), so one arriving during startup is still handled once the shutdown path is reached. Note that this does not abort startup, so a SIGTERM during CRD registration is only acted on once startup finishes. Add an orchestratord_is_leader gauge, set when this replica takes the lease and cleared when it stops holding it. Summed across the replicas it should be 1, so a sustained 0 means no replica can take the lease and the operator is reconciling nothing. That was previously visible only as a log line, which matters because the readiness probe checks the conversion webhook, so a replica that cannot reconcile at all still looks healthy. The most likely cause is a service account without permission on leases, which is what an installation managing its own RBAC ends up with after upgrading. Reset environmentd_needs_update when leadership is lost. The gauge is derived from what reconciliation observed and lives for the lifetime of the process, which now outlives its own leadership, since a replica keeps serving the conversion webhook after losing the lease. Left set, a former leader would publish its last observation forever, so summing the gauge across the replicas would count the same organizations once per past leader. Helm chart changes: default to 2 operator replicas, spread them across nodes with a soft pod anti-affinity, create a PodDisruptionBudget tolerating one unavailable replica when running more than one, and grant the operator get/create/update on coordination.k8s.io leases. The chart's unit tests cover the new budget, the anti-affinity and its interaction with a user-configured affinity, the leader election identity, and the termination grace period. Add a leader-failover workflow to test/orchestratord that verifies the lease is held by an operator pod, that leadership fails over when the leader pod is deleted, after a rolling restart of the operator deployment, and when another candidate takes the lease, and that reconciliation completes after each. Enable it in the Nightly pipeline. The candidate case requires the lease's transition count to grow, because the incumbent is still running and eligible there and would otherwise satisfy the check without ever having given leadership up. It also checks that exactly the lease holder reports itself leader through orchestratord_is_leader, scraped through the API server's pod proxy since the operator image has nothing to exec into. Once when leadership is first established, and again after the candidate case, which is the one where the replica that gave leadership up is still running and so has to stop reporting itself leader rather than being replaced by a fresh pod. The upgrade test asserted that exactly one operator pod exists, which breaks now that the chart defaults to two replicas. The released chart versions it upgrades from still install a single replica, so require that every operator pod runs the expected image instead of pinning a pod count. get_orchestratord_data now returns only the pods of the operator deployment's current revision, found by matching the deployment's revision annotation against the ReplicaSet that carries it. A rollout leaves pods of the previous generation running until their replacements are up, so the callers that inspect a single pod's spec could otherwise read the outgoing pod. It rejects an empty result, since that filtering can transiently leave nothing behind and its callers index into the pods. Co-Authored-By: Claude Fable 5 Co-Authored-By: Claude Opus 5 --- Cargo.lock | 6 +- Cargo.toml | 2 +- ci/nightly/pipeline.template.yml | 14 + doc/user/data/metrics.yml | 6 +- misc/helm-charts/operator/README.md | 3 + .../operator/templates/clusterrole.yaml | 7 + .../operator/templates/deployment.yaml | 31 +- .../templates/poddisruptionbudget.yaml | 23 + .../operator/tests/deployment_test.yaml | 50 ++ .../tests/poddisruptionbudget_test.yaml | 34 + misc/helm-charts/operator/values.yaml | 15 + src/orchestratord/Cargo.toml | 2 +- src/orchestratord/src/bin/orchestratord.rs | 662 +++++++++++------- src/orchestratord/src/controller/balancer.rs | 1 + src/orchestratord/src/controller/console.rs | 1 + .../src/controller/materialize.rs | 1 + src/orchestratord/src/gcp_node_upgrade.rs | 30 +- src/orchestratord/src/metrics.rs | 26 +- test/orchestratord/mzcompose.py | 469 ++++++++++++- 19 files changed, 1118 insertions(+), 265 deletions(-) create mode 100644 misc/helm-charts/operator/templates/poddisruptionbudget.yaml create mode 100644 misc/helm-charts/operator/tests/poddisruptionbudget_test.yaml diff --git a/Cargo.lock b/Cargo.lock index 599366bffadf1..f69d9c24ef065 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5303,18 +5303,20 @@ dependencies = [ [[package]] name = "k8s-controller" -version = "0.11.0" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9a1ba25014e0459f311fb50ece23c42905be914b2eaf66ae801feb07a80847c" +checksum = "548329de3a20cd4f697d9d6f99512f8a1b3b1b9526ed05186c70a7bb24a06b41" dependencies = [ "async-trait", "futures", + "k8s-openapi", "kube", "kube-runtime", "rand 0.9.4", "serde", "serde_json", "thiserror 2.0.18", + "tokio", "tracing", ] diff --git a/Cargo.toml b/Cargo.toml index cd86e168ab06d..d34ac9c23275a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -406,7 +406,7 @@ itertools = "0.14.0" jemalloc_pprof = "0.8.2" jsonwebtoken = { version = "10.3.0", features = ["aws_lc_rs"] } junit-report = "0.8.3" -k8s-controller = "0.11.0" +k8s-controller = "0.12.0" k8s-openapi = { version = "0.27.0", features = ["schemars", "v1_32"] } kube = { version = "3.1.0", default-features = false, features = ["client", "derive", "openssl-tls", "runtime", "ws"] } launchdarkly-server-sdk = { version = "3.1.1", default-features = false, features = ["hyper-rustls-native-roots", "crypto-aws-lc-rs"] } diff --git a/ci/nightly/pipeline.template.yml b/ci/nightly/pipeline.template.yml index 52964f2a29bf5..ad3863816c8c8 100644 --- a/ci/nightly/pipeline.template.yml +++ b/ci/nightly/pipeline.template.yml @@ -2644,6 +2644,20 @@ steps: agents: queue: hetzner-aarch64-16cpu-32gb + - id: orchestratord-leader-failover + topics: [self-managed] + label: Orchestratord leader election failover + artifact_paths: ["mz_debug_*.zip"] + depends_on: devel-docker-tags + timeout_in_minutes: 60 + plugins: + - ./ci/plugins/mzcompose: + composition: orchestratord + run: leader-failover + ci-builder: stable + agents: + queue: hetzner-aarch64-16cpu-32gb + - id: orchestratord-upgrade-operator topics: [self-managed, upgrade] label: Orchestratord operator upgrade diff --git a/doc/user/data/metrics.yml b/doc/user/data/metrics.yml index 920b6caf094cb..276de9f8288df 100644 --- a/doc/user/data/metrics.yml +++ b/doc/user/data/metrics.yml @@ -67,7 +67,7 @@ metrics: source: src/environmentd/src/http/metrics.rs visibility: internal - name: environmentd_needs_update - help: Count of organizations in this cluster which are running outdated pod templates + help: Count of organizations in this cluster which are running outdated pod templates. Only the operator replica holding the leadership lease reconciles, so the others report zero. source: src/orchestratord/src/metrics.rs visibility: internal - name: jemalloc_active @@ -4279,6 +4279,10 @@ metrics: - reason source: src/adapter/src/metrics.rs visibility: internal +- name: orchestratord_is_leader + help: Whether this operator replica holds the controller leadership lease, and is therefore the replica reconciling. Summed across the replicas this should be 1. A sustained 0 means no replica can take the lease, for instance because the service account lacks permission on leases, and the operator is reconciling nothing. + source: src/orchestratord/src/metrics.rs + visibility: internal - name: outer_join_lowering_cases help: How many times the different outer join lowering cases happened. labels: diff --git a/misc/helm-charts/operator/README.md b/misc/helm-charts/operator/README.md index 2d055dc024acd..96c018abea780 100644 --- a/misc/helm-charts/operator/README.md +++ b/misc/helm-charts/operator/README.md @@ -176,6 +176,9 @@ The following table lists the configurable parameters of the Materialize operato | `operator.image.repository` | The Docker repository for the operator image | ``"materialize/orchestratord"`` | | `operator.image.tag` | The tag/version of the operator image to be used | ``"v26.34.1"`` | | `operator.nodeSelector` | Node selector to use for the operator pod | ``{}`` | +| `operator.podDisruptionBudget.enabled` | Whether to create a PodDisruptionBudget for the operator. Only created when `replicas` is greater than 1, since a budget over a single replica either blocks node drains or protects nothing. | ``true`` | +| `operator.podDisruptionBudget.maxUnavailable` | Maximum number of operator pods that may be unavailable at once during voluntary disruptions. Expressed as a maximum rather than a minimum so that node drains are never blocked outright, they are only serialized. | ``1`` | +| `operator.replicas` | Number of operator replicas. The operator uses leader election so that only one replica reconciles at a time. Running more than one replica avoids downtime of the CRD conversion webhook during rollouts and node drains. | ``2`` | | `operator.resources.limits` | Resource limits for the operator's CPU and memory | ``{"memory":"512Mi"}`` | | `operator.resources.requests` | Resources requested by the operator for CPU and memory | ``{"cpu":"100m","memory":"512Mi"}`` | | `operator.secretsController` | Which secrets controller to use for storing secrets. Valid values are 'kubernetes' and 'aws-secrets-manager'. Setting 'aws-secrets-manager' requires a configured AWS cloud provider and IAM role for the environment with Secrets Manager permissions. | ``"kubernetes"`` | diff --git a/misc/helm-charts/operator/templates/clusterrole.yaml b/misc/helm-charts/operator/templates/clusterrole.yaml index 617293d59c48c..cb100ef595f25 100644 --- a/misc/helm-charts/operator/templates/clusterrole.yaml +++ b/misc/helm-charts/operator/templates/clusterrole.yaml @@ -37,6 +37,13 @@ rules: - get - list - watch +- apiGroups: ["coordination.k8s.io"] + resources: + - leases + verbs: + - get + - create + - update - apiGroups: ["networking.k8s.io"] resources: - networkpolicies diff --git a/misc/helm-charts/operator/templates/deployment.yaml b/misc/helm-charts/operator/templates/deployment.yaml index 14b4a73a53120..a132024b532c4 100644 --- a/misc/helm-charts/operator/templates/deployment.yaml +++ b/misc/helm-charts/operator/templates/deployment.yaml @@ -16,7 +16,7 @@ metadata: labels: {{- include "materialize-operator.labels" . | nindent 4 }} spec: - replicas: 1 + replicas: {{ .Values.operator.replicas }} selector: matchLabels: {{- include "materialize-operator.selectorLabels" . | nindent 6 }} @@ -32,6 +32,14 @@ spec: runAsUser: 65532 runAsNonRoot: true serviceAccountName: {{ include "materialize-operator.serviceAccountName" . }} + # On SIGTERM the operator releases its leadership lease, then keeps + # serving the conversion webhook for a few seconds while Kubernetes + # removes it from the webhook service's endpoints, then drains the + # in-flight conversion requests. Those steps run in sequence and are + # individually bounded in the binary, adding up to 15 seconds. This is + # set explicitly, rather than left to the Kubernetes default, so that the + # budget the binary assumes cannot drift away from what it is given. + terminationGracePeriodSeconds: 30 {{- if .Values.operator.nodeSelector }} nodeSelector: {{- toYaml .Values.operator.nodeSelector | nindent 8 }} @@ -39,6 +47,20 @@ spec: {{- if .Values.operator.affinity }} affinity: {{- toYaml .Values.operator.affinity | nindent 8 }} + {{- else if gt (int .Values.operator.replicas) 1 }} + # Spread the replicas across nodes by default, so that losing a single + # node can't take down every replica of the conversion webhook. This is + # a preference rather than a requirement, so that the operator still + # schedules on single-node clusters. + affinity: + podAntiAffinity: + preferredDuringSchedulingIgnoredDuringExecution: + - weight: 100 + podAffinityTerm: + topologyKey: kubernetes.io/hostname + labelSelector: + matchLabels: + {{- include "materialize-operator.selectorLabels" . | nindent 18 }} {{- end }} {{- if .Values.operator.tolerations }} tolerations: @@ -271,6 +293,13 @@ spec: - "--webhook-service-namespace" - {{ .Release.Namespace }} {{- end }} + env: + # Sets the operator's --leader-election-identity, which must be unique + # per replica. + - name: ORCHESTRATORD_LEADER_ELECTION_IDENTITY + valueFrom: + fieldRef: + fieldPath: metadata.name ports: - containerPort: 3100 name: metrics diff --git a/misc/helm-charts/operator/templates/poddisruptionbudget.yaml b/misc/helm-charts/operator/templates/poddisruptionbudget.yaml new file mode 100644 index 0000000000000..8d10b8c74019c --- /dev/null +++ b/misc/helm-charts/operator/templates/poddisruptionbudget.yaml @@ -0,0 +1,23 @@ +{{- if and .Values.operator.podDisruptionBudget.enabled (gt (int .Values.operator.replicas) 1) -}} +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# As of the Change Date specified in that file, in accordance with +# the Business Source License, use of this software will be governed +# by the Apache License, Version 2.0. + +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ include "materialize-operator.fullname" . }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "materialize-operator.labels" . | nindent 4 }} +spec: + maxUnavailable: {{ .Values.operator.podDisruptionBudget.maxUnavailable }} + selector: + matchLabels: + {{- include "materialize-operator.selectorLabels" . | nindent 6 }} +{{- end }} diff --git a/misc/helm-charts/operator/tests/deployment_test.yaml b/misc/helm-charts/operator/tests/deployment_test.yaml index d2cb0917996fe..66274c1ab8f11 100644 --- a/misc/helm-charts/operator/tests/deployment_test.yaml +++ b/misc/helm-charts/operator/tests/deployment_test.yaml @@ -244,3 +244,53 @@ tests: - notContains: path: spec.template.spec.containers[0].args content: "--disable-license-key-checks" + +- it: should pass the pod name for the leader election identity + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: ORCHESTRATORD_LEADER_ELECTION_IDENTITY + valueFrom: + fieldRef: + fieldPath: metadata.name + +- it: should spread replicas across nodes by default + asserts: + - equal: + path: spec.template.spec.affinity.podAntiAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].podAffinityTerm.topologyKey + value: kubernetes.io/hostname + - equal: + path: spec.template.spec.affinity.podAntiAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].podAffinityTerm.labelSelector.matchLabels["app.kubernetes.io/instance"] + value: RELEASE-NAME + +- it: should not spread a single replica + set: + operator.replicas: 1 + asserts: + - notExists: + path: spec.template.spec.affinity + +- it: should let a configured affinity win over the default + set: + operator.affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: kubernetes.io/os + operator: In + values: [linux] + asserts: + - notExists: + path: spec.template.spec.affinity.podAntiAffinity + - exists: + path: spec.template.spec.affinity.nodeAffinity + +# The operator's shutdown sequence is bounded to 15 seconds, so this must not +# drop below that. See the template for the breakdown. +- it: should give the operator time to shut down gracefully + asserts: + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 30 diff --git a/misc/helm-charts/operator/tests/poddisruptionbudget_test.yaml b/misc/helm-charts/operator/tests/poddisruptionbudget_test.yaml new file mode 100644 index 0000000000000..c6b1bd8bf5fc5 --- /dev/null +++ b/misc/helm-charts/operator/tests/poddisruptionbudget_test.yaml @@ -0,0 +1,34 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# As of the Change Date specified in that file, in accordance with +# the Business Source License, use of this software will be governed +# by the Apache License, Version 2.0. + +suite: test poddisruptionbudget +templates: +- poddisruptionbudget.yaml +tests: +- it: should create a pod disruption budget by default + asserts: + - isKind: + of: PodDisruptionBudget + - equal: + path: spec.maxUnavailable + value: 1 + +- it: should not create a pod disruption budget with a single replica + set: + operator.replicas: 1 + asserts: + - hasDocuments: + count: 0 + +- it: should not create a pod disruption budget when disabled + set: + operator.podDisruptionBudget.enabled: false + asserts: + - hasDocuments: + count: 0 diff --git a/misc/helm-charts/operator/values.yaml b/misc/helm-charts/operator/values.yaml index c2bbc7a1fde9c..bf62fdd86781f 100644 --- a/misc/helm-charts/operator/values.yaml +++ b/misc/helm-charts/operator/values.yaml @@ -17,6 +17,21 @@ operator: # -- Policy for pulling the image: "IfNotPresent" avoids unnecessary re-pulling of images pullPolicy: IfNotPresent + # -- Number of operator replicas. The operator uses leader election so that + # only one replica reconciles at a time. Running more than one replica + # avoids downtime of the CRD conversion webhook during rollouts and node + # drains. + replicas: 2 + + podDisruptionBudget: + # -- Whether to create a PodDisruptionBudget for the operator. Only created + # when `replicas` is greater than 1, since a budget over a single replica + # either blocks node drains or protects nothing. + enabled: true + # -- Maximum number of operator pods that may be unavailable at once during + # voluntary disruptions. Expressed as a maximum rather than a minimum so + # that node drains are never blocked outright, they are only serialized. + maxUnavailable: 1 args: # -- Log filtering settings for startup logs diff --git a/src/orchestratord/Cargo.toml b/src/orchestratord/Cargo.toml index 9671f9c2cbbbd..85e7f5896f5e2 100644 --- a/src/orchestratord/Cargo.toml +++ b/src/orchestratord/Cargo.toml @@ -14,7 +14,7 @@ anyhow.workspace = true async-trait.workspace = true axum.workspace = true axum-server.workspace = true -clap.workspace = true +clap = { workspace = true, features = ["env"] } futures.workspace = true gcp_auth.workspace = true http.workspace = true diff --git a/src/orchestratord/src/bin/orchestratord.rs b/src/orchestratord/src/bin/orchestratord.rs index 2421eec551e73..ffd089437c073 100644 --- a/src/orchestratord/src/bin/orchestratord.rs +++ b/src/orchestratord/src/bin/orchestratord.rs @@ -8,13 +8,15 @@ // by the Apache License, Version 2.0. use std::{ - future, + env, future, net::SocketAddr, + pin::pin, sync::{Arc, LazyLock}, time::Duration, }; use axum_server::tls_rustls::RustlsConfig; +use futures::future::{Either, join4, select}; use http::HeaderValue; use k8s_openapi::{ api::{ @@ -30,6 +32,8 @@ use k8s_openapi::{ }, }; use kube::{Api, api::ListParams, runtime::watcher}; +use tokio::signal::unix::{SignalKind, signal}; + use mz_cloud_provider::CloudProvider; use mz_cloud_resources::crd::generated::cert_manager::certificates::Certificate; use tracing::info; @@ -53,6 +57,29 @@ use mz_ore::{ const BUILD_INFO: BuildInfo = build_info!(); static VERSION: LazyLock = LazyLock::new(|| BUILD_INFO.human_version(None)); +/// How long to spend releasing the leadership lease during shutdown. Releasing +/// is best-effort, since the lease expires on its own anyway, and it must be +/// bounded because it precedes the webhook drain: the kube client's default +/// read timeout is far longer than the termination grace period, so an +/// unreachable API server would otherwise consume the whole grace period here +/// and leave no time to drain. +const LEASE_RELEASE_TIMEOUT: Duration = Duration::from_secs(5); + +/// How long to keep accepting conversion webhook requests after receiving +/// SIGTERM. Kubernetes removes the pod from the webhook service's endpoints +/// asynchronously, so requests are still routed here for a short while after +/// the signal, and refusing them would fail writes to Materialize resources. +const WEBHOOK_SHUTDOWN_DELAY: Duration = Duration::from_secs(5); + +/// How long to wait for in-flight conversion webhook requests to finish once +/// the server has stopped accepting connections. +/// +/// This, [`LEASE_RELEASE_TIMEOUT`] and [`WEBHOOK_SHUTDOWN_DELAY`] run in +/// sequence during shutdown, so their sum must stay well inside the pod's +/// termination grace period, otherwise Kubernetes kills the process mid-drain. +/// Our Helm chart sets that grace period explicitly to keep the two in sync. +const WEBHOOK_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); + #[derive(clap::Parser)] #[clap(name = "orchestratord", version = VERSION.as_str())] pub struct Args { @@ -93,6 +120,17 @@ pub struct Args { #[clap(long, default_value = "1h", value_parser = humantime::parse_duration)] webhook_cert_reload_interval: Duration, + /// Identity for the controllers' leader election, which must be unique per + /// replica. Replicas sharing an identity each mistake the others' lease + /// renewals for their own and all act as leader simultaneously, with no + /// error to indicate it. + /// + /// Defaults to `$HOSTNAME`, which in a pod is the pod name unless the pod + /// sets `spec.hostname`, and then to a random identity so that running + /// outside of a pod works too. + #[clap(long, env = "LEADER_ELECTION_IDENTITY")] + leader_election_identity: Option, + #[clap(long)] cloud_provider: CloudProvider, #[clap(long)] @@ -318,6 +356,20 @@ async fn main() { } async fn run(args: Args) -> Result<(), anyhow::Error> { + // Installed before anything else, because the default action for SIGTERM + // kills the process outright. The conversion webhook server starts below, + // and the pod's readiness probe only checks that server, so this pod can be + // in the webhook service's endpoints well before the rest of startup + // finishes. Without a handler in place, a SIGTERM in that window would drop + // the conversion requests already in flight. Tokio records a signal + // received before the first `recv()`, so one arriving during startup is + // still handled once the shutdown path below is reached. + // + // NOTE: installing the handler does not abort startup. A SIGTERM during + // `register_crds`, which is bounded at 120 seconds, is only acted on once + // startup finishes, which can outlast the pod's termination grace period. + let mut sigterm = signal(SignalKind::terminate()).expect("failed to install SIGTERM handler"); + let metrics_registry = MetricsRegistry::new(); args.tracing .configure_tracing( @@ -334,7 +386,10 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { let tls_cert = args.tls_cert; let tls_key = args.tls_key; let tls_ca = args.tls_ca; - let reload_config = if args.install_v1_crd { + // `reload_config` is the TLS config to reload rotated certificates into, + // and `webhook_server` is what shutdown uses to drain the server. Both are + // set only when the conversion webhook server is running. + let (reload_config, webhook_server) = if args.install_v1_crd { // Pin the rustls crypto provider to aws-lc-rs. `RustlsConfig` builds its // `ServerConfig` via `ServerConfig::builder()`, which resolves the // process-default provider. Installing it explicitly keeps the choice @@ -350,18 +405,23 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { let reload_config = config.clone(); let webhook_listen_address = args.webhook_listen_address; - mz_ore::task::spawn(|| "webhook server", async move { - if let Err(e) = axum_server::bind_rustls(webhook_listen_address, config) - .serve(webhook::router().into_make_service()) - .await - { - panic!("webhook server failed: {}", e.display_with_causes()); + let webhook_handle = axum_server::Handle::new(); + let webhook_task = mz_ore::task::spawn(|| "webhook server", { + let webhook_handle = webhook_handle.clone(); + async move { + if let Err(e) = axum_server::bind_rustls(webhook_listen_address, config) + .handle(webhook_handle) + .serve(webhook::router().into_make_service()) + .await + { + panic!("webhook server failed: {}", e.display_with_causes()); + } } }); - Some(reload_config) + (Some(reload_config), Some((webhook_handle, webhook_task))) } else { - None + (None, None) }; let (client, namespace) = create_client(args.kubernetes_context.clone()).await?; @@ -490,7 +550,12 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { }); } - if args.enable_gcp_node_upgrade_rollout_trigger { + // Validated and built here rather than where the watcher runs, so that a + // misconfiguration fails startup instead of only surfacing once this + // replica wins the election. The watcher itself runs under the leadership + // lease, since it triggers rollouts of Materialize instances and only one + // replica may do that. + let gcp_node_upgrade_config = if args.enable_gcp_node_upgrade_rollout_trigger { anyhow::ensure!( args.cloud_provider == CloudProvider::Gcp, "--enable-gcp-node-upgrade-rollout-trigger requires --cloud-provider=gcp" @@ -500,7 +565,7 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { "--enable-gcp-node-upgrade-rollout-trigger requires --install-v1-crd, since \ rollouts are triggered through the v1 Materialize CRD" ); - let config = gcp_node_upgrade::Config::new( + Some(gcp_node_upgrade::Config::new( args.gcp_node_upgrade_notification_subscription .clone() .expect("clap requires --gcp-node-upgrade-notification-subscription"), @@ -511,243 +576,354 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { .clone() .expect("clap requires --gcp-cluster-location"), args.gcp_node_upgrade_watched_node_pools.clone(), - )?; - mz_ore::task::spawn( - || "gcp node upgrade watcher", - gcp_node_upgrade::run(client.clone(), config), - ); - } + )?) + } else { + None + }; - mz_ore::task::spawn( - || "materialize controller", - k8s_controller::Controller::namespaced_all( - client.clone(), - controller::materialize::Context::new( - controller::materialize::Config { - cloud_provider: args.cloud_provider, - region: args.region, - create_balancers: args.create_balancers, - create_console: args.create_console, - helm_chart_version: args.helm_chart_version, - secrets_controller: args.secrets_controller, - collect_pod_metrics: args.collect_pod_metrics, - enable_prometheus_scrape_annotations: args.enable_prometheus_scrape_annotations, - segment_api_key: args.segment_api_key, - segment_client_side: args.segment_client_side, - console_image_tag_default: args.console_image_tag_default, - console_image_tag_map: args.console_image_tag_map, - aws_account_id: args.aws_account_id, - environmentd_iam_role_arn: args.environmentd_iam_role_arn, - environmentd_connection_role_arn: args.environmentd_connection_role_arn, - aws_secrets_controller_tags: args.aws_secrets_controller_tags, - environmentd_availability_zones: args.environmentd_availability_zones, - ephemeral_volume_class: args.ephemeral_volume_class, - scheduler_name: args.scheduler_name.clone(), - enable_security_context: args.enable_security_context, - enable_internal_statement_logging: args.enable_internal_statement_logging, - disable_statement_logging: args.disable_statement_logging, - orchestratord_pod_selector_labels: args.orchestratord_pod_selector_labels, - environmentd_node_selector: args.environmentd_node_selector, - environmentd_affinity: args.environmentd_affinity, - environmentd_tolerations: args.environmentd_tolerations, - environmentd_default_resources: args.environmentd_default_resources, - clusterd_node_selector: args.clusterd_node_selector, - clusterd_affinity: args.clusterd_affinity, - clusterd_tolerations: args.clusterd_tolerations, - image_pull_policy: args.image_pull_policy, - network_policies_internal_enabled: args.network_policies_internal_enabled, - network_policies_ingress_enabled: args.network_policies_ingress_enabled, - network_policies_ingress_cidrs: args.network_policies_ingress_cidrs.clone(), - network_policies_egress_enabled: args.network_policies_egress_enabled, - network_policies_egress_cidrs: args.network_policies_egress_cidrs, - environmentd_cluster_replica_sizes: args.environmentd_cluster_replica_sizes, - bootstrap_default_cluster_replica_size: args - .bootstrap_default_cluster_replica_size, - bootstrap_builtin_system_cluster_replica_size: args - .bootstrap_builtin_system_cluster_replica_size, - bootstrap_builtin_probe_cluster_replica_size: args - .bootstrap_builtin_probe_cluster_replica_size, - bootstrap_builtin_support_cluster_replica_size: args - .bootstrap_builtin_support_cluster_replica_size, - bootstrap_builtin_catalog_server_cluster_replica_size: args - .bootstrap_builtin_catalog_server_cluster_replica_size, - bootstrap_builtin_analytics_cluster_replica_size: args - .bootstrap_builtin_analytics_cluster_replica_size, - bootstrap_builtin_system_cluster_replication_factor: args - .bootstrap_builtin_system_cluster_replication_factor, - bootstrap_builtin_probe_cluster_replication_factor: args - .bootstrap_builtin_probe_cluster_replication_factor, - bootstrap_builtin_support_cluster_replication_factor: args - .bootstrap_builtin_support_cluster_replication_factor, - bootstrap_builtin_analytics_cluster_replication_factor: args - .bootstrap_builtin_analytics_cluster_replication_factor, - environmentd_allowed_origins: args.environmentd_allowed_origins, - internal_console_proxy_url: args.internal_console_proxy_url, - environmentd_sql_port: args.environmentd_sql_port, - environmentd_http_port: args.environmentd_http_port, - environmentd_internal_sql_port: args.environmentd_internal_sql_port, - environmentd_internal_http_port: args.environmentd_internal_http_port, - environmentd_internal_persist_pubsub_port: args - .environmentd_internal_persist_pubsub_port, - default_certificate_specs: args.default_certificate_specs.clone(), - disable_license_key_checks: args.disable_license_key_checks, - tracing: args.tracing, - orchestratord_namespace: namespace, - }, - Arc::clone(&metrics), - ), - watcher::Config::default().timeout(29), - ) - .with_controller(|controller| { - let controller = controller - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ); - if has_cert_manager { - controller.owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - } else { - controller - } - }) - .run(), + // Leader election allows running multiple replicas of orchestratord + // (e.g. to avoid downtime of the conversion webhook during rollouts) + // while ensuring that the controllers reconcile on only one replica at + // a time. A single lease guards all of the controllers, so that they + // can't end up scattered across replicas. + // + // See `Args::leader_election_identity` for what the identity must satisfy. + // Our Helm chart sets it from the pod name via the downward API. + let leader_election_identity = args + .leader_election_identity + .or_else(|| env::var("HOSTNAME").ok()) + .unwrap_or_else(|| format!("orchestratord-{}", uuid::Uuid::new_v4())); + let leader_election = k8s_controller::LeaderElection::new( + client.clone(), + &namespace, + "orchestratord", + &leader_election_identity, ); - mz_ore::task::spawn( - || "balancer controller", - k8s_controller::Controller::namespaced_all( - client.clone(), - controller::balancer::Context::new(controller::balancer::Config { - enable_security_context: args.enable_security_context, - enable_prometheus_scrape_annotations: args.enable_prometheus_scrape_annotations, - image_pull_policy: args.image_pull_policy, - scheduler_name: args.scheduler_name.clone(), - balancerd_node_selector: args.balancerd_node_selector, - balancerd_affinity: args.balancerd_affinity, - balancerd_tolerations: args.balancerd_tolerations, - balancerd_default_resources: args.balancerd_default_resources, - default_certificate_specs: args.default_certificate_specs.clone(), - environmentd_sql_port: args.environmentd_sql_port, - environmentd_http_port: args.environmentd_http_port, - balancerd_sql_port: args.balancerd_sql_port, - balancerd_http_port: args.balancerd_http_port, - balancerd_internal_http_port: args.balancerd_internal_http_port, - }), - watcher::Config::default().timeout(29), - ) - .with_controller(|controller| { - let controller = controller - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ); - if has_cert_manager { - controller.owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - } else { - controller + // Each of these is rebuilt every time this replica wins the election, since + // running a controller consumes it and we rejoin the election after losing + // the lease. + let make_materialize_controller = { + let client = client.clone(); + let metrics = Arc::clone(&metrics); + let config = controller::materialize::Config { + cloud_provider: args.cloud_provider, + region: args.region, + create_balancers: args.create_balancers, + create_console: args.create_console, + helm_chart_version: args.helm_chart_version, + secrets_controller: args.secrets_controller, + collect_pod_metrics: args.collect_pod_metrics, + enable_prometheus_scrape_annotations: args.enable_prometheus_scrape_annotations, + segment_api_key: args.segment_api_key, + segment_client_side: args.segment_client_side, + console_image_tag_default: args.console_image_tag_default, + console_image_tag_map: args.console_image_tag_map, + aws_account_id: args.aws_account_id, + environmentd_iam_role_arn: args.environmentd_iam_role_arn, + environmentd_connection_role_arn: args.environmentd_connection_role_arn, + aws_secrets_controller_tags: args.aws_secrets_controller_tags, + environmentd_availability_zones: args.environmentd_availability_zones, + ephemeral_volume_class: args.ephemeral_volume_class, + scheduler_name: args.scheduler_name.clone(), + enable_security_context: args.enable_security_context, + enable_internal_statement_logging: args.enable_internal_statement_logging, + disable_statement_logging: args.disable_statement_logging, + orchestratord_pod_selector_labels: args.orchestratord_pod_selector_labels, + environmentd_node_selector: args.environmentd_node_selector, + environmentd_affinity: args.environmentd_affinity, + environmentd_tolerations: args.environmentd_tolerations, + environmentd_default_resources: args.environmentd_default_resources, + clusterd_node_selector: args.clusterd_node_selector, + clusterd_affinity: args.clusterd_affinity, + clusterd_tolerations: args.clusterd_tolerations, + image_pull_policy: args.image_pull_policy, + network_policies_internal_enabled: args.network_policies_internal_enabled, + network_policies_ingress_enabled: args.network_policies_ingress_enabled, + network_policies_ingress_cidrs: args.network_policies_ingress_cidrs.clone(), + network_policies_egress_enabled: args.network_policies_egress_enabled, + network_policies_egress_cidrs: args.network_policies_egress_cidrs, + environmentd_cluster_replica_sizes: args.environmentd_cluster_replica_sizes, + bootstrap_default_cluster_replica_size: args.bootstrap_default_cluster_replica_size, + bootstrap_builtin_system_cluster_replica_size: args + .bootstrap_builtin_system_cluster_replica_size, + bootstrap_builtin_probe_cluster_replica_size: args + .bootstrap_builtin_probe_cluster_replica_size, + bootstrap_builtin_support_cluster_replica_size: args + .bootstrap_builtin_support_cluster_replica_size, + bootstrap_builtin_catalog_server_cluster_replica_size: args + .bootstrap_builtin_catalog_server_cluster_replica_size, + bootstrap_builtin_analytics_cluster_replica_size: args + .bootstrap_builtin_analytics_cluster_replica_size, + bootstrap_builtin_system_cluster_replication_factor: args + .bootstrap_builtin_system_cluster_replication_factor, + bootstrap_builtin_probe_cluster_replication_factor: args + .bootstrap_builtin_probe_cluster_replication_factor, + bootstrap_builtin_support_cluster_replication_factor: args + .bootstrap_builtin_support_cluster_replication_factor, + bootstrap_builtin_analytics_cluster_replication_factor: args + .bootstrap_builtin_analytics_cluster_replication_factor, + environmentd_allowed_origins: args.environmentd_allowed_origins, + internal_console_proxy_url: args.internal_console_proxy_url, + environmentd_sql_port: args.environmentd_sql_port, + environmentd_http_port: args.environmentd_http_port, + environmentd_internal_sql_port: args.environmentd_internal_sql_port, + environmentd_internal_http_port: args.environmentd_internal_http_port, + environmentd_internal_persist_pubsub_port: args + .environmentd_internal_persist_pubsub_port, + default_certificate_specs: args.default_certificate_specs.clone(), + disable_license_key_checks: args.disable_license_key_checks, + tracing: args.tracing, + orchestratord_namespace: namespace, + }; + move || { + k8s_controller::Controller::namespaced_all( + client.clone(), + controller::materialize::Context::new(config.clone(), Arc::clone(&metrics)), + watcher::Config::default().timeout(29), + ) + .with_controller(|controller| { + let controller = controller + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ); + if has_cert_manager { + controller.owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + } else { + controller + } + }) + } + }; + let make_balancer_controller = { + let client = client.clone(); + let config = controller::balancer::Config { + enable_security_context: args.enable_security_context, + enable_prometheus_scrape_annotations: args.enable_prometheus_scrape_annotations, + image_pull_policy: args.image_pull_policy, + scheduler_name: args.scheduler_name.clone(), + balancerd_node_selector: args.balancerd_node_selector, + balancerd_affinity: args.balancerd_affinity, + balancerd_tolerations: args.balancerd_tolerations, + balancerd_default_resources: args.balancerd_default_resources, + default_certificate_specs: args.default_certificate_specs.clone(), + environmentd_sql_port: args.environmentd_sql_port, + environmentd_http_port: args.environmentd_http_port, + balancerd_sql_port: args.balancerd_sql_port, + balancerd_http_port: args.balancerd_http_port, + balancerd_internal_http_port: args.balancerd_internal_http_port, + }; + move || { + k8s_controller::Controller::namespaced_all( + client.clone(), + controller::balancer::Context::new(config.clone()), + watcher::Config::default().timeout(29), + ) + .with_controller(|controller| { + let controller = controller + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ); + if has_cert_manager { + controller.owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + } else { + controller + } + }) + } + }; + let make_console_controller = { + let client = client.clone(); + let config = controller::console::Config { + enable_security_context: args.enable_security_context, + enable_prometheus_scrape_annotations: args.enable_prometheus_scrape_annotations, + image_pull_policy: args.image_pull_policy, + scheduler_name: args.scheduler_name, + console_node_selector: args.console_node_selector, + console_affinity: args.console_affinity, + console_tolerations: args.console_tolerations, + console_default_resources: args.console_default_resources, + network_policies_ingress_enabled: args.network_policies_ingress_enabled, + network_policies_ingress_cidrs: args.network_policies_ingress_cidrs, + default_certificate_specs: args.default_certificate_specs, + console_http_port: args.console_http_port, + balancerd_http_port: args.balancerd_http_port, + }; + move || { + k8s_controller::Controller::namespaced_all( + client.clone(), + controller::console::Context::new(config.clone()), + watcher::Config::default().timeout(29), + ) + .with_controller(|controller| { + let controller = controller + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + .owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ); + if has_cert_manager { + controller.owns( + Api::::all(client.clone()), + watcher::Config::default() + .labels("materialize.cloud/mz-resource-id") + .timeout(29), + ) + } else { + controller + } + }) + } + }; + let make_gcp_node_upgrade_watcher = { + let client = client.clone(); + move || { + let client = client.clone(); + let config = gcp_node_upgrade_config.clone(); + async move { + match config { + Some(config) => gcp_node_upgrade::run(client, config).await, + // The trigger is disabled, so there is nothing to watch. + // This future must still never complete, since completing + // is how the joined futures below report that they stopped. + None => future::pending().await, + } } - }) - .run(), - ); - - mz_ore::task::spawn( - || "console controller", - k8s_controller::Controller::namespaced_all( - client.clone(), - controller::console::Context::new(controller::console::Config { - enable_security_context: args.enable_security_context, - enable_prometheus_scrape_annotations: args.enable_prometheus_scrape_annotations, - image_pull_policy: args.image_pull_policy, - scheduler_name: args.scheduler_name, - console_node_selector: args.console_node_selector, - console_affinity: args.console_affinity, - console_tolerations: args.console_tolerations, - console_default_resources: args.console_default_resources, - network_policies_ingress_enabled: args.network_policies_ingress_enabled, - network_policies_ingress_cidrs: args.network_policies_ingress_cidrs, - default_certificate_specs: args.default_certificate_specs, - console_http_port: args.console_http_port, - balancerd_http_port: args.balancerd_http_port, - }), - watcher::Config::default().timeout(29), - ) - .with_controller(|controller| { - let controller = controller - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ) - .owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), - ); - if has_cert_manager { - controller.owns( - Api::::all(client.clone()), - watcher::Config::default() - .labels("materialize.cloud/mz-resource-id") - .timeout(29), + } + }; + mz_ore::task::spawn(|| "controllers", async move { + loop { + // All of the controllers, and the node upgrade watcher, run under + // the single leadership lease, on whichever replica holds it. The + // watcher belongs here because it triggers rollouts, which only the + // leader may do. None of these futures ever complete on their own. + // + // `with_lease` polls this future only once it holds the lease, so + // building the controllers here, rather than passing them in, ties + // both them and the metric to the moment leadership is acquired. + let controllers = Box::pin(leader_election.with_lease(async { + metrics.leadership_acquired(); + join4( + make_materialize_controller().run(), + make_balancer_controller().run(), + make_console_controller().run(), + make_gcp_node_upgrade_watcher(), ) - } else { - controller + .await + })); + match select(controllers, pin!(sigterm.recv())).await { + Either::Left((None, _)) => { + // Losing the lease dropped the futures above, which cancels + // the in-flight reconciliations and the node upgrade watch. + // Both are polled inline rather than spawned, so no work + // outlives the lease and we can rejoin the election in + // process instead of restarting. Staying up keeps this + // replica serving the conversion webhook. + // + // The leader-only gauges describe what this replica saw + // while it was reconciling, and this process outlives its + // own leadership, so leaving them set would keep publishing + // a stale count alongside the new leader's. + metrics.leadership_lost(); + tracing::warn!("lost leadership lease, rejoining the election"); + } + Either::Left((Some(_), _)) => { + // The controllers only finish if their watches end, which + // leaves this process unable to reconcile anything. + // `with_lease` released the lease before returning, so a + // standby has taken leadership over already. Let Kubernetes + // restart us rather than rejoining the election only to + // fall over the same way again. + tracing::error!("controllers exited unexpectedly, exiting"); + std::process::exit(1); + } + Either::Right((_, controllers)) => { + // Stop reconciling before releasing the lease, so that no + // reconciliation outlives our leadership. Releasing hands + // leadership over immediately rather than making the other + // replicas wait for the lease to expire. + drop(controllers); + // The webhook keeps serving for a while below, so this + // replica is still scraped after it stops reconciling. + metrics.leadership_lost(); + let _ = tokio::time::timeout(LEASE_RELEASE_TIMEOUT, leader_election.release()) + .await; + if let Some((webhook_handle, webhook_task)) = webhook_server { + // Removing this pod from the webhook service's + // endpoints is not synchronous with SIGTERM, so + // conversion requests keep arriving for a while after + // it. Keep accepting them until the endpoints have + // caught up, then stop accepting and let the in-flight + // requests finish. + tokio::time::sleep(WEBHOOK_SHUTDOWN_DELAY).await; + webhook_handle.graceful_shutdown(Some(WEBHOOK_DRAIN_TIMEOUT)); + let _ = webhook_task.await; + } + info!("received SIGTERM, shut down cleanly"); + std::process::exit(0); + } } - }) - .run(), - ); + } + }); info!("All tasks started successfully."); diff --git a/src/orchestratord/src/controller/balancer.rs b/src/orchestratord/src/controller/balancer.rs index 8687e99dd57d9..367151240b5d0 100644 --- a/src/orchestratord/src/controller/balancer.rs +++ b/src/orchestratord/src/controller/balancer.rs @@ -46,6 +46,7 @@ use mz_cloud_resources::crd::{ use mz_orchestrator_kubernetes::KubernetesImagePullPolicy; use mz_ore::{cli::KeyValueArg, instrument}; +#[derive(Clone)] pub struct Config { pub enable_security_context: bool, pub enable_prometheus_scrape_annotations: bool, diff --git a/src/orchestratord/src/controller/console.rs b/src/orchestratord/src/controller/console.rs index 0fd88209e495e..173bef895c2e1 100644 --- a/src/orchestratord/src/controller/console.rs +++ b/src/orchestratord/src/controller/console.rs @@ -51,6 +51,7 @@ use mz_orchestrator_kubernetes::KubernetesImagePullPolicy; use mz_ore::{cli::KeyValueArg, instrument}; use mz_server_core::listeners::AuthenticatorKind; +#[derive(Clone)] pub struct Config { pub enable_security_context: bool, pub enable_prometheus_scrape_annotations: bool, diff --git a/src/orchestratord/src/controller/materialize.rs b/src/orchestratord/src/controller/materialize.rs index b257c1599bcc7..823ba5ba0282c 100644 --- a/src/orchestratord/src/controller/materialize.rs +++ b/src/orchestratord/src/controller/materialize.rs @@ -54,6 +54,7 @@ use mz_ore::{cast::CastFrom, cli::KeyValueArg, instrument}; pub mod generation; pub mod global; +#[derive(Clone)] pub struct Config { pub cloud_provider: CloudProvider, pub region: String, diff --git a/src/orchestratord/src/gcp_node_upgrade.rs b/src/orchestratord/src/gcp_node_upgrade.rs index 531657d78f64e..d4c0ba8319d37 100644 --- a/src/orchestratord/src/gcp_node_upgrade.rs +++ b/src/orchestratord/src/gcp_node_upgrade.rs @@ -50,6 +50,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use anyhow::{Context as _, bail}; +use futures::future; use k8s_openapi::api::core::v1::{Node, Pod}; use kube::{ Api, Client, @@ -213,6 +214,10 @@ impl ArmedPools { } /// Runs the GCP node upgrade watcher forever. Errors are logged and retried. +/// +/// Never returns. Dropping this future stops all of its work: it spawns no +/// tasks, so it is safe to run under a leadership lease, which it must be, +/// since it triggers rollouts and only one replica may do that. pub async fn run(client: Client, config: Config) { info!( subscription = config.notification_subscription, @@ -226,13 +231,26 @@ pub async fn run(client: Client, config: Config) { let gcp = Arc::new(GcpApiClient::new().await); - mz_ore::task::spawn(|| "gcp node upgrade notification subscriber", { - let config = config.clone(); - let armed = Arc::clone(&armed); - let gcp = Arc::clone(&gcp); - async move { subscriber_loop(gcp, config, armed).await } - }); + // The subscriber is polled inline rather than spawned, so that dropping + // this future stops it too. A subscriber that outlived it would keep + // pulling notifications, and acking them, out from under whichever + // replica holds the lease next. + future::join( + subscriber_loop(Arc::clone(&gcp), config.clone(), Arc::clone(&armed)), + scan_loop(client, gcp, config, armed), + ) + .await; +} +/// Polls the armed node pools, and the GKE API for pools that should be +/// armed, triggering rollouts for the instances on pools whose blue nodes +/// have all been cordoned. Never returns. +async fn scan_loop( + client: Client, + gcp: Arc, + config: Config, + armed: Arc>, +) { let mut scan_interval = tokio::time::interval(config.scan_interval); scan_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let mut last_gke_poll: Option = None; diff --git a/src/orchestratord/src/metrics.rs b/src/orchestratord/src/metrics.rs index 2143d7720104e..3125d7f7f41a1 100644 --- a/src/orchestratord/src/metrics.rs +++ b/src/orchestratord/src/metrics.rs @@ -21,19 +21,43 @@ use mz_ore::metrics::{MetricsRegistry, UIntGauge}; #[derive(Debug)] pub struct Metrics { + pub is_leader: UIntGauge, pub environmentd_needs_update: UIntGauge, } impl Metrics { pub fn register_into(registry: &MetricsRegistry) -> Self { Self { + is_leader: registry.register( + metric! { + name: "orchestratord_is_leader", + help: "Whether this operator replica holds the controller leadership lease, and is therefore the replica reconciling. Summed across the replicas this should be 1. A sustained 0 means no replica can take the lease, for instance because the service account lacks permission on leases, and the operator is reconciling nothing.", + }), environmentd_needs_update: registry.register( metric! { name: "environmentd_needs_update", - help: "Count of organizations in this cluster which are running outdated pod templates", + help: "Count of organizations in this cluster which are running outdated pod templates. Only the operator replica holding the leadership lease reconciles, so the others report zero.", }), } } + + /// Records that this replica has taken the leadership lease. + pub fn leadership_acquired(&self) { + self.is_leader.set(1); + } + + /// Records that this replica no longer holds the leadership lease, and + /// resets the metrics that only mean anything while it is reconciling. + /// + /// Those are derived from what reconciliation observed, and the process + /// outlives its own leadership: it keeps serving the conversion webhook + /// after losing the lease. Without this, a former leader would go on + /// publishing its last observation forever, so summing across the + /// replicas would count the same organizations once per past leader. + pub fn leadership_lost(&self) { + self.is_leader.set(0); + self.environmentd_needs_update.set(0); + } } pub fn router(registry: MetricsRegistry) -> Router { diff --git a/test/orchestratord/mzcompose.py b/test/orchestratord/mzcompose.py index dd2d9c1a952c0..a5ed9a4f0435e 100644 --- a/test/orchestratord/mzcompose.py +++ b/test/orchestratord/mzcompose.py @@ -41,8 +41,8 @@ Service, WorkflowArgumentParser, ) +from materialize.mzcompose.service import Service as ServiceDefinition from materialize.mzcompose.services.balancerd import Balancerd -from materialize.mzcompose.services.clusterd import Clusterd from materialize.mzcompose.services.environmentd import Environmentd from materialize.mzcompose.services.mz import Mz from materialize.mzcompose.services.mz_debug import MzDebug @@ -65,7 +65,15 @@ Testdrive(), Orchestratord(), Environmentd(), - Clusterd(), + # orchestratord deploys clusterd and console pods with the standalone + # materialize/clusterd and materialize/console images (at the same tag + # as environmentd), so those exact images must be built and loaded into + # kind. The Clusterd compose service would resolve to the `materialized` + # fat image instead, leaving materialize/clusterd: nowhere to be + # found when the tag was never published (e.g. when running against a + # local commit). + ServiceDefinition(name="clusterd", config={"mzbuild": "clusterd"}), + ServiceDefinition(name="console", config={"mzbuild": "console"}), Balancerd(), MzDebug(), Mz(app_password=""), @@ -161,11 +169,79 @@ def get_pod_data( ) +def get_current_operator_pod_template_hash() -> str | None: + """The `pod-template-hash` of the operator deployment's current revision, + or None if it can't be determined because the operator isn't installed. + + A Deployment records the revision it is currently rolling out in an + annotation, and each ReplicaSet it owns records the revision it was created + for, so the ReplicaSet whose revision matches is the current one. The + deployment is looked up by label rather than by name so that this keeps + working across the chart versions the upgrade tests install. + """ + + def get_operator_resources(kind: str) -> list[dict[str, Any]]: + return json.loads( + spawn.capture( + [ + "kubectl", + "get", + kind, + "-l", + "app.kubernetes.io/instance=operator", + "-n", + "materialize", + "-o", + "json", + ] + ) + )["items"] + + def revision(resource: dict[str, Any]) -> str | None: + annotations = resource["metadata"].get("annotations") or {} + return annotations.get("deployment.kubernetes.io/revision") + + deployments = get_operator_resources("deployment") + if len(deployments) != 1: + return None + current_revision = revision(deployments[0]) + if current_revision is None: + return None + + for replicaset in get_operator_resources("replicaset"): + if revision(replicaset) == current_revision: + return replicaset["metadata"]["labels"].get("pod-template-hash") + return None + + def get_orchestratord_data() -> dict[str, Any]: - return get_pod_data( + data = get_pod_data( labels={"app.kubernetes.io/instance": "operator"}, namespace="materialize", ) + # A rollout of the operator leaves pods of the previous generation around, + # both terminating ones and (while the replacements come up) ones that are + # still running. Callers only care about the pods that are here to stay, so + # keep only the current generation. Falling back to dropping just the + # terminating pods keeps this usable before the deployment exists. + pod_template_hash = get_current_operator_pod_template_hash() + data["items"] = [ + pod + for pod in data["items"] + if pod["metadata"].get("deletionTimestamp") is None + and ( + pod_template_hash is None + or pod["metadata"]["labels"].get("pod-template-hash") == pod_template_hash + ) + ] + # Callers index into the pods, and the filtering above can transiently leave + # nothing behind: a new ReplicaSet carries the deployment's revision before + # its pods exist. Fail with something readable rather than an IndexError. + assert data["items"], ( + "found no operator pods of the current deployment revision " + f"(pod-template-hash {pod_template_hash})" + ) + return data def get_balancerd_data() -> dict[str, Any]: @@ -3445,16 +3521,21 @@ def image_tag(image: str): def check_orchestratord_version(version: MzVersion): def check(): data = get_orchestratord_data() - assert len(data["items"]) == 1, f"got {len(data['items'])} items" - - got_image = data["items"][0]["spec"]["containers"][0]["image"] + # The number of operator replicas differs between chart versions, + # so we require that all of the pods that exist run the expected + # image rather than pinning an exact pod count. Pods of the + # previous version linger while they terminate, so this also waits + # for the rollout to finish. get_orchestratord_data() already + # rejects an empty set of pods, so this can't pass vacuously. expected_image = get_image( c.compose["services"]["orchestratord"]["image"], str(version), ) - assert image_tag(got_image) == image_tag( - expected_image - ), f"{got_image} != {expected_image}" + for pod in data["items"]: + got_image = pod["spec"]["containers"][0]["image"] + assert image_tag(got_image) == image_tag( + expected_image + ), f"{got_image} != {expected_image}" retry(check, 60) @@ -3973,6 +4054,375 @@ def check_single_generation() -> None: retry(check_single_generation, 120) +OPERATOR_DEPLOYMENT = "operator-materialize-operator" +# The lease name is fixed in the orchestratord binary. A single lease guards +# all of its controllers, in the namespace orchestratord runs in. +LEADER_ELECTION_LEASE = "orchestratord" + + +def get_operator_pod_names() -> list[str]: + """Names of the ready orchestratord pods of the current generation. + + Readiness rather than the `Running` phase, because only ready pods are in + the webhook service's endpoints, and a pod whose readiness probe fails + stays in the `Running` phase. + """ + return [ + pod["metadata"]["name"] + for pod in get_orchestratord_data()["items"] + if any( + condition["type"] == "Ready" and condition["status"] == "True" + for condition in pod["status"].get("conditions") or [] + ) + ] + + +def get_operator_pod_restart_counts() -> dict[str, int]: + """Container restart count per orchestratord pod of the current generation.""" + return { + pod["metadata"]["name"]: sum( + container["restartCount"] + for container in pod["status"].get("containerStatuses") or [] + ) + for pod in get_orchestratord_data()["items"] + } + + +def get_lease_holder(lease_name: str) -> str | None: + """The holderIdentity of a lease, or None if the lease doesn't exist or + has no holder. orchestratord uses its pod name as the identity.""" + try: + holder = spawn.capture( + [ + "kubectl", + "get", + "lease", + lease_name, + "-n", + "materialize", + "-o", + "jsonpath={.spec.holderIdentity}", + ], + stderr=subprocess.DEVNULL, + ).strip() + except subprocess.CalledProcessError: + return None + return holder or None + + +def get_lease_transitions(lease_name: str) -> int: + """How many times a lease has changed hands, or 0 if the lease doesn't + exist yet. + + Every takeover increments this and renewals by the incumbent leave it + alone, so requiring it to grow across an event proves that leadership + really moved, rather than that the incumbent simply held on. + """ + try: + transitions = spawn.capture( + [ + "kubectl", + "get", + "lease", + lease_name, + "-n", + "materialize", + "-o", + "jsonpath={.spec.leaseTransitions}", + ], + stderr=subprocess.DEVNULL, + ).strip() + except subprocess.CalledProcessError: + return 0 + return int(transitions) if transitions else 0 + + +def get_operator_is_leader_metric(pod: str) -> float: + """The `orchestratord_is_leader` gauge scraped from an operator pod. + + Scraped through the API server's pod proxy rather than a port-forward, + since the operator image has no shell or HTTP client to exec into. + """ + metrics = spawn.capture( + [ + "kubectl", + "get", + "--raw", + f"/api/v1/namespaces/materialize/pods/{pod}:3100/proxy/metrics", + ] + ) + for line in metrics.splitlines(): + if line.startswith("orchestratord_is_leader "): + return float(line.split(" ", 1)[1]) + raise AssertionError(f"operator pod {pod} does not report orchestratord_is_leader") + + +def check_is_leader_metric_agrees_with_lease() -> None: + """Assert that exactly the lease holder reports itself as the leader. + + The metric is what alerting keys off, so it has to track the lease rather + than merely be present. A total of zero across the replicas means nobody + is reconciling, which is otherwise only visible in the operator's logs. + """ + holder = get_lease_holder(LEADER_ELECTION_LEASE) + reported = { + pod: get_operator_is_leader_metric(pod) for pod in get_operator_pod_names() + } + assert holder is not None and reported.get(holder) == 1, ( + f"lease {LEADER_ELECTION_LEASE} is held by {holder!r}, " + f"but the replicas report {reported}" + ) + assert ( + sum(reported.values()) == 1 + ), f"expected exactly one operator replica to report itself leader, got {reported}" + + +def check_lease_held_by_operator_pod() -> None: + """Assert that the controller lease is held by a running operator pod.""" + pods = get_operator_pod_names() + holder = get_lease_holder(LEADER_ELECTION_LEASE) + assert ( + holder in pods + ), f"lease {LEADER_ELECTION_LEASE} is held by {holder!r}, expected one of {pods}" + + +def request_rollout_via_v1(definition: dict[str, Any]) -> None: + """Change the v1 spec so the next apply triggers a rollout. + + The v1 CRD has no requestRollout field (setting it would be silently + pruned by the API server). Instead the conversion webhook derives + requestRollout from a hash of the v1 spec, so changing any hashed field + (environmentdExtraEnv here) requests a new rollout. + """ + definition["materialize"]["spec"]["environmentdExtraEnv"] = [ + {"name": "LEADER_FAILOVER_TEST_NONCE", "value": str(uuid.uuid4())} + ] + + +def workflow_leader_failover( + c: Composition, + parser: WorkflowArgumentParser, +) -> None: + """Test that orchestratord leader election fails over between replicas. + + orchestratord runs with multiple replicas (the chart default), and its + controllers (materialize, balancer, console) are all guarded by a single + coordination.k8s.io Lease so that only one replica reconciles at a time, + while the other replicas wait to take the lease over. The conversion + webhook is served by every replica (that's the motivation for running + more than one), so the v1 CRD is enabled and the webhook must stay + functional across both failover paths: + + 1. Leader pod deletion: the lease stops being renewed, so a standby + replica must take it over after the lease expires and then actually + reconcile (proven by driving a rollout to completion), while the + surviving replica keeps serving the conversion webhook. + + 2. Rolling restart of the operator deployment (the production rollout + scenario that motivates running multiple replicas): after the + rollout the lease must be held by a new pod, reconciliation must + still work, and the webhook must still work. + + 3. Losing the lease to another candidate: the replica that lost it must + rejoin the election in process rather than restarting, so that it + keeps serving the conversion webhook throughout. + """ + parser.add_argument( + "--recreate-cluster", + action=argparse.BooleanOptionalAction, + help="Recreate cluster if it exists already", + ) + parser.add_argument( + "--tag", + type=str, + help="Custom version tag to use", + ) + parser.add_argument( + "--orchestratord-override", + default=True, + action=argparse.BooleanOptionalAction, + help="Override orchestratord tag", + ) + args = parser.parse_args() + + definition = setup(c, args) + replicas = definition["operator"]["operator"]["replicas"] + assert ( + replicas >= 2 + ), f"leader failover requires at least 2 operator replicas, chart default is {replicas}" + + # Serve the v1 CRD and apply the Materialize resource at v1, so that + # writes go through the conversion webhook and webhook availability + # across failovers is actually exercised. + enable_v1_crd(definition) + definition["materialize"]["apiVersion"] = "materialize.cloud/v1" + + init(definition) + run(definition, expect_fail=False) + + def webhook_works() -> None: + assert conversion_webhook_works(), "conversion webhook is not working" + + retry(webhook_works, 120) + print("Conversion webhook works") + + # The balancer and console controllers must have reconciled their pods + # into a running state, proving all three controllers work under the + # lease. run() already proved this for the materialize controller. + def balancer_and_console_running() -> None: + for name, data in ( + ("balancerd", get_balancerd_data()), + ("console", get_console_data()), + ): + pods = data["items"] + assert pods and all( + pod["status"]["phase"] == "Running" for pod in pods + ), f"expected all {name} pods to be running, got {[(pod['metadata']['name'], pod['status']['phase']) for pod in pods]}" + + retry(balancer_and_console_running, 300) + print("Balancerd and console pods are running") + + # The chart must create a PodDisruptionBudget so that node drains can't + # take down all operator replicas at once. + spawn.runv(["kubectl", "get", "pdb", OPERATOR_DEPLOYMENT, "-n", "materialize"]) + + def all_replicas_ready() -> None: + pods = get_operator_pod_names() + assert ( + len(pods) == replicas + ), f"expected {replicas} ready operator pods, got {pods}" + + retry(all_replicas_ready, 120) + retry(check_lease_held_by_operator_pod, 120) + # Retried because the leader sets the gauge when it starts running the + # controllers, which is just after it writes the lease. + retry(check_is_leader_metric_agrees_with_lease, 60) + print( + "The controller lease is held by an operator pod, which reports itself leader" + ) + + # --- Case 1: leader pod deletion ----------------------------------------- + old_leader = get_lease_holder(LEADER_ELECTION_LEASE) + assert old_leader is not None + print(f"Deleting leader pod {old_leader}") + spawn.runv( + ["kubectl", "delete", "pod", old_leader, "-n", "materialize", "--wait=false"] + ) + + # The dead leader stops renewing, so after the lease duration (15s by + # default) a standby replica (or the replacement pod) must take over. + def new_leader_elected() -> None: + holder = get_lease_holder(LEADER_ELECTION_LEASE) + pods = get_operator_pod_names() + assert ( + holder != old_leader and holder in pods + ), f"lease {LEADER_ELECTION_LEASE} is held by {holder!r}, expected a pod other than {old_leader} among {pods}" + + retry(new_leader_elected, 120) + print(f"New leader elected: {get_lease_holder(LEADER_ELECTION_LEASE)}") + + # The surviving replica must keep serving the conversion webhook while + # the deleted pod's replacement comes up. + retry(webhook_works, 60) + + # Prove the new leader actually reconciles: request a fresh rollout and + # wait for it to complete. + request_rollout_via_v1(definition) + run(definition, expect_fail=False) + print("Reconciliation works after leader pod deletion") + + # --- Case 2: rolling restart of the operator deployment ------------------ + spawn.runv( + [ + "kubectl", + "rollout", + "restart", + f"deployment/{OPERATOR_DEPLOYMENT}", + "-n", + "materialize", + ] + ) + spawn.runv( + [ + "kubectl", + "rollout", + "status", + f"deployment/{OPERATOR_DEPLOYMENT}", + "-n", + "materialize", + "--timeout=300s", + ] + ) + + # All pods were replaced, so the lease must be taken over by a new pod + # once the old holder's lease expires, and the webhook must be served by + # the new pods. + retry(check_lease_held_by_operator_pod, 120) + retry(webhook_works, 60) + request_rollout_via_v1(definition) + run(definition, expect_fail=False) + print("Reconciliation works after a rolling restart of the operator") + + # --- Case 3: losing the lease to another candidate ----------------------- + # Hand the lease to an identity that belongs to nobody. The leader notices + # on its next renewal that it is no longer the holder and gives leadership + # up, and since the fake holder never renews, the lease is taken over + # again once it expires. + restarts_before = get_operator_pod_restart_counts() + transitions_before = get_lease_transitions(LEADER_ELECTION_LEASE) + spawn.runv( + [ + "kubectl", + "patch", + "lease", + LEADER_ELECTION_LEASE, + "-n", + "materialize", + "--type=merge", + "-p", + json.dumps({"spec": {"holderIdentity": "not-a-real-candidate"}}), + ] + ) + + # An operator pod holding the lease is not enough here: unlike the cases + # above, the incumbent leader is still running and eligible, so it could + # satisfy that without ever having noticed the patch. Requiring a new + # transition proves leadership was actually given up and taken back. + def lease_taken_over_again() -> None: + check_lease_held_by_operator_pod() + transitions = get_lease_transitions(LEADER_ELECTION_LEASE) + assert transitions > transitions_before, ( + f"lease {LEADER_ELECTION_LEASE} is at {transitions} transitions, " + f"expected more than {transitions_before}" + ) + + retry(lease_taken_over_again, 120) + retry(webhook_works, 60) + # This is the case that exercises clearing the gauge: unlike the failovers + # above, the replica that gave leadership up is still running, so it has to + # stop reporting itself leader rather than being replaced by a fresh pod. + retry(check_is_leader_metric_agrees_with_lease, 60) + request_rollout_via_v1(definition) + run(definition, expect_fail=False) + + # Losing the lease must not restart the process: the replica rejoins the + # election in place, which is what keeps it serving the conversion webhook + # the whole time. + restarts_after = get_operator_pod_restart_counts() + for pod, restarts in restarts_before.items(): + # A pod that vanished entirely is a failure too, so require that every + # pod is still there rather than defaulting a missing one to a pass. + assert ( + pod in restarts_after + ), f"operator pod {pod} disappeared after losing the lease" + assert restarts_after[pod] == restarts, ( + f"operator pod {pod} restarted after losing the lease " + f"({restarts} -> {restarts_after[pod]})" + ) + print("Reconciliation works after losing the lease, without a restart") + print("leader failover test PASSED") + + def workflow_default(c: Composition, parser: WorkflowArgumentParser) -> None: parser.add_argument( "--recreate-cluster", @@ -4135,6 +4585,7 @@ def setup(c: Composition, args) -> dict[str, Any]: "orchestratord", "environmentd", "clusterd", + "console", "balancerd", ] c.pull_images(*services) From aa101ecf7ec909e0dfb2edd6ee566360104efe2f Mon Sep 17 00:00:00 2001 From: Alex Hunt Date: Thu, 30 Jul 2026 10:26:25 +0100 Subject: [PATCH 2/2] orchestratord: run the leased components as tasks The controllers and the node upgrade watcher are joined rather than spawned, so that losing the leadership lease drops them. Joining raw futures puts all of them, and the lease renewal that `with_lease` polls alongside them, on a single task, where one of them blocking its poll stalls the rest. A stalled renewal is the worse case, since it can miss the renew deadline and cost this replica its leadership. Spawn each of them as a task instead, and join the abort-on-drop handles, which keeps them independently scheduled while still stopping all of them when the lease is lost. Co-Authored-By: Claude Opus 5 --- src/orchestratord/src/bin/orchestratord.rs | 51 ++++++++++++++++------ src/orchestratord/src/gcp_node_upgrade.rs | 29 ++++++++---- 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/src/orchestratord/src/bin/orchestratord.rs b/src/orchestratord/src/bin/orchestratord.rs index ffd089437c073..1016782c85b1a 100644 --- a/src/orchestratord/src/bin/orchestratord.rs +++ b/src/orchestratord/src/bin/orchestratord.rs @@ -860,24 +860,46 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { // `with_lease` polls this future only once it holds the lease, so // building the controllers here, rather than passing them in, ties // both them and the metric to the moment leadership is acquired. + // + // Each of them runs as its own task, so that they are scheduled + // independently and none can hold up the poll of the others, or of + // the lease renewal that `with_lease` runs alongside them. Their + // handles abort the tasks when dropped, so that none of them + // outlives the lease. let controllers = Box::pin(leader_election.with_lease(async { metrics.leadership_acquired(); join4( - make_materialize_controller().run(), - make_balancer_controller().run(), - make_console_controller().run(), - make_gcp_node_upgrade_watcher(), + mz_ore::task::spawn( + || "materialize controller", + make_materialize_controller().run(), + ) + .abort_on_drop(), + mz_ore::task::spawn(|| "balancer controller", make_balancer_controller().run()) + .abort_on_drop(), + mz_ore::task::spawn(|| "console controller", make_console_controller().run()) + .abort_on_drop(), + mz_ore::task::spawn( + || "gcp node upgrade watcher", + make_gcp_node_upgrade_watcher(), + ) + .abort_on_drop(), ) .await })); match select(controllers, pin!(sigterm.recv())).await { Either::Left((None, _)) => { - // Losing the lease dropped the futures above, which cancels - // the in-flight reconciliations and the node upgrade watch. - // Both are polled inline rather than spawned, so no work - // outlives the lease and we can rejoin the election in - // process instead of restarting. Staying up keeps this - // replica serving the conversion webhook. + // Losing the lease dropped the handles above, which aborts + // the in-flight reconciliations and the node upgrade watch, + // so we can rejoin the election in process instead of + // restarting. Staying up keeps this replica serving the + // conversion webhook. + // + // An aborted task runs until its next await point, so a + // reconciliation can still finish a request it had already + // issued. That is within the margin the lease timings + // leave: the renew deadline is shorter than the lease + // duration, so we give leadership up before another + // candidate is allowed to take the lease over. // // The leader-only gauges describe what this replica saw // while it was reconciling, and this process outlives its @@ -897,10 +919,11 @@ async fn run(args: Args) -> Result<(), anyhow::Error> { std::process::exit(1); } Either::Right((_, controllers)) => { - // Stop reconciling before releasing the lease, so that no - // reconciliation outlives our leadership. Releasing hands - // leadership over immediately rather than making the other - // replicas wait for the lease to expire. + // Abort the controllers before releasing the lease, so that + // reconciliation stops before another replica picks + // leadership up. Releasing hands leadership over + // immediately rather than making the other replicas wait + // for the lease to expire. drop(controllers); // The webhook keeps serving for a while below, so this // replica is still scraped after it stops reconciling. diff --git a/src/orchestratord/src/gcp_node_upgrade.rs b/src/orchestratord/src/gcp_node_upgrade.rs index d4c0ba8319d37..dfcc3c5dbb88b 100644 --- a/src/orchestratord/src/gcp_node_upgrade.rs +++ b/src/orchestratord/src/gcp_node_upgrade.rs @@ -215,9 +215,10 @@ impl ArmedPools { /// Runs the GCP node upgrade watcher forever. Errors are logged and retried. /// -/// Never returns. Dropping this future stops all of its work: it spawns no -/// tasks, so it is safe to run under a leadership lease, which it must be, -/// since it triggers rollouts and only one replica may do that. +/// Never returns. Dropping this future aborts all of its work, which it must, +/// since it triggers rollouts and only the replica holding the leadership +/// lease may do that. The abort is not synchronous with the drop: a task +/// caught mid-poll runs until its next await point. pub async fn run(client: Client, config: Config) { info!( subscription = config.notification_subscription, @@ -231,13 +232,23 @@ pub async fn run(client: Client, config: Config) { let gcp = Arc::new(GcpApiClient::new().await); - // The subscriber is polled inline rather than spawned, so that dropping - // this future stops it too. A subscriber that outlived it would keep - // pulling notifications, and acking them, out from under whichever - // replica holds the lease next. + // The two loops run as separate tasks, so that they are scheduled + // independently and neither can hold the other's poll up. Their handles + // are held here and abort the tasks when dropped, so neither loop + // outlives this future. A subscriber that did would keep pulling + // notifications, and acking them, out from under whichever replica holds + // the lease next. future::join( - subscriber_loop(Arc::clone(&gcp), config.clone(), Arc::clone(&armed)), - scan_loop(client, gcp, config, armed), + mz_ore::task::spawn( + || "gcp node upgrade notification subscriber", + subscriber_loop(Arc::clone(&gcp), config.clone(), Arc::clone(&armed)), + ) + .abort_on_drop(), + mz_ore::task::spawn( + || "gcp node upgrade scan", + scan_loop(client, gcp, config, armed), + ) + .abort_on_drop(), ) .await; }