Support scale down by updating zk node in the verified zk controller (#…

…186) The key is to update the `CLUSTER_SIZE` zk node to make it consistent with the expected replicas number so that the (unverified) teardown script inside the zookeeper image can correctly decide whether to remove the pod from the zookeeper cluster before the pod deletion. To implement this, this PR also introduces some miscellaneous changes including (1) introducing one more field to `ZookeeperReconcileState` to store a statefulset object from the get request, (2) implementing the update branch of the external `reconcile_zk_node` function, (3) adding two more fields to the statefulset spec. --------- Signed-off-by: Wenjie Ma <[email protected]> Signed-off-by: ZichengMa <[email protected]> Co-authored-by: Wenjie Ma <[email protected]>
anvil-verifier · Jul 27, 2023 · 9170c80 · 9170c80
1 parent cc37f5c
commit 9170c80
Show file tree

Hide file tree

Showing 10 changed files with 477 additions and 234 deletions.
diff --git a/src/controller_examples/rabbitmq_controller/proof/safety.rs b/src/controller_examples/rabbitmq_controller/proof/safety.rs
diff --git a/src/controller_examples/zookeeper_controller/common/mod.rs b/src/controller_examples/zookeeper_controller/common/mod.rs
@@ -16,12 +16,14 @@ pub enum ZookeeperReconcileStep {
     AfterCreateStatefulSet,
     AfterUpdateStatefulSet,
     AfterCreateZKNode,
+    AfterUpdateZKNode,
     Done,
     Error,
 }
 
 pub enum Error {
     ClusterSizeZKNodeCreationFailed,
+    ClusterSizeZKNodeUpdateFailed,
 }
 
 }
diff --git a/src/controller_examples/zookeeper_controller/exec/reconciler.rs b/src/controller_examples/zookeeper_controller/exec/reconciler.rs
@@ -24,12 +24,17 @@ pub struct ZookeeperReconcileState {
     // reconcile_step, like a program counter, is used to track the progress of reconcile_core
     // since reconcile_core is frequently "trapped" into the controller_runtime spec.
     pub reconcile_step: ZookeeperReconcileStep,
+    pub sts_from_get: Option<StatefulSet>,
 }
 
 impl ZookeeperReconcileState {
     pub open spec fn to_view(&self) -> zk_spec::ZookeeperReconcileState {
         zk_spec::ZookeeperReconcileState {
             reconcile_step: self.reconcile_step,
+            sts_from_get: match &self.sts_from_get {
+                Some(sts) => Option::Some(sts@),
+                None => Option::None,
+            },
         }
     }
 }
@@ -68,6 +73,7 @@ pub fn reconcile_init_state() -> (state: ZookeeperReconcileState)
 {
     ZookeeperReconcileState {
         reconcile_step: ZookeeperReconcileStep::Init,
+        sts_from_get: Option::None,
     }
 }
 
@@ -174,23 +180,31 @@ pub fn reconcile_core(
                 let stateful_set = make_stateful_set(zk);
                 let get_sts_resp = resp_o.unwrap().into_k_response().into_get_response().res;
                 if get_sts_resp.is_ok() {
-                    // update
                     let found_stateful_set = StatefulSet::from_dynamic_object(get_sts_resp.unwrap());
-                    if found_stateful_set.is_ok() {
-                        let mut new_stateful_set = found_stateful_set.unwrap();
-                        new_stateful_set.set_spec(stateful_set.spec().unwrap());
-                        let req_o = KubeAPIRequest::UpdateRequest(KubeUpdateRequest {
-                            api_resource: StatefulSet::api_resource(),
-                            name: stateful_set.metadata().name().unwrap(),
-                            namespace: zk.metadata().namespace().unwrap(),
-                            obj: new_stateful_set.to_dynamic_object(),
-                        });
+                    // Before update sts itself, firstly update ZKNode
+                    // since zk cluster may be down if majority of the servers are down
+                    // Updating ZKNode firstly can make sure the peer list will change as the sever goes down.
+                    // Details can be found in https://github.com/vmware-research/verifiable-controllers/issues/174
+                    if found_stateful_set.is_ok(){
                         let state_prime = ZookeeperReconcileState {
-                            reconcile_step: ZookeeperReconcileStep::AfterUpdateStatefulSet,
+                            reconcile_step: ZookeeperReconcileStep::AfterUpdateZKNode,
+                            // Save the old sts from get request.
+                            // Then, later when we want to update sts, we can use the old sts as the base.
+                            // We do not need to call GetRequest again.
+                            sts_from_get: Some(found_stateful_set.unwrap()),
                             ..state
                         };
-                        return (state_prime, Option::Some(Request::KRequest(req_o)));
+                        let path = cluster_size_zk_node_path(zk);
+                        let uri = zk_service_uri(zk);
+                        let replicas = i32_to_string(zk.spec().replicas());
+                        let ext_req = ZKSupportInput::ReconcileZKNode(path, uri, replicas);
+                        proof {
+                            zk_support_input_to_view_match(path, uri, replicas);
+                        }
+                        // Call external APIs to update the content in ZKNode
+                        return (state_prime, Option::Some(Request::ExternalRequest(ext_req)));
                     }
+
                 } else if get_sts_resp.unwrap_err().is_object_not_found() {
                     // create
                     let req_o = KubeAPIRequest::CreateRequest(KubeCreateRequest {
@@ -240,6 +254,35 @@ pub fn reconcile_core(
             }
             return (state_prime, Option::Some(Request::ExternalRequest(ext_req)));
         },
+        ZookeeperReconcileStep::AfterUpdateZKNode => {
+            // update sts
+            if resp_o.is_some() && resp_o.as_ref().unwrap().is_external_response()
+            && resp_o.as_ref().unwrap().as_external_response_ref().is_reconcile_zk_node()
+            && state.sts_from_get.is_some() {
+                let found_stateful_set = state.sts_from_get;
+                let mut new_stateful_set = found_stateful_set.unwrap();
+                let stateful_set = make_stateful_set(zk);
+                new_stateful_set.set_spec(stateful_set.spec().unwrap());
+                let req_o = KubeAPIRequest::UpdateRequest(KubeUpdateRequest {
+                    api_resource: StatefulSet::api_resource(),
+                    name: stateful_set.metadata().name().unwrap(),
+                    namespace: zk.metadata().namespace().unwrap(),
+                    obj: new_stateful_set.to_dynamic_object(),
+                });
+                let state_prime = ZookeeperReconcileState {
+                    reconcile_step: ZookeeperReconcileStep::AfterUpdateStatefulSet,
+                    sts_from_get: Option::None
+                };
+                return (state_prime, Option::Some(Request::KRequest(req_o)));
+            } else {
+                // return error state
+                let state_prime = ZookeeperReconcileState {
+                    reconcile_step: ZookeeperReconcileStep::Error,
+                    ..state
+                };
+                return (state_prime, Option::None);
+            }
+        },
         ZookeeperReconcileStep::AfterCreateZKNode => {
             if resp_o.is_some() && resp_o.as_ref().unwrap().is_external_response()
             && resp_o.as_ref().unwrap().as_external_response_ref().is_reconcile_zk_node() {
@@ -546,6 +589,12 @@ fn make_stateful_set(zk: &ZookeeperCluster) -> (stateful_set: StatefulSet)
             });
             selector
         });
+        stateful_set_spec.set_pvc_retention_policy({
+            let mut policy = StatefulSetPersistentVolumeClaimRetentionPolicy::default();
+            policy.set_when_deleted(new_strlit("Delete").to_string());
+            policy.set_when_scaled(new_strlit("Delete").to_string());
+            policy
+        });
         // Set the template used for creating pods
         stateful_set_spec.set_template({
             let mut pod_template_spec = PodTemplateSpec::default();
@@ -635,6 +684,25 @@ fn make_zk_pod_spec(zk: &ZookeeperCluster) -> (pod_spec: PodSpec)
                 command.push(new_strlit("/usr/local/bin/zookeeperStart.sh").to_string());
                 command
             });
+            zk_container.set_lifecycle({
+                let mut lifecycle = Lifecycle::default();
+                lifecycle.set_pre_stop({
+                    let mut pre_stop = LifecycleHandler::default();
+                    pre_stop.set_exec({
+                        let mut commands = Vec::new();
+                        commands.push(new_strlit("zookeeperTeardown.sh").to_string());
+                        proof {
+                            assert_seqs_equal!(
+                                [email protected]_values(|command: String| command@),
+                                zk_spec::make_zk_pod_spec(zk@).containers[0].lifecycle.get_Some_0().pre_stop.get_Some_0().exec_.get_Some_0()
+                            );
+                        }
+                        commands
+                    });
+                    pre_stop
+                });
+                lifecycle
+            });
             zk_container.set_image_pull_policy(new_strlit("Always").to_string());
             zk_container.set_volume_mounts({
                 let mut volume_mounts = Vec::new();

diff --git a/src/controller_examples/zookeeper_controller/exec/zookeeper_lib/helper_funcs.rs b/src/controller_examples/zookeeper_controller/exec/zookeeper_lib/helper_funcs.rs
@@ -37,24 +37,40 @@ pub fn reconcile_zk_node(path: String, uri: String, replicas: String) -> ZKNodeR
     let path_look_up = path_look_up.unwrap();
     match path_look_up {
         Some(_) => {
-            return ZKNodeResult{ res: Ok(())};
+            // Node already exists
+            // Update the cluster size
+            // This CLUSTER_SIZE will be used in zkTeardown.sh when scale down happens
+            if zk_client.set_data(
+                    path.as_rust_string_ref(),
+                    new_strlit("CLUSTER_SIZE=").to_string().concat(replicas.as_str())
+                    .as_str()
+                    .into_rust_str()
+                    .as_bytes()
+                    .to_vec(),
+                    Some(-1),
+                )
+            .is_err() {
+                return ZKNodeResult{ res: Err(Error::ClusterSizeZKNodeUpdateFailed) };
+            }
+            return ZKNodeResult{ res: Ok(()) };
         },
         None => {
+            // Node does not exist yet
             // First create the parent node
-            if zk_client
-                .create(
+            if zk_client.create(
                     "/zookeeper-operator",
                     new_strlit("")
                         .into_rust_str()
                         .as_bytes()
                         .to_vec(),
                     Acl::open_unsafe().clone(),
                     CreateMode::Persistent,
-                ).is_err() {
+                )
+            .is_err() {
                 return ZKNodeResult{ res: Err(Error::ClusterSizeZKNodeCreationFailed) };
             }
-            if zk_client
-                .create(
+            // Then create the node with correct zk.name and the cluster size
+            if zk_client.create(
                     path.as_rust_string_ref(),
                     new_strlit("CLUSTER_SIZE=").to_string().concat(replicas.as_str())
                         .as_str()
@@ -63,8 +79,9 @@ pub fn reconcile_zk_node(path: String, uri: String, replicas: String) -> ZKNodeR
                         .to_vec(),
                     Acl::open_unsafe().clone(),
                     CreateMode::Persistent,
-                ).is_err() {
-                return ZKNodeResult{ res: Err(Error::ClusterSizeZKNodeCreationFailed)};
+                )
+            .is_err() {
+                return ZKNodeResult{ res: Err(Error::ClusterSizeZKNodeCreationFailed) };
             }
             return ZKNodeResult{ res: Ok(()) };
         }

diff --git a/src/controller_examples/zookeeper_controller/exec/zookeeper_lib/lib.rs b/src/controller_examples/zookeeper_controller/exec/zookeeper_lib/lib.rs
@@ -24,15 +24,15 @@ impl ToView for ZKSupportInput {
     type V = ZKSupportInputView;
     spec fn to_view(&self) -> ZKSupportInputView {
         match self {
-            ZKSupportInput::ReconcileZKNode(path, uri, replicas) 
+            ZKSupportInput::ReconcileZKNode(path, uri, replicas)
                 => ZKSupportInputView::ReconcileZKNode(path@, uri@, replicas@),
         }
     }
 }
 
 pub proof fn zk_support_input_to_view_match(path: String, uri: String, replicas: String)
     ensures
-        ZKSupportInput::ReconcileZKNode(path, uri, replicas).to_view() 
+        ZKSupportInput::ReconcileZKNode(path, uri, replicas).to_view()
             == ZKSupportInputView::ReconcileZKNode(path@, uri@, replicas@) {}
 
 

diff --git a/src/controller_examples/zookeeper_controller/mod.rs b/src/controller_examples/zookeeper_controller/mod.rs
@@ -2,5 +2,5 @@
 // SPDX-License-Identifier: MIT
 pub mod common;
 pub mod exec;
-pub mod proof;
+// pub mod proof;
 pub mod spec;
diff --git a/src/controller_examples/zookeeper_controller/spec/reconciler.rs b/src/controller_examples/zookeeper_controller/spec/reconciler.rs
@@ -20,6 +20,7 @@ verus! {
 
 pub struct ZookeeperReconcileState {
     pub reconcile_step: ZookeeperReconcileStep,
+    pub sts_from_get: Option<StatefulSetView>,
 }
 
 pub struct ZookeeperReconciler {}
@@ -55,6 +56,7 @@ impl Reconciler<ZookeeperClusterView> for ZookeeperReconciler {
 pub open spec fn reconcile_init_state() -> ZookeeperReconcileState {
     ZookeeperReconcileState {
         reconcile_step: ZookeeperReconcileStep::Init,
+        sts_from_get: Option::None,
     }
 }
 
@@ -151,15 +153,15 @@ pub open spec fn reconcile_core(
                     // update
                     if StatefulSetView::from_dynamic_object(get_sts_resp.get_Ok_0()).is_Ok() {
                         let found_stateful_set = StatefulSetView::from_dynamic_object(get_sts_resp.get_Ok_0()).get_Ok_0();
-                        let req_o = APIRequest::UpdateRequest(UpdateRequest {
-                            key: make_stateful_set_key(zk.object_ref()),
-                            obj: update_stateful_set(zk, found_stateful_set).to_dynamic_object(),
-                        });
                         let state_prime = ZookeeperReconcileState {
-                            reconcile_step: ZookeeperReconcileStep::AfterUpdateStatefulSet,
+                            reconcile_step: ZookeeperReconcileStep::AfterUpdateZKNode,
+                            sts_from_get: Option::Some(found_stateful_set),
                             ..state
                         };
-                        (state_prime, Option::Some(RequestView::KRequest(req_o)))
+                        let ext_req = ZKSupportInputView::ReconcileZKNode(
+                            cluster_size_zk_node_path(zk), zk_service_uri(zk), int_to_string_view(zk.spec.replicas)
+                        );
+                        (state_prime, Option::Some(RequestView::ExternalRequest(ext_req)))
                     } else {
                         let state_prime = ZookeeperReconcileState {
                             reconcile_step: ZookeeperReconcileStep::Error,
@@ -213,6 +215,28 @@ pub open spec fn reconcile_core(
             );
             (state_prime, Option::Some(RequestView::ExternalRequest(ext_req)))
         },
+        ZookeeperReconcileStep::AfterUpdateZKNode => {
+            if resp_o.is_Some() && resp_o.get_Some_0().is_ExternalResponse()
+            && resp_o.get_Some_0().get_ExternalResponse_0().is_ReconcileZKNode()
+            && state.sts_from_get.is_Some(){
+                let found_stateful_set = state.sts_from_get.get_Some_0();
+                let req_o = APIRequest::UpdateRequest(UpdateRequest {
+                    key: make_stateful_set_key(zk.object_ref()),
+                    obj: update_stateful_set(zk, found_stateful_set).to_dynamic_object(),
+                });
+                let state_prime = ZookeeperReconcileState {
+                    reconcile_step: ZookeeperReconcileStep::AfterUpdateStatefulSet,
+                    sts_from_get: Option::None,
+                };
+                (state_prime,  Option::Some(RequestView::KRequest(req_o)))
+            } else {
+                let state_prime = ZookeeperReconcileState {
+                    reconcile_step: ZookeeperReconcileStep::Error,
+                    ..state
+                };
+                (state_prime, Option::None)
+            }
+        },
         ZookeeperReconcileStep::AfterCreateZKNode => {
             if resp_o.is_Some() && resp_o.get_Some_0().is_ExternalResponse()
             && resp_o.get_Some_0().get_ExternalResponse_0().is_ReconcileZKNode() {
@@ -467,6 +491,10 @@ pub open spec fn make_stateful_set(zk: ZookeeperClusterView) -> StatefulSetView
             )
             .set_spec(make_zk_pod_spec(zk))
         )
+        .set_pvc_retention_policy(StatefulSetPersistentVolumeClaimRetentionPolicyView::default()
+            .set_when_deleted(new_strlit("Delete")@)
+            .set_when_scaled(new_strlit("Delete")@)
+        )
         .set_volume_claim_templates(seq![
             PersistentVolumeClaimView::default()
                 .set_metadata(ObjectMetaView::default()
@@ -491,6 +519,11 @@ pub open spec fn make_zk_pod_spec(zk: ZookeeperClusterView) -> PodSpecView
             ContainerView::default()
                 .set_name(new_strlit("zookeeper")@)
                 .set_image(new_strlit("pravega/zookeeper:0.2.14")@)
+                .set_lifecycle(LifecycleView::default()
+                    .set_pre_stop(LifecycleHandlerView::default()
+                        .set_exec(seq![new_strlit("zookeeperTeardown.sh")@])
+                    )
+                )
                 .set_volume_mounts(seq![
                     VolumeMountView::default()
                         .set_name(new_strlit("data")@)

diff --git a/src/kubernetes_api_objects/persistent_volume_claim.rs b/src/kubernetes_api_objects/persistent_volume_claim.rs
@@ -313,4 +313,5 @@ impl Marshalable for PersistentVolumeClaimSpecView {
     proof fn marshal_preserves_integrity() {}
 }
 
+
 }