Merge pull request #34587 from MaterializeInc/push-nxowpluyvpsm

fix upgrading existing balancer and console deployments

Author: doy-materialize

ci/nightly/pipeline.template.yml

@@ -2504,3 +2504,16 @@ steps:
               ci-builder: stable
         agents:
           queue: hetzner-aarch64-16cpu-32gb
+
+      - id: orchestratord-upgrade-operator
+        label: "Orchestratord operator upgrade tests"
+        artifact_paths: ["mz_debug_*.zip"]
+        depends_on: devel-docker-tags
+        timeout_in_minutes: 120
+        plugins:
+          - ./ci/plugins/mzcompose:
+              composition: orchestratord
+              run: orchestratord-upgrade
+              ci-builder: stable
+        agents:
+          queue: hetzner-aarch64-16cpu-32gb

misc/python/materialize/cli/mzcompose.py

@@ -176,16 +176,20 @@ def load_composition(args: argparse.Namespace) -> Composition:
                 hint="If you believe this is a mistake, contact the QA team. While not recommended, --ignore-docker-version can be used to ignore this version check.",
             )
 
-        compose_local_version = Version.parse(
-            spawn.capture(["docker", "compose", "version", "--short"])
+        compose_local_version_s = spawn.capture(
+            ["docker", "compose", "version", "--short"]
         )
-        compose_ci_version = Version.parse("2.15.1")
-        if compose_local_version < compose_ci_version:
-            raise UIError(
-                f"Your Docker Compose version is {compose_local_version} while the version used in CI is {compose_ci_version}, please upgrade your local Docker Compose version to prevent unexpected breakages.",
-                hint="If you believe this is a mistake, contact the QA team. While not recommended, --ignore-docker-version can be used to ignore this version check.",
-            )
-            sys.exit(1)
+        try:
+            compose_local_version = Version.parse(compose_local_version_s)
+            compose_ci_version = Version.parse("2.15.1")
+            if compose_local_version < compose_ci_version:
+                raise UIError(
+                    f"Your Docker Compose version is {compose_local_version} while the version used in CI is {compose_ci_version}, please upgrade your local Docker Compose version to prevent unexpected breakages.",
+                    hint="If you believe this is a mistake, contact the QA team. While not recommended, --ignore-docker-version can be used to ignore this version check.",
+                )
+                sys.exit(1)
+        except ValueError:
+            pass
 
     repo = mzbuild.Repository.from_arguments(MZ_ROOT, args)
     try:

src/cloud-resources/src/crd/balancer.rs

@@ -13,7 +13,7 @@ use anyhow::bail;
 use k8s_openapi::{
     api::core::v1::ResourceRequirements, apimachinery::pkg::apis::meta::v1::Condition,
 };
-use kube::CustomResource;
+use kube::{CustomResource, Resource, ResourceExt};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
@@ -94,6 +94,10 @@ pub mod v1alpha1 {
             &self.status.as_ref().unwrap().resource_id
         }
 
+        pub fn namespace(&self) -> String {
+            self.meta().namespace.clone().unwrap()
+        }
+
         pub fn deployment_name(&self) -> String {
             self.name_prefixed("balancerd")
         }
@@ -151,6 +155,14 @@ pub mod v1alpha1 {
     impl ManagedResource for Balancer {
         fn default_labels(&self) -> BTreeMap<String, String> {
             BTreeMap::from_iter([
+                (
+                    "materialize.cloud/organization-name".to_owned(),
+                    self.name_unchecked(),
+                ),
+                (
+                    "materialize.cloud/organization-namespace".to_owned(),
+                    self.namespace(),
+                ),
                 (
                     "materialize.cloud/mz-resource-id".to_owned(),
                     self.resource_id().to_owned(),

src/cloud-resources/src/crd/console.rs

@@ -12,7 +12,7 @@ use std::collections::BTreeMap;
 use k8s_openapi::{
     api::core::v1::ResourceRequirements, apimachinery::pkg::apis::meta::v1::Condition,
 };
-use kube::CustomResource;
+use kube::{CustomResource, Resource, ResourceExt};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
@@ -91,6 +91,10 @@ pub mod v1alpha1 {
             &self.status.as_ref().unwrap().resource_id
         }
 
+        pub fn namespace(&self) -> String {
+            self.meta().namespace.clone().unwrap()
+        }
+
         pub fn deployment_name(&self) -> String {
             self.name_prefixed("console")
         }
@@ -143,6 +147,14 @@ pub mod v1alpha1 {
     impl ManagedResource for Console {
         fn default_labels(&self) -> BTreeMap<String, String> {
             BTreeMap::from_iter([
+                (
+                    "materialize.cloud/organization-name".to_owned(),
+                    self.name_unchecked(),
+                ),
+                (
+                    "materialize.cloud/organization-namespace".to_owned(),
+                    self.namespace(),
+                ),
                 (
                     "materialize.cloud/mz-resource-id".to_owned(),
                     self.resource_id().to_owned(),

src/orchestratord/src/controller/balancer.rs

@@ -7,6 +7,8 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::sync::Arc;
+
 use anyhow::bail;
 use k8s_openapi::{
     api::{
@@ -25,18 +27,20 @@ use k8s_openapi::{
 };
 use kube::{
     Api, Client, Resource, ResourceExt,
-    api::{ObjectMeta, PostParams},
+    api::{DeleteParams, ObjectMeta, PostParams},
     runtime::{
+        conditions::is_deployment_completed,
         controller::Action,
         reflector::{ObjectRef, Store},
+        wait::await_condition,
     },
 };
 use maplit::btreemap;
-use tracing::trace;
+use tracing::{trace, warn};
 
 use crate::{
     Error,
-    k8s::{apply_resource, make_reflector},
+    k8s::{apply_resource, make_reflector, replace_resource},
     tls::{DefaultCertificateSpecs, create_certificate, issuer_ref_defined},
 };
 use mz_cloud_resources::crd::{
@@ -85,7 +89,7 @@ impl Context {
         client: &Client,
         balancer: &Balancer,
     ) -> Result<(), kube::Error> {
-        let namespace = balancer.namespace().unwrap();
+        let namespace = balancer.namespace();
         let balancer_api: Api<Balancer> = Api::namespaced(client.clone(), &namespace);
 
         let Some(deployment) = self
@@ -448,6 +452,109 @@ impl Context {
             status: None,
         }
     }
+
+    // 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) = self
+            .deployments
+            .get(
+                &ObjectRef::new(&new_deployment.name_unchecked())
+                    .within(&new_deployment.namespace().unwrap()),
+            )
+            .map(Arc::unwrap_or_clone)
+        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]
@@ -478,7 +585,7 @@ impl k8s_controller::Context for Context {
             return Ok(None);
         }
 
-        let namespace = balancer.namespace().unwrap();
+        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);
@@ -489,6 +596,7 @@ impl k8s_controller::Context for Context {
         }
 
         let deployment = self.create_deployment_object(balancer)?;
+        self.fix_deployment(&deployment_api, &deployment).await?;
         trace!("creating new balancerd deployment");
         apply_resource(&deployment_api, &deployment).await?;