1use 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
49pub 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 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 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 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 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_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 max_surge: Some(IntOrString::String("100%".into())),
405 ..Default::default()
406 }),
407 ..Default::default()
408 }),
409 template: PodTemplateSpec {
410 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 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 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 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 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 }
572 Err(e) => return Err(anyhow::anyhow!(e).into()),
573 }
574
575 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 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}