orchestratord: Add PodDisruptionBudgets for balancerd and console

misc/helm-charts/operator/templates/clusterrole.yaml

@@ -61,6 +61,17 @@ rules:
   - get
   - list
   - watch
+- apiGroups: ["policy"]
+  resources:
+  - poddisruptionbudgets
+  verbs:
+  - create
+  - update
+  - patch
+  - delete
+  - get
+  - list
+  - watch
 - apiGroups: ["rbac.authorization.k8s.io"]
   resources:
   - roles

src/orchestratord/src/bin/orchestratord.rs

@@ -26,6 +26,7 @@ use k8s_openapi::{
             Affinity, ConfigMap, Probe, ResourceRequirements, Service, ServiceAccount, Toleration,
         },
         networking::v1::NetworkPolicy,
+        policy::v1::PodDisruptionBudget,
         rbac::v1::{Role, RoleBinding},
     },
     apiextensions_apiserver::pkg::apis::apiextensions::v1::{
@@ -872,6 +873,12 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
                         watcher::Config::default()
                             .labels("materialize.cloud/mz-resource-id")
                             .timeout(29),
+                    )
+                    .owns(
+                        Api::<PodDisruptionBudget>::all(client.clone()),
+                        watcher::Config::default()
+                            .labels("materialize.cloud/mz-resource-id")
+                            .timeout(29),
                     );
                 if has_cert_manager {
                     controller.owns(
@@ -929,6 +936,12 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
                             .labels("materialize.cloud/mz-resource-id")
                             .timeout(29),
                     )
+                    .owns(
+                        Api::<PodDisruptionBudget>::all(client.clone()),
+                        watcher::Config::default()
+                            .labels("materialize.cloud/mz-resource-id")
+                            .timeout(29),
+                    )
                     .owns(
                         Api::<NetworkPolicy>::all(client.clone()),
                         watcher::Config::default()

src/orchestratord/src/controller/balancer.rs

@@ -18,6 +18,7 @@ use k8s_openapi::{
             SeccompProfile, SecretVolumeSource, SecurityContext, Service, ServicePort, ServiceSpec,
             Toleration, Volume, VolumeMount,
         },
+        policy::v1::PodDisruptionBudget,
     },
     apimachinery::pkg::{
         apis::meta::v1::{Condition, LabelSelector, Time},
@@ -35,7 +36,10 @@ use tracing::{trace, warn};
 
 use crate::{
     Error,
-    k8s::{FIELD_MANAGER, apply_resource, get_resource, recommended_k8s_labels, replace_resource},
+    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::{
@@ -663,6 +667,7 @@ impl k8s_controller::Context for Context {
         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)? {
@@ -685,6 +690,14 @@ impl k8s_controller::Context for Context {
         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");
@@ -768,6 +781,38 @@ mod tests {
         }
     }
 
+    #[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};

src/orchestratord/src/controller/console.rs

@@ -21,6 +21,7 @@ use k8s_openapi::{
             IPBlock, NetworkPolicy, NetworkPolicyIngressRule, NetworkPolicyPeer, NetworkPolicyPort,
             NetworkPolicySpec,
         },
+        policy::v1::PodDisruptionBudget,
     },
     apimachinery::pkg::{
         apis::meta::v1::{Condition, LabelSelector, Time},
@@ -39,7 +40,10 @@ use tracing::{trace, warn};
 
 use crate::{
     Error,
-    k8s::{apply_resource, get_resource, recommended_k8s_labels, replace_resource},
+    k8s::{
+        apply_resource, get_resource, pod_disruption_budget, recommended_k8s_labels,
+        replace_resource,
+    },
     tls::{DefaultCertificateSpecs, create_certificate, issuer_ref_defined},
 };
 use mz_cloud_resources::crd::{
@@ -646,6 +650,7 @@ impl k8s_controller::Context for Context {
         let deployment_api: Api<Deployment> = Api::namespaced(client.clone(), &namespace);
         let service_api: Api<Service> = Api::namespaced(client.clone(), &namespace);
         let certificate_api: Api<Certificate> = Api::namespaced(client.clone(), &namespace);
+        let pdb_api: Api<PodDisruptionBudget> = Api::namespaced(client.clone(), &namespace);
 
         let step = metadata.step("network_policies");
         trace!("creating new network policies");
@@ -673,6 +678,14 @@ impl k8s_controller::Context for Context {
         apply_resource(&deployment_api, &console_deployment).await?;
         step.finish(Outcome::Completed);
 
+        let step = metadata.step("pod_disruption_budget");
+        let pdb = pod_disruption_budget(
+            console.managed_resource_meta(console.deployment_name()),
+            &console_deployment,
+        );
+        apply_resource(&pdb_api, &pdb).await?;
+        step.finish(Outcome::Completed);
+
         let step = metadata.step("service");
         trace!("creating new console service");
         let console_service = self.create_console_service_object(console);

src/orchestratord/src/k8s.rs

@@ -12,14 +12,19 @@ use std::time::Duration;
 
 use k8s_openapi::{
     ByteString,
+    api::{
+        apps::v1::Deployment,
+        policy::v1::{PodDisruptionBudget, PodDisruptionBudgetSpec},
+    },
     apiextensions_apiserver::pkg::apis::apiextensions::v1::{
         CustomResourceColumnDefinition, CustomResourceConversion, ServiceReference,
         WebhookClientConfig, WebhookConversion,
     },
+    apimachinery::pkg::util::intstr::IntOrString,
 };
 use kube::{
     Api, Client, CustomResourceExt, Resource, ResourceExt,
-    api::{DeleteParams, Patch, PatchParams, PostParams},
+    api::{DeleteParams, ObjectMeta, Patch, PatchParams, PostParams},
     core::Status,
     core::response::StatusSummary,
 };
@@ -100,6 +105,24 @@ where
     }
 }
 
+/// A PodDisruptionBudget letting at most one of `deployment`'s pods be evicted
+/// at a time. With a single replica it never blocks a drain.
+pub fn pod_disruption_budget(metadata: ObjectMeta, deployment: &Deployment) -> PodDisruptionBudget {
+    PodDisruptionBudget {
+        metadata,
+        spec: Some(PodDisruptionBudgetSpec {
+            max_unavailable: Some(IntOrString::Int(1)),
+            selector: deployment.spec.as_ref().map(|spec| spec.selector.clone()),
+            // The default, `IfHealthyBudget`, refuses to evict unready pods
+            // once none are healthy, so a crashlooping Deployment would block
+            // node drains.
+            unhealthy_pod_eviction_policy: Some("AlwaysAllow".into()),
+            ..Default::default()
+        }),
+        status: None,
+    }
+}
+
 /// Configuration for the conversion webhook that serves the v1 version of the
 /// Materialize CRD. When present, the v1 version is registered alongside
 /// v1alpha1 with webhook conversion between them; when absent, only v1alpha1

test/orchestratord/mzcompose.py

@@ -712,6 +712,48 @@ def validate(self, mods: dict[type[Modification], Any]) -> None:
             ), f"Expected {expected} in environmentd args, but only found {args}"
 
 
+def check_pod_disruption_budget(app: str, enabled: bool) -> None:
+    namespace = "materialize-environment"
+    pdbs = json.loads(
+        spawn.capture(
+            [
+                "kubectl",
+                "get",
+                "poddisruptionbudget",
+                "-l",
+                f"app.kubernetes.io/name={app}",
+                "-n",
+                namespace,
+                "-o",
+                "json",
+            ]
+        )
+    )["items"]
+    if not enabled:
+        assert not pdbs, f"Unexpected {app} PDBs: {pdbs}"
+        return
+    assert len(pdbs) == 1, f"Expected one {app} PDB, got {pdbs}"
+    pdb = pdbs[0]
+    name = pdb["metadata"]["name"]
+    assert pdb["spec"]["maxUnavailable"] == 1, pdb["spec"]
+    assert pdb["spec"]["unhealthyPodEvictionPolicy"] == "AlwaysAllow", pdb["spec"]
+    deployment = json.loads(
+        spawn.capture(
+            ["kubectl", "get", "deployment", name, "-n", namespace, "-o", "json"]
+        )
+    )
+    pod_labels = deployment["spec"]["template"]["metadata"]["labels"]
+    match_labels = pdb["spec"]["selector"]["matchLabels"]
+    assert (
+        match_labels and match_labels.items() <= pod_labels.items()
+    ), f"{app} PDB selector {match_labels} does not match pod labels {pod_labels}"
+    # With every pod healthy, one may be evicted whatever the replica count,
+    # so the PDB never blocks a drain.
+    status = pdb.get("status", {})
+    assert status.get("expectedPods", 0) >= 1, status
+    assert status.get("disruptionsAllowed") == 1, status
+
+
 class BalancerdEnabled(Modification):
     @classmethod
     def values(cls, version: MzVersion) -> list[Any]:
@@ -748,6 +790,7 @@ def check() -> None:
                 assert result, f"Unexpected result: {result}"
             else:
                 assert not result, f"Unexpected result: {result}"
+            check_pod_disruption_budget("balancerd", self.value)
 
         # Balancerd can take a while to start up
         retry(check, 240)
@@ -894,7 +937,8 @@ def validate(self, mods: dict[type[Modification], Any]) -> None:
             return
 
         def check() -> None:
-            pass  # TODO: https://linear.app/materializeinc/issue/DB-105
+            check_pod_disruption_budget("console", self.value)
+            # TODO: https://linear.app/materializeinc/issue/DB-105
             # result = spawn.capture(
             #     [
             #         "kubectl",