Skip to main content

mz_orchestratord/controller/
balancer.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use anyhow::bail;
11use k8s_controller::{Outcome, TraceMetadata};
12use k8s_openapi::{
13    api::{
14        apps::v1::{Deployment, DeploymentSpec, DeploymentStrategy, RollingUpdateDeployment},
15        core::v1::{
16            Affinity, Capabilities, ConfigMapVolumeSource, Container, ContainerPort, HTTPGetAction,
17            KeyToPath, PodSecurityContext, PodSpec, PodTemplateSpec, Probe, ResourceRequirements,
18            SeccompProfile, SecretVolumeSource, SecurityContext, Service, ServicePort, ServiceSpec,
19            Toleration, Volume, VolumeMount,
20        },
21    },
22    apimachinery::pkg::{
23        apis::meta::v1::{Condition, LabelSelector, Time},
24        util::intstr::IntOrString,
25    },
26    jiff::Timestamp,
27};
28use kube::{
29    Api, Client, Resource, ResourceExt,
30    api::{DeleteParams, ObjectMeta, PostParams},
31    runtime::{conditions::is_deployment_completed, controller::Action, wait::await_condition},
32};
33use maplit::btreemap;
34use tracing::{trace, warn};
35
36use crate::{
37    Error,
38    k8s::{apply_resource, get_resource, recommended_k8s_labels, replace_resource},
39    tls::{DefaultCertificateSpecs, create_certificate, issuer_ref_defined},
40};
41use mz_cloud_resources::crd::{
42    ManagedResource,
43    balancer::v1alpha1::{Balancer, Routing},
44    generated::cert_manager::certificates::{Certificate, CertificatePrivateKeyAlgorithm},
45};
46use mz_orchestrator_kubernetes::KubernetesImagePullPolicy;
47use mz_ore::{cli::KeyValueArg, instrument};
48
49/// The name identifying this controller in its reconciliation metrics and in
50/// the reporter of the events it publishes.
51pub const CONTROLLER_NAME: &str = "balancer";
52
53#[derive(Clone)]
54pub struct Config {
55    pub enable_security_context: bool,
56    pub enable_prometheus_scrape_annotations: bool,
57
58    pub image_pull_policy: KubernetesImagePullPolicy,
59    pub scheduler_name: Option<String>,
60    pub balancerd_node_selector: Vec<KeyValueArg<String, String>>,
61    pub balancerd_affinity: Option<Affinity>,
62    pub balancerd_tolerations: Option<Vec<Toleration>>,
63    pub balancerd_default_resources: Option<ResourceRequirements>,
64
65    pub default_certificate_specs: DefaultCertificateSpecs,
66
67    pub environmentd_sql_port: u16,
68    pub environmentd_http_port: u16,
69    pub balancerd_sql_port: u16,
70    pub balancerd_http_port: u16,
71    pub balancerd_internal_http_port: u16,
72}
73
74pub struct Context {
75    config: Config,
76}
77
78impl Context {
79    pub fn new(config: Config) -> Self {
80        Self { config }
81    }
82
83    async fn sync_deployment_status(
84        &self,
85        client: &Client,
86        balancer: &Balancer,
87    ) -> Result<(), Error> {
88        let namespace = balancer.namespace();
89        let balancer_api: Api<Balancer> = Api::namespaced(client.clone(), &namespace);
90        let deployment_api: Api<Deployment> = Api::namespaced(client.clone(), &namespace);
91
92        let Some(deployment) = get_resource(&deployment_api, &balancer.deployment_name()).await?
93        else {
94            return Ok(());
95        };
96
97        let Some(deployment_conditions) = &deployment
98            .status
99            .as_ref()
100            .and_then(|status| status.conditions.as_ref())
101        else {
102            // if the deployment doesn't have any conditions set yet, there
103            // is nothing to sync
104            return Ok(());
105        };
106
107        let ready = deployment_conditions
108            .iter()
109            .any(|condition| condition.type_ == "Available" && condition.status == "True");
110        let ready_str = if ready { "True" } else { "False" };
111
112        let mut status = balancer.status.clone().unwrap();
113        if status
114            .conditions
115            .iter()
116            .any(|condition| condition.type_ == "Ready" && condition.status == ready_str)
117        {
118            // if the deployment status is already set correctly, we don't
119            // need to set it again (this prevents us from getting stuck in
120            // a reconcile loop)
121            return Ok(());
122        }
123
124        status.conditions = vec![Condition {
125            type_: "Ready".to_string(),
126            status: ready_str.to_string(),
127            last_transition_time: Time(Timestamp::now()),
128            message: format!(
129                "balancerd deployment is{} ready",
130                if ready { "" } else { " not" }
131            ),
132            observed_generation: None,
133            reason: "DeploymentStatus".to_string(),
134        }];
135        let mut new_balancer = balancer.clone();
136        new_balancer.status = Some(status);
137
138        balancer_api
139            .replace_status(
140                &balancer.name_unchecked(),
141                &PostParams::default(),
142                &new_balancer,
143            )
144            .await?;
145
146        Ok(())
147    }
148
149    fn create_external_certificate_object(
150        &self,
151        balancer: &Balancer,
152    ) -> Result<Option<Certificate>, anyhow::Error> {
153        create_certificate(
154            self.config
155                .default_certificate_specs
156                .balancerd_external
157                .clone(),
158            balancer,
159            balancer.spec.external_certificate_spec.clone(),
160            balancer.external_certificate_name(),
161            balancer.external_certificate_secret_name(),
162            None,
163            CertificatePrivateKeyAlgorithm::Ecdsa,
164            Some(256),
165        )
166    }
167
168    fn create_deployment_object(&self, balancer: &Balancer) -> anyhow::Result<Deployment> {
169        let security_context = if self.config.enable_security_context {
170            // Since we want to adhere to the most restrictive security context, all
171            // of these fields have to be set how they are.
172            // See https://kubernetes.io/docs/concepts/security/pod-security-standards/#restricted
173            Some(SecurityContext {
174                run_as_non_root: Some(true),
175                capabilities: Some(Capabilities {
176                    drop: Some(vec!["ALL".to_string()]),
177                    ..Default::default()
178                }),
179                seccomp_profile: Some(SeccompProfile {
180                    type_: "RuntimeDefault".to_string(),
181                    ..Default::default()
182                }),
183                allow_privilege_escalation: Some(false),
184                ..Default::default()
185            })
186        } else {
187            None
188        };
189
190        let pod_template_annotations = if self.config.enable_prometheus_scrape_annotations {
191            Some(btreemap! {
192                "prometheus.io/scrape".to_owned() => "true".to_string(),
193                "prometheus.io/port".to_owned() => self.config.balancerd_internal_http_port.to_string(),
194                "prometheus.io/path".to_owned() => "/metrics".to_string(),
195                "prometheus.io/scheme".to_owned() => "http".to_string(),
196            })
197        } else {
198            None
199        };
200        let mut pod_template_labels = balancer.default_labels();
201        pod_template_labels.insert(
202            "materialize.cloud/name".to_owned(),
203            balancer.deployment_name(),
204        );
205        pod_template_labels.insert("app".to_owned(), "balancerd".to_string());
206        pod_template_labels.insert("materialize.cloud/app".to_owned(), balancer.app_name());
207
208        let ports = vec![
209            ContainerPort {
210                container_port: self.config.balancerd_sql_port.into(),
211                name: Some("pgwire".into()),
212                protocol: Some("TCP".into()),
213                ..Default::default()
214            },
215            ContainerPort {
216                container_port: self.config.balancerd_http_port.into(),
217                name: Some("http".into()),
218                protocol: Some("TCP".into()),
219                ..Default::default()
220            },
221            ContainerPort {
222                container_port: self.config.balancerd_internal_http_port.into(),
223                name: Some("internal-http".into()),
224                protocol: Some("TCP".into()),
225                ..Default::default()
226            },
227        ];
228
229        let mut args = vec![
230            "service".to_string(),
231            format!(
232                "--pgwire-listen-addr=0.0.0.0:{}",
233                self.config.balancerd_sql_port
234            ),
235            format!(
236                "--https-listen-addr=0.0.0.0:{}",
237                self.config.balancerd_http_port
238            ),
239            format!(
240                "--internal-http-listen-addr=0.0.0.0:{}",
241                self.config.balancerd_internal_http_port
242            ),
243        ];
244        match balancer.routing()? {
245            Routing::Static(static_routing_config) => {
246                args.extend([
247                    format!(
248                        "--https-resolver-template={}.{}.svc.cluster.local:{}",
249                        static_routing_config.environmentd_service_name,
250                        static_routing_config.environmentd_namespace,
251                        self.config.environmentd_http_port
252                    ),
253                    format!(
254                        "--static-resolver-addr={}.{}.svc.cluster.local:{}",
255                        static_routing_config.environmentd_service_name,
256                        static_routing_config.environmentd_namespace,
257                        self.config.environmentd_sql_port
258                    ),
259                ]);
260            }
261            Routing::Frontegg(_frontegg_routing_config) => {
262                bail!("frontegg routing is not yet implemented");
263            }
264        }
265
266        if issuer_ref_defined(
267            &self.config.default_certificate_specs.internal,
268            &balancer.spec.internal_certificate_spec,
269        ) {
270            args.push("--internal-tls".to_owned())
271        }
272
273        let mut volumes = Vec::new();
274        let mut volume_mounts = Vec::new();
275        if let Some(name) = &balancer.spec.configmap_name {
276            volumes.push(Volume {
277                name: "dynamic-config".to_owned(),
278                config_map: Some(ConfigMapVolumeSource {
279                    name: name.clone(),
280                    // The file must exist at startup or balancerd skips the sync loop.
281                    optional: Some(false),
282                    items: Some(vec![KeyToPath {
283                        key: "config.json".to_owned(),
284                        path: "config.json".to_owned(),
285                        ..Default::default()
286                    }]),
287                    ..Default::default()
288                }),
289                ..Default::default()
290            });
291            volume_mounts.push(VolumeMount {
292                name: "dynamic-config".to_owned(),
293                // Mount the directory: subPath mounts do not receive ConfigMap updates.
294                mount_path: "/etc/balancerd".to_owned(),
295                read_only: Some(true),
296                ..Default::default()
297            });
298            args.extend([
299                "--config-sync-file-path=/etc/balancerd/config.json".to_owned(),
300                "--config-sync-loop-interval=1s".to_owned(),
301            ]);
302        }
303        if issuer_ref_defined(
304            &self.config.default_certificate_specs.balancerd_external,
305            &balancer.spec.external_certificate_spec,
306        ) {
307            volumes.push(Volume {
308                name: "external-certificate".to_owned(),
309                secret: Some(SecretVolumeSource {
310                    default_mode: Some(0o400),
311                    secret_name: Some(balancer.external_certificate_secret_name()),
312                    items: None,
313                    optional: Some(false),
314                }),
315                ..Default::default()
316            });
317            volume_mounts.push(VolumeMount {
318                name: "external-certificate".to_owned(),
319                mount_path: "/etc/external_tls".to_owned(),
320                read_only: Some(true),
321                ..Default::default()
322            });
323            args.extend([
324                "--tls-mode=require".into(),
325                "--tls-cert=/etc/external_tls/tls.crt".into(),
326                "--tls-key=/etc/external_tls/tls.key".into(),
327            ]);
328        } else {
329            args.push("--tls-mode=disable".to_string());
330        }
331
332        let startup_probe = Probe {
333            http_get: Some(HTTPGetAction {
334                port: IntOrString::Int(self.config.balancerd_internal_http_port.into()),
335                path: Some("/api/readyz".into()),
336                ..Default::default()
337            }),
338            failure_threshold: Some(20),
339            initial_delay_seconds: Some(3),
340            period_seconds: Some(3),
341            success_threshold: Some(1),
342            timeout_seconds: Some(1),
343            ..Default::default()
344        };
345        let readiness_probe = Probe {
346            http_get: Some(HTTPGetAction {
347                port: IntOrString::Int(self.config.balancerd_internal_http_port.into()),
348                path: Some("/api/readyz".into()),
349                ..Default::default()
350            }),
351            failure_threshold: Some(3),
352            period_seconds: Some(10),
353            success_threshold: Some(1),
354            timeout_seconds: Some(1),
355            ..Default::default()
356        };
357        let liveness_probe = Probe {
358            http_get: Some(HTTPGetAction {
359                port: IntOrString::Int(self.config.balancerd_internal_http_port.into()),
360                path: Some("/api/livez".into()),
361                ..Default::default()
362            }),
363            failure_threshold: Some(3),
364            initial_delay_seconds: Some(8),
365            period_seconds: Some(10),
366            success_threshold: Some(1),
367            timeout_seconds: Some(1),
368            ..Default::default()
369        };
370
371        let container = Container {
372            name: "balancerd".to_owned(),
373            image: Some(balancer.spec.balancerd_image_ref.clone()),
374            image_pull_policy: Some(self.config.image_pull_policy.to_string()),
375            ports: Some(ports),
376            args: Some(args),
377            startup_probe: Some(startup_probe),
378            readiness_probe: Some(readiness_probe),
379            liveness_probe: Some(liveness_probe),
380            resources: balancer
381                .spec
382                .resource_requirements
383                .clone()
384                .or_else(|| self.config.balancerd_default_resources.clone()),
385            security_context: security_context.clone(),
386            volume_mounts: Some(volume_mounts),
387            ..Default::default()
388        };
389
390        let match_labels = pod_template_labels.clone();
391        pod_template_labels.extend(recommended_k8s_labels(balancer.app_name()));
392
393        let deployment_spec = DeploymentSpec {
394            replicas: Some(balancer.replicas()),
395            selector: LabelSelector {
396                match_labels: Some(match_labels),
397                ..Default::default()
398            },
399            strategy: Some(DeploymentStrategy {
400                rolling_update: Some(RollingUpdateDeployment {
401                    // Allow a complete set of new pods at once, to minimize the
402                    // chances of a new connection going to a pod that will be
403                    // immediately drained
404                    max_surge: Some(IntOrString::String("100%".into())),
405                    ..Default::default()
406                }),
407                ..Default::default()
408            }),
409            template: PodTemplateSpec {
410                // not using managed_resource_meta because the pod should be
411                // owned by the deployment, not the materialize instance
412                metadata: Some(ObjectMeta {
413                    annotations: pod_template_annotations,
414                    labels: Some(pod_template_labels),
415                    ..Default::default()
416                }),
417                spec: Some(PodSpec {
418                    containers: vec![container],
419                    node_selector: Some(
420                        self.config
421                            .balancerd_node_selector
422                            .iter()
423                            .map(|selector| (selector.key.clone(), selector.value.clone()))
424                            .collect(),
425                    ),
426                    affinity: self.config.balancerd_affinity.clone(),
427                    tolerations: self.config.balancerd_tolerations.clone(),
428                    security_context: Some(PodSecurityContext {
429                        fs_group: Some(999),
430                        run_as_user: Some(999),
431                        run_as_group: Some(999),
432                        ..Default::default()
433                    }),
434                    scheduler_name: self.config.scheduler_name.clone(),
435                    volumes: Some(volumes),
436                    ..Default::default()
437                }),
438            },
439            ..Default::default()
440        };
441
442        Ok(Deployment {
443            metadata: balancer.managed_resource_meta(balancer.deployment_name()),
444            spec: Some(deployment_spec),
445            status: None,
446        })
447    }
448
449    fn create_service_object(&self, balancer: &Balancer) -> Service {
450        let selector =
451            btreemap! {"materialize.cloud/name".to_string() => balancer.deployment_name()};
452
453        let ports = vec![
454            ServicePort {
455                name: Some("http".to_string()),
456                protocol: Some("TCP".to_string()),
457                port: self.config.balancerd_http_port.into(),
458                target_port: Some(IntOrString::Int(self.config.balancerd_http_port.into())),
459                ..Default::default()
460            },
461            ServicePort {
462                name: Some("pgwire".to_string()),
463                protocol: Some("TCP".to_string()),
464                port: self.config.balancerd_sql_port.into(),
465                target_port: Some(IntOrString::Int(self.config.balancerd_sql_port.into())),
466                ..Default::default()
467            },
468        ];
469
470        let spec = ServiceSpec {
471            type_: Some("ClusterIP".to_string()),
472            cluster_ip: Some("None".to_string()),
473            selector: Some(selector),
474            ports: Some(ports),
475            ..Default::default()
476        };
477
478        Service {
479            metadata: balancer.managed_resource_meta(balancer.service_name()),
480            spec: Some(spec),
481            status: None,
482        }
483    }
484
485    // TODO: remove this once everyone is upgraded to an orchestratord
486    // version with the separate balancer operator
487    async fn fix_deployment(
488        &self,
489        deployment_api: &Api<Deployment>,
490        new_deployment: &Deployment,
491    ) -> Result<(), Error> {
492        let Some(mut existing_deployment) =
493            get_resource(deployment_api, &new_deployment.name_unchecked()).await?
494        else {
495            return Ok(());
496        };
497
498        if existing_deployment.spec.as_ref().unwrap().selector
499            == new_deployment.spec.as_ref().unwrap().selector
500        {
501            return Ok(());
502        }
503
504        warn!("found existing deployment with old label selector, fixing");
505
506        // this is sufficient because the new labels are a superset of the
507        // old labels, so the existing label selector should still be valid
508        existing_deployment
509            .spec
510            .as_mut()
511            .unwrap()
512            .template
513            .metadata
514            .as_mut()
515            .unwrap()
516            .labels = new_deployment
517            .spec
518            .as_ref()
519            .unwrap()
520            .template
521            .metadata
522            .as_ref()
523            .unwrap()
524            .labels
525            .clone();
526
527        // using await_condition is not ideal in a controller loop, but this
528        // is very temporary and will only ever happen once, so this feels
529        // simpler than trying to introduce an entire state machine here
530        replace_resource(deployment_api, &existing_deployment).await?;
531        await_condition(
532            deployment_api.clone(),
533            &existing_deployment.name_unchecked(),
534            |deployment: Option<&Deployment>| {
535                let observed_generation = deployment
536                    .and_then(|deployment| deployment.status.as_ref())
537                    .and_then(|status| status.observed_generation)
538                    .unwrap_or(0);
539                let current_generation = deployment
540                    .and_then(|deployment| deployment.meta().generation)
541                    .unwrap_or(0);
542                let previous_generation = existing_deployment.meta().generation.unwrap_or(0);
543                observed_generation == current_generation
544                    && current_generation > previous_generation
545            },
546        )
547        .await
548        .map_err(|e| anyhow::anyhow!(e))?;
549        await_condition(
550            deployment_api.clone(),
551            &existing_deployment.name_unchecked(),
552            is_deployment_completed(),
553        )
554        .await
555        .map_err(|e| anyhow::anyhow!(e))?;
556
557        // delete the deployment but leave the pods around (via
558        // DeleteParams::orphan)
559        match kube::runtime::wait::delete::delete_and_finalize(
560            deployment_api.clone(),
561            &existing_deployment.name_unchecked(),
562            &DeleteParams::orphan(),
563        )
564        .await
565        {
566            Ok(_) => {}
567            Err(kube::runtime::wait::delete::Error::Delete(kube::Error::Api(e)))
568                if e.code == 404 =>
569            {
570                // the resource already doesn't exist
571            }
572            Err(e) => return Err(anyhow::anyhow!(e).into()),
573        }
574
575        // now, the normal apply of the new deployment (in the main loop)
576        // will take over the existing pods from the old deployment we just
577        // deleted, since we already updated the pod labels to be the same as
578        // the new label selector
579
580        Ok(())
581    }
582}
583
584#[async_trait::async_trait]
585impl k8s_controller::Context for Context {
586    type Resource = Balancer;
587    type Error = Error;
588
589    #[instrument(fields())]
590    async fn apply(
591        &self,
592        client: Client,
593        balancer: &Self::Resource,
594        metadata: &mut TraceMetadata,
595    ) -> Result<Option<Action>, Self::Error> {
596        if balancer.status.is_none() {
597            let step = metadata.step("initialize_status");
598            let balancer_api: Api<Balancer> =
599                Api::namespaced(client.clone(), &balancer.meta().namespace.clone().unwrap());
600            let mut new_balancer = balancer.clone();
601            new_balancer.status = Some(balancer.status());
602            balancer_api
603                .replace_status(
604                    &balancer.name_unchecked(),
605                    &PostParams::default(),
606                    &new_balancer,
607                )
608                .await?;
609            step.finish(Outcome::Completed);
610            // Updating the status should trigger a reconciliation
611            // which will include a status this time.
612            return Ok(None);
613        }
614
615        let namespace = balancer.namespace();
616        let certificate_api: Api<Certificate> = Api::namespaced(client.clone(), &namespace);
617        let deployment_api: Api<Deployment> = Api::namespaced(client.clone(), &namespace);
618        let service_api: Api<Service> = Api::namespaced(client.clone(), &namespace);
619
620        let step = metadata.step("certificate");
621        if let Some(external_certificate) = self.create_external_certificate_object(balancer)? {
622            trace!("creating new balancerd external certificate");
623            apply_resource(&certificate_api, &external_certificate).await?;
624            step.finish(Outcome::Completed);
625        } else {
626            step.finish(Outcome::Skipped);
627        }
628
629        let step = metadata.step("deployment");
630        let deployment = self.create_deployment_object(balancer)?;
631        self.fix_deployment(&deployment_api, &deployment).await?;
632        trace!("creating new balancerd deployment");
633        apply_resource(&deployment_api, &deployment).await?;
634        step.finish(Outcome::Completed);
635
636        let step = metadata.step("service");
637        let service = self.create_service_object(balancer);
638        trace!("creating new balancerd service");
639        apply_resource(&service_api, &service).await?;
640        step.finish(Outcome::Completed);
641
642        let step = metadata.step("sync_status");
643        self.sync_deployment_status(&client, balancer).await?;
644        step.finish(Outcome::Completed);
645
646        Ok(None)
647    }
648}
649
650#[cfg(test)]
651mod tests {
652    use super::*;
653    use mz_cloud_resources::crd::balancer::v1alpha1::{BalancerSpec, StaticRoutingConfig};
654
655    #[mz_ore::test]
656    fn configmap_reference_controls_file_sync() {
657        let context = Context::new(Config {
658            enable_security_context: false,
659            enable_prometheus_scrape_annotations: false,
660            image_pull_policy: KubernetesImagePullPolicy::IfNotPresent,
661            scheduler_name: None,
662            balancerd_node_selector: vec![],
663            balancerd_affinity: None,
664            balancerd_tolerations: None,
665            balancerd_default_resources: None,
666            default_certificate_specs: DefaultCertificateSpecs::default(),
667            environmentd_sql_port: 6875,
668            environmentd_http_port: 6876,
669            balancerd_sql_port: 6875,
670            balancerd_http_port: 6876,
671            balancerd_internal_http_port: 6878,
672        });
673        for configmap_name in [None, Some("balancerd-settings".to_owned())] {
674            let mut balancer = Balancer::new(
675                "test",
676                BalancerSpec {
677                    configmap_name: configmap_name.clone(),
678                    static_routing: Some(StaticRoutingConfig {
679                        environmentd_namespace: "test".to_owned(),
680                        environmentd_service_name: "environmentd".to_owned(),
681                    }),
682                    ..Default::default()
683                },
684            );
685            balancer.metadata.namespace = Some("test".to_owned());
686            balancer.metadata.uid = Some("test".to_owned());
687            balancer.status = Some(balancer.status());
688            let pod = context
689                .create_deployment_object(&balancer)
690                .unwrap()
691                .spec
692                .unwrap()
693                .template
694                .spec
695                .unwrap();
696            let container = &pod.containers[0];
697            let args = container.args.as_ref().unwrap();
698            let volumes = pod.volumes.as_ref().unwrap();
699            let mounts = container.volume_mounts.as_ref().unwrap();
700            if let Some(name) = configmap_name {
701                let source = volumes[0].config_map.as_ref().unwrap();
702                assert_eq!(source.name, name);
703                assert_eq!(source.optional, Some(false));
704                assert_eq!(source.items.as_ref().unwrap()[0].key, "config.json");
705                assert_eq!(mounts[0].read_only, Some(true));
706                assert!(mounts[0].sub_path.is_none());
707                assert!(args.contains(&format!(
708                    "--config-sync-file-path={}/config.json",
709                    mounts[0].mount_path
710                )));
711                assert!(args.contains(&"--config-sync-loop-interval=1s".to_owned()));
712            } else {
713                assert!(volumes.is_empty());
714                assert!(mounts.is_empty());
715                assert!(!args.iter().any(|arg| arg.starts_with("--config-sync-")));
716            }
717        }
718    }
719}