mz_orchestratord/controller/

balancer.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// 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.

use anyhow::bail;
use k8s_controller::{Outcome, TraceMetadata};
use k8s_openapi::{
    api::{
        apps::v1::{Deployment, DeploymentSpec, DeploymentStrategy, RollingUpdateDeployment},
        core::v1::{
            Affinity, Capabilities, ConfigMapVolumeSource, Container, ContainerPort, HTTPGetAction,
            KeyToPath, PodSecurityContext, PodSpec, PodTemplateSpec, Probe, ResourceRequirements,
            SeccompProfile, SecretVolumeSource, SecurityContext, Service, ServicePort, ServiceSpec,
            Toleration, Volume, VolumeMount,
        },
        policy::v1::PodDisruptionBudget,
    },
    apimachinery::pkg::{
        apis::meta::v1::{Condition, LabelSelector, Time},
        util::intstr::IntOrString,
    },
    jiff::Timestamp,
};
use kube::{
    Api, Client, Resource, ResourceExt,
    api::{DeleteParams, ObjectMeta, Patch, PatchParams, PostParams},
    runtime::{conditions::is_deployment_completed, controller::Action, wait::await_condition},
};
use maplit::btreemap;
use tracing::{trace, warn};

use crate::{
    Error,
    k8s::{
        FIELD_MANAGER, apply_resource, get_resource, pod_disruption_budget, recommended_k8s_labels,
        replace_resource,
    },
    tls::{DefaultCertificateSpecs, create_certificate, issuer_ref_defined},
};
use mz_cloud_resources::crd::{
    ManagedResource,
    balancer::v1alpha1::{Balancer, Routing},
    generated::cert_manager::certificates::{Certificate, CertificatePrivateKeyAlgorithm},
};
use mz_orchestrator_kubernetes::KubernetesImagePullPolicy;
use mz_ore::{cli::KeyValueArg, instrument};

/// The name identifying this controller in its reconciliation metrics and in
/// the reporter of the events it publishes.
pub const CONTROLLER_NAME: &str = "balancer";

#[derive(Clone)]
pub struct Config {
    pub enable_security_context: bool,
    /// Whether to run pods in a Kubernetes user namespace (`hostUsers: false`).
    pub enable_user_namespaces: bool,
    pub enable_prometheus_scrape_annotations: bool,

    pub image_pull_policy: KubernetesImagePullPolicy,
    pub scheduler_name: Option<String>,
    pub balancerd_node_selector: Vec<KeyValueArg<String, String>>,
    pub balancerd_affinity: Option<Affinity>,
    pub balancerd_tolerations: Option<Vec<Toleration>>,
    pub balancerd_default_resources: Option<ResourceRequirements>,

    pub default_certificate_specs: DefaultCertificateSpecs,

    pub environmentd_sql_port: u16,
    pub environmentd_http_port: u16,
    pub balancerd_sql_port: u16,
    pub balancerd_http_port: u16,
    pub balancerd_internal_http_port: u16,
    /// `/api/readyz` takes a balancerd at its connection high watermark out of
    /// the Service endpoints. `/api/livez` keeps it in, for load balancers that
    /// read `/api/readyz` in their own health check instead.
    pub balancerd_readiness_probe_path: String,
    pub balancerd_termination_grace_period_seconds: i64,
}

pub struct Context {
    config: Config,
}

impl Context {
    pub fn new(config: Config) -> Self {
        Self { config }
    }

    async fn sync_deployment_status(
        &self,
        client: &Client,
        balancer: &Balancer,
    ) -> Result<(), Error> {
        let namespace = balancer.namespace();
        let balancer_api: Api<Balancer> = Api::namespaced(client.clone(), &namespace);
        let deployment_api: Api<Deployment> = Api::namespaced(client.clone(), &namespace);

        let Some(deployment) = get_resource(&deployment_api, &balancer.deployment_name()).await?
        else {
            return Ok(());
        };

        let Some(deployment_conditions) = &deployment
            .status
            .as_ref()
            .and_then(|status| status.conditions.as_ref())
        else {
            // if the deployment doesn't have any conditions set yet, there
            // is nothing to sync
            return Ok(());
        };

        let ready = deployment_conditions
            .iter()
            .any(|condition| condition.type_ == "Available" && condition.status == "True");
        let ready_str = if ready { "True" } else { "False" };

        let mut status = balancer.status.clone().unwrap();
        if status
            .conditions
            .iter()
            .any(|condition| condition.type_ == "Ready" && condition.status == ready_str)
        {
            // if the deployment status is already set correctly, we don't
            // need to set it again (this prevents us from getting stuck in
            // a reconcile loop)
            return Ok(());
        }

        status.conditions = vec![Condition {
            type_: "Ready".to_string(),
            status: ready_str.to_string(),
            last_transition_time: Time(Timestamp::now()),
            message: format!(
                "balancerd deployment is{} ready",
                if ready { "" } else { " not" }
            ),
            observed_generation: None,
            reason: "DeploymentStatus".to_string(),
        }];
        let mut new_balancer = balancer.clone();
        new_balancer.status = Some(status);

        balancer_api
            .replace_status(
                &balancer.name_unchecked(),
                &PostParams::default(),
                &new_balancer,
            )
            .await?;

        Ok(())
    }

    fn create_external_certificate_object(
        &self,
        balancer: &Balancer,
    ) -> Result<Option<Certificate>, anyhow::Error> {
        create_certificate(
            self.config
                .default_certificate_specs
                .balancerd_external
                .clone(),
            balancer,
            balancer.spec.external_certificate_spec.clone(),
            balancer.external_certificate_name(),
            balancer.external_certificate_secret_name(),
            None,
            CertificatePrivateKeyAlgorithm::Ecdsa,
            Some(256),
        )
    }

    fn create_deployment_object(&self, balancer: &Balancer) -> anyhow::Result<Deployment> {
        let security_context = if self.config.enable_security_context {
            // Since we want to adhere to the most restrictive security context, all
            // of these fields have to be set how they are.
            // See https://kubernetes.io/docs/concepts/security/pod-security-standards/#restricted
            Some(SecurityContext {
                run_as_non_root: Some(true),
                capabilities: Some(Capabilities {
                    drop: Some(vec!["ALL".to_string()]),
                    ..Default::default()
                }),
                seccomp_profile: Some(SeccompProfile {
                    type_: "RuntimeDefault".to_string(),
                    ..Default::default()
                }),
                allow_privilege_escalation: Some(false),
                ..Default::default()
            })
        } else {
            None
        };

        let pod_template_annotations = if self.config.enable_prometheus_scrape_annotations {
            Some(btreemap! {
                "prometheus.io/scrape".to_owned() => "true".to_string(),
                "prometheus.io/port".to_owned() => self.config.balancerd_internal_http_port.to_string(),
                "prometheus.io/path".to_owned() => "/metrics".to_string(),
                "prometheus.io/scheme".to_owned() => "http".to_string(),
            })
        } else {
            None
        };
        let mut pod_template_labels = balancer.default_labels();
        pod_template_labels.insert(
            "materialize.cloud/name".to_owned(),
            balancer.deployment_name(),
        );
        pod_template_labels.insert("app".to_owned(), "balancerd".to_string());
        pod_template_labels.insert("materialize.cloud/app".to_owned(), balancer.app_name());

        let ports = vec![
            ContainerPort {
                container_port: self.config.balancerd_sql_port.into(),
                name: Some("pgwire".into()),
                protocol: Some("TCP".into()),
                ..Default::default()
            },
            ContainerPort {
                container_port: self.config.balancerd_http_port.into(),
                name: Some("http".into()),
                protocol: Some("TCP".into()),
                ..Default::default()
            },
            ContainerPort {
                container_port: self.config.balancerd_internal_http_port.into(),
                name: Some("internal-http".into()),
                protocol: Some("TCP".into()),
                ..Default::default()
            },
        ];

        let mut args = vec![
            "service".to_string(),
            format!(
                "--pgwire-listen-addr=0.0.0.0:{}",
                self.config.balancerd_sql_port
            ),
            format!(
                "--https-listen-addr=0.0.0.0:{}",
                self.config.balancerd_http_port
            ),
            format!(
                "--internal-http-listen-addr=0.0.0.0:{}",
                self.config.balancerd_internal_http_port
            ),
        ];
        match balancer.routing()? {
            Routing::Static(static_routing_config) => {
                args.extend([
                    format!(
                        "--https-resolver-template={}.{}.svc.cluster.local:{}",
                        static_routing_config.environmentd_service_name,
                        static_routing_config.environmentd_namespace,
                        self.config.environmentd_http_port
                    ),
                    format!(
                        "--static-resolver-addr={}.{}.svc.cluster.local:{}",
                        static_routing_config.environmentd_service_name,
                        static_routing_config.environmentd_namespace,
                        self.config.environmentd_sql_port
                    ),
                ]);
            }
            Routing::Frontegg(_frontegg_routing_config) => {
                bail!("frontegg routing is not yet implemented");
            }
        }

        if issuer_ref_defined(
            &self.config.default_certificate_specs.internal,
            &balancer.spec.internal_certificate_spec,
        ) {
            args.push("--internal-tls".to_owned())
        }

        let mut volumes = Vec::new();
        let mut volume_mounts = Vec::new();
        if let Some(name) = &balancer.spec.configmap_name {
            volumes.push(Volume {
                name: "dynamic-config".to_owned(),
                config_map: Some(ConfigMapVolumeSource {
                    name: name.clone(),
                    // The file must exist at startup or balancerd skips the sync loop.
                    optional: Some(false),
                    items: Some(vec![KeyToPath {
                        key: "config.json".to_owned(),
                        path: "config.json".to_owned(),
                        ..Default::default()
                    }]),
                    ..Default::default()
                }),
                ..Default::default()
            });
            volume_mounts.push(VolumeMount {
                name: "dynamic-config".to_owned(),
                // Mount the directory: subPath mounts do not receive ConfigMap updates.
                mount_path: "/etc/balancerd".to_owned(),
                read_only: Some(true),
                ..Default::default()
            });
            args.extend([
                "--config-sync-file-path=/etc/balancerd/config.json".to_owned(),
                "--config-sync-loop-interval=1s".to_owned(),
            ]);
        }
        if issuer_ref_defined(
            &self.config.default_certificate_specs.balancerd_external,
            &balancer.spec.external_certificate_spec,
        ) {
            volumes.push(Volume {
                name: "external-certificate".to_owned(),
                secret: Some(SecretVolumeSource {
                    default_mode: Some(0o400),
                    secret_name: Some(balancer.external_certificate_secret_name()),
                    items: None,
                    optional: Some(false),
                }),
                ..Default::default()
            });
            volume_mounts.push(VolumeMount {
                name: "external-certificate".to_owned(),
                mount_path: "/etc/external_tls".to_owned(),
                read_only: Some(true),
                ..Default::default()
            });
            args.extend([
                "--tls-mode=require".into(),
                "--tls-cert=/etc/external_tls/tls.crt".into(),
                "--tls-key=/etc/external_tls/tls.key".into(),
            ]);
        } else {
            args.push("--tls-mode=disable".to_string());
        }

        let startup_probe = Probe {
            http_get: Some(HTTPGetAction {
                port: IntOrString::Int(self.config.balancerd_internal_http_port.into()),
                path: Some("/api/readyz".into()),
                ..Default::default()
            }),
            failure_threshold: Some(20),
            initial_delay_seconds: Some(3),
            period_seconds: Some(3),
            success_threshold: Some(1),
            timeout_seconds: Some(1),
            ..Default::default()
        };
        // Tighter than liveness so that, on `/api/readyz`, a balancerd at its
        // connection high watermark leaves the Service endpoints within seconds.
        let readiness_probe = Probe {
            http_get: Some(HTTPGetAction {
                port: IntOrString::Int(self.config.balancerd_internal_http_port.into()),
                path: Some(self.config.balancerd_readiness_probe_path.clone()),
                ..Default::default()
            }),
            failure_threshold: Some(2),
            period_seconds: Some(2),
            success_threshold: Some(1),
            timeout_seconds: Some(1),
            ..Default::default()
        };
        let liveness_probe = Probe {
            http_get: Some(HTTPGetAction {
                port: IntOrString::Int(self.config.balancerd_internal_http_port.into()),
                path: Some("/api/livez".into()),
                ..Default::default()
            }),
            failure_threshold: Some(3),
            initial_delay_seconds: Some(8),
            period_seconds: Some(10),
            success_threshold: Some(1),
            timeout_seconds: Some(1),
            ..Default::default()
        };

        let container = Container {
            name: "balancerd".to_owned(),
            image: Some(balancer.spec.balancerd_image_ref.clone()),
            image_pull_policy: Some(self.config.image_pull_policy.to_string()),
            ports: Some(ports),
            args: Some(args),
            startup_probe: Some(startup_probe),
            readiness_probe: Some(readiness_probe),
            liveness_probe: Some(liveness_probe),
            resources: balancer
                .spec
                .resource_requirements
                .clone()
                .or_else(|| self.config.balancerd_default_resources.clone()),
            security_context: security_context.clone(),
            volume_mounts: Some(volume_mounts),
            ..Default::default()
        };

        let match_labels = pod_template_labels.clone();
        pod_template_labels.extend(recommended_k8s_labels(balancer.app_name()));

        let deployment_spec = DeploymentSpec {
            replicas: balancer.replicas(),
            selector: LabelSelector {
                match_labels: Some(match_labels),
                ..Default::default()
            },
            strategy: Some(DeploymentStrategy {
                rolling_update: Some(RollingUpdateDeployment {
                    // Allow a complete set of new pods at once, to minimize the
                    // chances of a new connection going to a pod that will be
                    // immediately drained
                    max_surge: Some(IntOrString::String("100%".into())),
                    ..Default::default()
                }),
                ..Default::default()
            }),
            template: PodTemplateSpec {
                // not using managed_resource_meta because the pod should be
                // owned by the deployment, not the materialize instance
                metadata: Some(ObjectMeta {
                    annotations: pod_template_annotations,
                    labels: Some(pod_template_labels),
                    ..Default::default()
                }),
                spec: Some(PodSpec {
                    containers: vec![container],
                    node_selector: Some(
                        self.config
                            .balancerd_node_selector
                            .iter()
                            .map(|selector| (selector.key.clone(), selector.value.clone()))
                            .collect(),
                    ),
                    affinity: self.config.balancerd_affinity.clone(),
                    tolerations: self.config.balancerd_tolerations.clone(),
                    security_context: Some(PodSecurityContext {
                        fs_group: Some(999),
                        run_as_user: Some(999),
                        run_as_group: Some(999),
                        ..Default::default()
                    }),
                    scheduler_name: self.config.scheduler_name.clone(),
                    termination_grace_period_seconds: Some(
                        self.config.balancerd_termination_grace_period_seconds,
                    ),
                    volumes: Some(volumes),
                    // `None` rather than `Some(true)` when disabled, so the
                    // pod spec is unchanged from before user namespace support.
                    host_users: self.config.enable_user_namespaces.then_some(false),
                    ..Default::default()
                }),
            },
            ..Default::default()
        };

        Ok(Deployment {
            metadata: balancer.managed_resource_meta(balancer.deployment_name()),
            spec: Some(deployment_spec),
            status: None,
        })
    }

    fn create_service_object(&self, balancer: &Balancer) -> Service {
        let selector =
            btreemap! {"materialize.cloud/name".to_string() => balancer.deployment_name()};

        let ports = vec![
            ServicePort {
                name: Some("http".to_string()),
                protocol: Some("TCP".to_string()),
                port: self.config.balancerd_http_port.into(),
                target_port: Some(IntOrString::Int(self.config.balancerd_http_port.into())),
                ..Default::default()
            },
            ServicePort {
                name: Some("pgwire".to_string()),
                protocol: Some("TCP".to_string()),
                port: self.config.balancerd_sql_port.into(),
                target_port: Some(IntOrString::Int(self.config.balancerd_sql_port.into())),
                ..Default::default()
            },
        ];

        let spec = ServiceSpec {
            type_: Some("ClusterIP".to_string()),
            cluster_ip: Some("None".to_string()),
            selector: Some(selector),
            ports: Some(ports),
            ..Default::default()
        };

        Service {
            metadata: balancer.managed_resource_meta(balancer.service_name()),
            spec: Some(spec),
            status: None,
        }
    }

    /// Hands `spec.replicas` to [`REPLICAS_HANDOVER_MANAGER`] at its live value
    /// while [`FIELD_MANAGER`] still owns it, creating the Deployment with
    /// `initial_replicas` first if it does not exist yet.
    ///
    /// NOTE: Dropping a field from an apply when no other manager owns it
    /// deletes it, which defaults the Deployment to 1 replica. An autoscaler
    /// cannot become a co-owner first, because each of its writes triggers a
    /// reconcile whose forced apply reclaims the field. The handover apply is
    /// not forced, so a concurrent scale by the autoscaler conflicts and the
    /// reconcile retries.
    async fn release_replicas(
        &self,
        deployment_api: &Api<Deployment>,
        deployment: &Deployment,
        initial_replicas: i32,
    ) -> Result<(), Error> {
        let name = deployment.name_unchecked();
        let existing = match get_resource(deployment_api, &name).await? {
            Some(existing) => existing,
            None => {
                // The handover apply carries no selector or template, so it
                // cannot create the Deployment. Create it with the initial
                // count under FIELD_MANAGER, then hand that count over below.
                let mut seeded = deployment.clone();
                seeded.spec.as_mut().unwrap().replicas = Some(initial_replicas);
                apply_resource(deployment_api, &seeded).await?
            }
        };
        if !owns_replicas(&existing) {
            return Ok(());
        }
        let replicas = existing.spec.and_then(|spec| spec.replicas);
        deployment_api
            .patch(
                &name,
                &PatchParams::apply(REPLICAS_HANDOVER_MANAGER),
                &Patch::Apply(serde_json::json!({
                    "apiVersion": "apps/v1",
                    "kind": "Deployment",
                    "metadata": { "name": name },
                    "spec": { "replicas": replicas },
                })),
            )
            .await?;
        Ok(())
    }

    // TODO: remove this once everyone is upgraded to an orchestratord
    // version with the separate balancer operator
    async fn fix_deployment(
        &self,
        deployment_api: &Api<Deployment>,
        new_deployment: &Deployment,
    ) -> Result<(), Error> {
        let Some(mut existing_deployment) =
            get_resource(deployment_api, &new_deployment.name_unchecked()).await?
        else {
            return Ok(());
        };

        if existing_deployment.spec.as_ref().unwrap().selector
            == new_deployment.spec.as_ref().unwrap().selector
        {
            return Ok(());
        }

        warn!("found existing deployment with old label selector, fixing");

        // this is sufficient because the new labels are a superset of the
        // old labels, so the existing label selector should still be valid
        existing_deployment
            .spec
            .as_mut()
            .unwrap()
            .template
            .metadata
            .as_mut()
            .unwrap()
            .labels = new_deployment
            .spec
            .as_ref()
            .unwrap()
            .template
            .metadata
            .as_ref()
            .unwrap()
            .labels
            .clone();

        // using await_condition is not ideal in a controller loop, but this
        // is very temporary and will only ever happen once, so this feels
        // simpler than trying to introduce an entire state machine here
        replace_resource(deployment_api, &existing_deployment).await?;
        await_condition(
            deployment_api.clone(),
            &existing_deployment.name_unchecked(),
            |deployment: Option<&Deployment>| {
                let observed_generation = deployment
                    .and_then(|deployment| deployment.status.as_ref())
                    .and_then(|status| status.observed_generation)
                    .unwrap_or(0);
                let current_generation = deployment
                    .and_then(|deployment| deployment.meta().generation)
                    .unwrap_or(0);
                let previous_generation = existing_deployment.meta().generation.unwrap_or(0);
                observed_generation == current_generation
                    && current_generation > previous_generation
            },
        )
        .await
        .map_err(|e| anyhow::anyhow!(e))?;
        await_condition(
            deployment_api.clone(),
            &existing_deployment.name_unchecked(),
            is_deployment_completed(),
        )
        .await
        .map_err(|e| anyhow::anyhow!(e))?;

        // delete the deployment but leave the pods around (via
        // DeleteParams::orphan)
        match kube::runtime::wait::delete::delete_and_finalize(
            deployment_api.clone(),
            &existing_deployment.name_unchecked(),
            &DeleteParams::orphan(),
        )
        .await
        {
            Ok(_) => {}
            Err(kube::runtime::wait::delete::Error::Delete(kube::Error::Api(e)))
                if e.code == 404 =>
            {
                // the resource already doesn't exist
            }
            Err(e) => return Err(anyhow::anyhow!(e).into()),
        }

        // now, the normal apply of the new deployment (in the main loop)
        // will take over the existing pods from the old deployment we just
        // deleted, since we already updated the pod labels to be the same as
        // the new label selector

        Ok(())
    }
}

#[async_trait::async_trait]
impl k8s_controller::Context for Context {
    type Resource = Balancer;
    type Error = Error;

    #[instrument(fields())]
    async fn apply(
        &self,
        client: Client,
        balancer: &Self::Resource,
        metadata: &mut TraceMetadata,
    ) -> Result<Option<Action>, Self::Error> {
        if balancer.status.is_none() {
            let step = metadata.step("initialize_status");
            let balancer_api: Api<Balancer> =
                Api::namespaced(client.clone(), &balancer.meta().namespace.clone().unwrap());
            let mut new_balancer = balancer.clone();
            new_balancer.status = Some(balancer.status());
            balancer_api
                .replace_status(
                    &balancer.name_unchecked(),
                    &PostParams::default(),
                    &new_balancer,
                )
                .await?;
            step.finish(Outcome::Completed);
            // Updating the status should trigger a reconciliation
            // which will include a status this time.
            return Ok(None);
        }

        let namespace = balancer.namespace();
        let certificate_api: Api<Certificate> = Api::namespaced(client.clone(), &namespace);
        let deployment_api: Api<Deployment> = Api::namespaced(client.clone(), &namespace);
        let service_api: Api<Service> = Api::namespaced(client.clone(), &namespace);
        let pdb_api: Api<PodDisruptionBudget> = Api::namespaced(client.clone(), &namespace);

        let step = metadata.step("certificate");
        if let Some(external_certificate) = self.create_external_certificate_object(balancer)? {
            trace!("creating new balancerd external certificate");
            apply_resource(&certificate_api, &external_certificate).await?;
            step.finish(Outcome::Completed);
        } else {
            step.finish(Outcome::Skipped);
        }

        let step = metadata.step("deployment");
        let deployment = self.create_deployment_object(balancer)?;
        self.fix_deployment(&deployment_api, &deployment).await?;
        // Must precede the apply that omits `replicas`; see `release_replicas`.
        if balancer.replicas().is_none() {
            self.release_replicas(&deployment_api, &deployment, balancer.initial_replicas())
                .await?;
        }
        trace!("creating new balancerd deployment");
        apply_resource(&deployment_api, &deployment).await?;
        step.finish(Outcome::Completed);

        let step = metadata.step("pod_disruption_budget");
        let pdb = pod_disruption_budget(
            balancer.managed_resource_meta(balancer.deployment_name()),
            &deployment,
        );
        apply_resource(&pdb_api, &pdb).await?;
        step.finish(Outcome::Completed);

        let step = metadata.step("service");
        let service = self.create_service_object(balancer);
        trace!("creating new balancerd service");
        apply_resource(&service_api, &service).await?;
        step.finish(Outcome::Completed);

        let step = metadata.step("sync_status");
        self.sync_deployment_status(&client, balancer).await?;
        step.finish(Outcome::Completed);

        Ok(None)
    }
}

const REPLICAS_HANDOVER_MANAGER: &str = "orchestratord.materialize.cloud/replicas-handover";

fn owns_replicas(deployment: &Deployment) -> bool {
    deployment
        .metadata
        .managed_fields
        .iter()
        .flatten()
        .filter(|entry| entry.manager.as_deref() == Some(FIELD_MANAGER))
        .filter_map(|entry| entry.fields_v1.as_ref())
        .any(|fields| fields.0.pointer("/f:spec/f:replicas").is_some())
}

#[cfg(test)]
mod tests {
    use super::*;
    use mz_cloud_resources::crd::balancer::v1alpha1::{BalancerSpec, StaticRoutingConfig};

    fn test_context() -> Context {
        Context::new(Config {
            enable_security_context: false,
            enable_user_namespaces: false,
            enable_prometheus_scrape_annotations: false,
            image_pull_policy: KubernetesImagePullPolicy::IfNotPresent,
            scheduler_name: None,
            balancerd_node_selector: vec![],
            balancerd_affinity: None,
            balancerd_tolerations: None,
            balancerd_default_resources: None,
            default_certificate_specs: DefaultCertificateSpecs::default(),
            environmentd_sql_port: 6875,
            environmentd_http_port: 6876,
            balancerd_sql_port: 6875,
            balancerd_http_port: 6876,
            balancerd_internal_http_port: 6878,
            balancerd_readiness_probe_path: "/api/readyz".to_owned(),
            balancerd_termination_grace_period_seconds: 600,
        })
    }

    fn test_balancer(spec: BalancerSpec) -> Balancer {
        let mut balancer = Balancer::new(
            "test",
            BalancerSpec {
                static_routing: Some(StaticRoutingConfig {
                    environmentd_namespace: "test".to_owned(),
                    environmentd_service_name: "environmentd".to_owned(),
                }),
                ..spec
            },
        );
        balancer.metadata.namespace = Some("test".to_owned());
        balancer.metadata.uid = Some("test".to_owned());
        balancer.status = Some(balancer.status());
        balancer
    }

    #[mz_ore::test]
    fn externally_scaled_omits_replicas() {
        let context = test_context();
        for (externally_scaled, expected) in [(None, Some(3)), (Some(true), None)] {
            let balancer = test_balancer(BalancerSpec {
                replicas: Some(3),
                externally_scaled,
                ..Default::default()
            });
            let deployment = context.create_deployment_object(&balancer).unwrap();
            assert_eq!(deployment.spec.unwrap().replicas, expected);
            assert_eq!(balancer.initial_replicas(), 3);
        }
    }

    #[mz_ore::test]
    fn user_namespaces_set_host_users() {
        let balancer = test_balancer(BalancerSpec::default());
        for (enabled, expected) in [(false, None), (true, Some(false))] {
            let mut context = test_context();
            context.config.enable_user_namespaces = enabled;
            let pod = context
                .create_deployment_object(&balancer)
                .unwrap()
                .spec
                .unwrap()
                .template
                .spec
                .unwrap();
            assert_eq!(pod.host_users, expected);
        }
    }

    #[mz_ore::test]
    fn pdb_selects_deployment_pods() {
        let balancer = test_balancer(BalancerSpec::default());
        let deployment = test_context().create_deployment_object(&balancer).unwrap();
        let pdb = pod_disruption_budget(
            balancer.managed_resource_meta(balancer.deployment_name()),
            &deployment,
        );
        let spec = pdb.spec.unwrap();
        assert_eq!(spec.max_unavailable, Some(IntOrString::Int(1)));
        assert_eq!(
            spec.unhealthy_pod_eviction_policy.as_deref(),
            Some("AlwaysAllow")
        );
        let match_labels = spec.selector.unwrap().match_labels.unwrap();
        let pod_labels = deployment
            .spec
            .unwrap()
            .template
            .metadata
            .unwrap()
            .labels
            .unwrap();
        assert!(!match_labels.is_empty());
        assert!(
            match_labels
                .iter()
                .all(|(k, v)| pod_labels.get(k) == Some(v)),
            "PDB selector {match_labels:?} does not match pod labels {pod_labels:?}",
        );
    }

    #[mz_ore::test]
    fn owns_replicas_reads_managed_fields() {
        use k8s_openapi::apimachinery::pkg::apis::meta::v1::{FieldsV1, ManagedFieldsEntry};
        let deployment = |manager: &str| Deployment {
            metadata: ObjectMeta {
                managed_fields: Some(vec![ManagedFieldsEntry {
                    manager: Some(manager.to_owned()),
                    fields_v1: Some(FieldsV1(
                        serde_json::json!({ "f:spec": { "f:replicas": {} } }),
                    )),
                    ..Default::default()
                }]),
                ..Default::default()
            },
            ..Default::default()
        };
        assert!(owns_replicas(&deployment(FIELD_MANAGER)));
        assert!(!owns_replicas(&deployment(REPLICAS_HANDOVER_MANAGER)));
        assert!(!owns_replicas(&Deployment::default()));
    }

    #[mz_ore::test]
    fn configmap_reference_controls_file_sync() {
        let context = test_context();
        for configmap_name in [None, Some("balancerd-settings".to_owned())] {
            let balancer = test_balancer(BalancerSpec {
                configmap_name: configmap_name.clone(),
                ..Default::default()
            });
            let pod = context
                .create_deployment_object(&balancer)
                .unwrap()
                .spec
                .unwrap()
                .template
                .spec
                .unwrap();
            let container = &pod.containers[0];
            let args = container.args.as_ref().unwrap();
            let volumes = pod.volumes.as_ref().unwrap();
            let mounts = container.volume_mounts.as_ref().unwrap();
            if let Some(name) = configmap_name {
                let source = volumes[0].config_map.as_ref().unwrap();
                assert_eq!(source.name, name);
                assert_eq!(source.optional, Some(false));
                assert_eq!(source.items.as_ref().unwrap()[0].key, "config.json");
                assert_eq!(mounts[0].read_only, Some(true));
                assert!(mounts[0].sub_path.is_none());
                assert!(args.contains(&format!(
                    "--config-sync-file-path={}/config.json",
                    mounts[0].mount_path
                )));
                assert!(args.contains(&"--config-sync-loop-interval=1s".to_owned()));
            } else {
                assert!(volumes.is_empty());
                assert!(mounts.is_empty());
                assert!(!args.iter().any(|arg| arg.starts_with("--config-sync-")));
            }
        }
    }
}