Create pdb for clusters (#25)

* Create pdb for clusters

Fixes #24

* Ensure deletions get reconciled

Author: jackkleeman

charts/restate-operator-helm/templates/rbac.yaml

@@ -47,6 +47,7 @@ rules:
       - services
       - configmaps
       - serviceaccounts
+      - poddisruptionbudgets
       - networkpolicies
       - statefulsets
       - persistentvolumeclaims
@@ -66,6 +67,7 @@ rules:
       - ''
       - batch
       - apps
+      - policy
       - networking.k8s.io
       - vpcresources.k8s.aws
       - secrets-store.csi.x-k8s.io

src/controllers/restatecluster/controller.rs

@@ -10,6 +10,7 @@ use k8s_openapi::api::core::v1::{
     ConfigMap, Namespace, PersistentVolumeClaim, Service, ServiceAccount, ServiceSpec,
 };
 use k8s_openapi::api::networking::v1::NetworkPolicy;
+use k8s_openapi::api::policy::v1::PodDisruptionBudget;
 use k8s_openapi::apimachinery::pkg::apis::meta::v1::{APIGroup, ObjectMeta};
 
 use kube::core::object::HasStatus;
@@ -373,6 +374,7 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
     let pvc_api = Api::<PersistentVolumeClaim>::all(client.clone());
     let svc_api = Api::<Service>::all(client.clone());
     let svcacc_api = Api::<ServiceAccount>::all(client.clone());
+    let pdb_api = Api::<PodDisruptionBudget>::all(client.clone());
     let cm_api = Api::<ConfigMap>::all(client.clone());
     let np_api = Api::<NetworkPolicy>::all(client.clone());
     let pia_api = Api::<PodIdentityAssociation>::all(client.clone());
@@ -401,28 +403,39 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
 
     let (ss_store, ss_writer) = reflector::store();
     let ss_reflector = reflector(ss_writer, watcher(ss_api, cfg.clone()))
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         .default_backoff()
         .predicate_filter(changed_predicate.combine(status_predicate_serde));
 
     let np_watcher = metadata_watcher(np_api, cfg.clone())
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         .predicate_filter(changed_predicate);
 
     let ns_watcher = metadata_watcher(ns_api, cfg.clone())
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         .predicate_filter(changed_predicate);
 
     let svcacc_watcher = metadata_watcher(svcacc_api, cfg.clone())
+        .map(|event| ensure_deletion_change(event))
+        .touched_objects()
+        .predicate_filter(changed_predicate);
+
+    let pdb_watcher = metadata_watcher(pdb_api, cfg.clone())
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         .predicate_filter(changed_predicate);
 
     let svc_watcher = watcher(svc_api, cfg.clone())
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         // svc has no generation so we hash the spec to check for changes
         .predicate_filter(changed_predicate.combine(spec_predicate_serde));
 
     let cm_watcher = watcher(cm_api, cfg.clone())
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         // cm has no generation so we hash the data to check for changes
         .predicate_filter(changed_predicate.combine(spec_predicate));
@@ -433,6 +446,7 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
         .owns_stream(cm_watcher)
         .owns_stream(ns_watcher)
         .owns_stream(svcacc_watcher)
+        .owns_stream(pdb_watcher)
         .owns_stream(np_watcher)
         .owns_stream(ss_reflector)
         .watches_stream(
@@ -451,6 +465,7 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
         );
     let controller = if pod_identity_association_installed {
         let pia_watcher = watcher(pia_api, cfg.clone())
+            .map(|event| ensure_deletion_change(event))
             .touched_objects()
             // avoid apply loops that seem to happen with crds
             .predicate_filter(changed_predicate.combine(status_predicate));
@@ -461,6 +476,7 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
             job_api,
             Config::default().labels("app.kubernetes.io/name=restate-pia-canary"),
         )
+        .map(|event| ensure_deletion_change(event))
         .touched_objects()
         .predicate_filter(changed_predicate);
 
@@ -470,6 +486,7 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
     };
     let controller = if security_group_policy_installed {
         let sgp_watcher = metadata_watcher(sgp_api, cfg.clone())
+            .map(|event| ensure_deletion_change(event))
             .touched_objects()
             // avoid apply loops that seem to happen with crds
             .predicate_filter(changed_predicate);
@@ -480,6 +497,7 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
     };
     let controller = if secret_provider_class_installed {
         let spc_watcher = metadata_watcher(spc_api, cfg.clone())
+            .map(|event| ensure_deletion_change(event))
             .touched_objects()
             // avoid apply loops that seem to happen with crds
             .predicate_filter(changed_predicate);
@@ -507,6 +525,21 @@ pub async fn run(client: Client, metrics: Metrics, state: State) {
         .await;
 }
 
+// deletion apparently doesn't lead to any change in metadata otherwise, which means the changed_predicate
+// would drop them.
+fn ensure_deletion_change<K: Resource, E>(
+    mut event: Result<kube::runtime::watcher::Event<K>, E>,
+) -> Result<kube::runtime::watcher::Event<K>, E> {
+    if let Ok(kube::runtime::watcher::Event::Delete(ref mut object)) = event {
+        let meta = object.meta_mut();
+        meta.generation = match meta.generation {
+            Some(val) => Some(val + 1),
+            None => Some(0),
+        }
+    }
+    event
+}
+
 fn changed_predicate<K: Resource>(obj: &K) -> Option<u64> {
     let mut hasher = DefaultHasher::new();
     if let Some(g) = obj.meta().generation {

src/controllers/restatecluster/reconcilers/compute.rs

@@ -11,6 +11,7 @@ use k8s_openapi::api::core::v1::{
     ServiceAccount, ServicePort, ServiceSpec, Toleration, Volume, VolumeMount,
     VolumeResourceRequirements,
 };
+use k8s_openapi::api::policy::v1::{PodDisruptionBudget, PodDisruptionBudgetSpec};
 use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
 use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
 use k8s_openapi::apimachinery::pkg::util::intstr::IntOrString;
@@ -177,6 +178,23 @@ fn restate_cluster_service(base_metadata: &ObjectMeta) -> Service {
     }
 }
 
+fn restate_pod_disruption_budget(base_metadata: &ObjectMeta) -> PodDisruptionBudget {
+    PodDisruptionBudget {
+        metadata: object_meta(base_metadata, "restate"),
+        spec: Some(PodDisruptionBudgetSpec {
+            // 1 is a sane default for clusters of all sizes:
+            // cluster size one it will allow downtime, but this is unavoidable when draining nodes
+            // cluster size of 3 with r=2, it will prevent rollouts leading to unavailability
+            // cluster size of 5 with r=3, it is conservative but not unreasonable
+            max_unavailable: Some(IntOrString::Int(1)),
+            min_available: None,
+            selector: Some(label_selector(base_metadata)),
+            unhealthy_pod_eviction_policy: None,
+        }),
+        status: None,
+    }
+}
+
 fn env(cluster_name: &str, custom: Option<&[EnvVar]>) -> Vec<EnvVar> {
     let defaults = [
         ("RESTATE_LOG_FORMAT", "json"),
@@ -439,6 +457,7 @@ pub async fn reconcile_compute(
     let job_api: Api<Job> = Api::namespaced(ctx.client.clone(), namespace);
     let pod_api: Api<Pod> = Api::namespaced(ctx.client.clone(), namespace);
     let sgp_api: Api<SecurityGroupPolicy> = Api::namespaced(ctx.client.clone(), namespace);
+    let pdb_api: Api<PodDisruptionBudget> = Api::namespaced(ctx.client.clone(), namespace);
 
     apply_service_account(
         namespace,
@@ -556,6 +575,13 @@ pub async fn reconcile_compute(
     let restate_cluster_service = restate_cluster_service(base_metadata);
     apply_service(namespace, &svc_api, restate_cluster_service).await?;
 
+    apply_pod_disruption_budget(
+        namespace,
+        &pdb_api,
+        restate_pod_disruption_budget(base_metadata),
+    )
+    .await?;
+
     resize_statefulset_storage(
         namespace,
         base_metadata,
@@ -1002,3 +1028,17 @@ fn validate_stateful_set_status(
 
     Ok(())
 }
+
+async fn apply_pod_disruption_budget(
+    namespace: &str,
+    pdb_api: &Api<PodDisruptionBudget>,
+    pdb: PodDisruptionBudget,
+) -> Result<PodDisruptionBudget, Error> {
+    let name = pdb.metadata.name.as_ref().unwrap();
+    let params: PatchParams = PatchParams::apply("restate-operator").force();
+    debug!(
+        "Applying Pod Disruption Budget {} in namespace {}",
+        name, namespace
+    );
+    Ok(pdb_api.patch(name, &params, &Patch::Apply(&pdb)).await?)
+}