diff --git a/CHANGELOG.md b/CHANGELOG.md index e3356a5d8..804ab3a85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +### Added + +- Extended the `PodBuilder` with `pod_affinity`, `pod_anti_affinity`, `node_selector` and their `*_opt` variants ([#520]) + +[#520]: https://github.com/stackabletech/operator-rs/pull/520 + ## [0.29.0] - 2022-12-16 ### Added diff --git a/src/builder/pod/mod.rs b/src/builder/pod/mod.rs index 108e1a3c2..ddb70f224 100644 --- a/src/builder/pod/mod.rs +++ b/src/builder/pod/mod.rs @@ -6,20 +6,18 @@ use crate::builder::meta::ObjectMetaBuilder; use crate::commons::product_image_selection::ResolvedProductImage; use crate::error::{Error, OperatorResult}; +use super::{ListenerOperatorVolumeSourceBuilder, ListenerReference}; use k8s_openapi::{ api::core::v1::{ Affinity, Container, LocalObjectReference, NodeAffinity, NodeSelector, - NodeSelectorRequirement, NodeSelectorTerm, Pod, PodAffinity, PodCondition, + NodeSelectorRequirement, NodeSelectorTerm, Pod, PodAffinity, PodAntiAffinity, PodCondition, PodSecurityContext, PodSpec, PodStatus, PodTemplateSpec, Toleration, Volume, }, apimachinery::pkg::apis::meta::v1::{LabelSelector, LabelSelectorRequirement, ObjectMeta}, }; use std::collections::BTreeMap; -use super::{ListenerOperatorVolumeSourceBuilder, ListenerReference}; - -/// A builder to build [`Pod`] objects. -/// +/// A builder to build [`Pod`] or [`PodTemplateSpec`] objects. #[derive(Clone, Default)] pub struct PodBuilder { containers: Vec, @@ -29,6 +27,7 @@ pub struct PodBuilder { node_name: Option, node_selector: Option, pod_affinity: Option, + pod_anti_affinity: Option, status: Option, security_context: Option, tolerations: Option>, @@ -87,11 +86,31 @@ impl PodBuilder { self } + pub fn pod_affinity_opt(&mut self, affinity: Option) -> &mut Self { + self.pod_affinity = affinity; + self + } + + pub fn pod_anti_affinity(&mut self, anti_affinity: PodAntiAffinity) -> &mut Self { + self.pod_anti_affinity = Some(anti_affinity); + self + } + + pub fn pod_anti_affinity_opt(&mut self, anti_affinity: Option) -> &mut Self { + self.pod_anti_affinity = anti_affinity; + self + } + pub fn node_selector(&mut self, node_selector: LabelSelector) -> &mut Self { self.node_selector = Some(node_selector); self } + pub fn node_selector_opt(&mut self, node_selector: Option) -> &mut Self { + self.node_selector = node_selector; + self + } + pub fn phase(&mut self, phase: &str) -> &mut Self { let mut status = self.status.get_or_insert_with(PodStatus::default); status.phase = Some(phase.to_string()); @@ -354,18 +373,11 @@ impl PodBuilder { init_containers: self.init_containers.clone(), node_name: self.node_name.clone(), node_selector: node_selector_labels, - affinity: node_affinity - .map(|node_affinity| Affinity { - node_affinity: Some(node_affinity), - pod_affinity: self.pod_affinity.clone(), - ..Affinity::default() - }) - .or_else(|| { - Some(Affinity { - pod_affinity: self.pod_affinity.clone(), - ..Affinity::default() - }) - }), + affinity: Some(Affinity { + node_affinity, + pod_affinity: self.pod_affinity.clone(), + pod_anti_affinity: self.pod_anti_affinity.clone(), + }), security_context: self.security_context.clone(), tolerations: self.tolerations.clone(), volumes: self.volumes.clone(), @@ -406,33 +418,71 @@ impl PodBuilder { #[cfg(test)] mod tests { + use super::*; use crate::builder::{ meta::ObjectMetaBuilder, - pod::{container::ContainerBuilder, volume::VolumeBuilder, PodBuilder}, + pod::{container::ContainerBuilder, volume::VolumeBuilder}, }; use k8s_openapi::{ api::core::v1::{LocalObjectReference, PodAffinity, PodAffinityTerm}, apimachinery::pkg::apis::meta::v1::{LabelSelector, LabelSelectorRequirement}, }; + use rstest::*; - #[test] - fn test_pod_builder() { - let container = ContainerBuilder::new("containername") + // A simple [`Container`] with a name and image. + #[fixture] + fn dummy_container() -> Container { + ContainerBuilder::new("container") .expect("ContainerBuilder not created") - .image("stackable/zookeeper:2.4.14") - .command(vec!["zk-server-start.sh".to_string()]) - .args(vec!["stackable/conf/zk.properties".to_string()]) - .add_volume_mount("zk-worker-1", "conf/") - .build(); + .image("private-company/product:2.4.14") + .build() + } - let init_container = ContainerBuilder::new("init-containername") - .expect("ContainerBuilder not created") - .image("stackable/zookeeper:2.4.14") - .command(vec!["wrapper.sh".to_string()]) - .args(vec!["12345".to_string()]) - .build(); + /// A [`PodBuilder`] that already contains the minum setup to build a Pod (name and container). + #[fixture] + fn pod_builder_with_name_and_container(dummy_container: Container) -> PodBuilder { + let mut builder = PodBuilder::new(); + builder + .metadata(ObjectMetaBuilder::new().name("testpod").build()) + .add_container(dummy_container); + builder + } + + // A fixture for a node selector to use on a Pod, and the resulting node selector labels and node affinity. + #[fixture] + fn node_selector1() -> ( + LabelSelector, + Option>, + Option, + ) { + let labels = BTreeMap::from([("key".to_owned(), "value".to_owned())]); + let label_selector = LabelSelector { + match_expressions: Some(vec![LabelSelectorRequirement { + key: "security".to_owned(), + operator: "In".to_owned(), + values: Some(vec!["S1".to_owned(), "S2".to_owned()]), + }]), + match_labels: Some(labels.clone()), + }; + let affinity = Some(NodeAffinity { + required_during_scheduling_ignored_during_execution: Some(NodeSelector { + node_selector_terms: vec![NodeSelectorTerm { + match_expressions: Some(vec![NodeSelectorRequirement { + key: "security".to_owned(), + operator: "In".to_owned(), + values: Some(vec!["S1".to_owned(), "S2".to_owned()]), + }]), + ..Default::default() + }], + }), + ..Default::default() + }); + (label_selector, Some(labels), affinity) + } - let pod_affinity = PodAffinity { + #[fixture] + fn pod_affinity() -> PodAffinity { + PodAffinity { preferred_during_scheduling_ignored_during_execution: None, required_during_scheduling_ignored_during_execution: Some(vec![PodAffinityTerm { label_selector: Some(LabelSelector { @@ -446,12 +496,39 @@ mod tests { topology_key: "topology.kubernetes.io/zone".to_string(), ..Default::default() }]), - }; + } + } + + #[fixture] + fn pod_anti_affinity(pod_affinity: PodAffinity) -> PodAntiAffinity { + PodAntiAffinity { + preferred_during_scheduling_ignored_during_execution: None, + required_during_scheduling_ignored_during_execution: pod_affinity + .required_during_scheduling_ignored_during_execution, + } + } + + #[rstest] + fn test_pod_builder_pod_name() { + let pod = PodBuilder::new() + .metadata_builder(|builder| builder.name("foo")) + .build() + .unwrap(); + + assert_eq!(pod.metadata.name.unwrap(), "foo"); + } + + #[rstest] + fn test_pod_builder(pod_affinity: PodAffinity, dummy_container: Container) { + let init_container = ContainerBuilder::new("init-containername") + .expect("ContainerBuilder not created") + .image("stackable/zookeeper:2.4.14") + .build(); let pod = PodBuilder::new() .pod_affinity(pod_affinity.clone()) .metadata(ObjectMetaBuilder::new().name("testpod").build()) - .add_container(container) + .add_container(dummy_container) .add_init_container(init_container) .node_name("worker-1.stackable.demo") .add_volume( @@ -486,25 +563,11 @@ mod tests { .and_then(|volume| volume.config_map.as_ref()?.name.clone())), Some("configmap".to_string()) ); - - let pod = PodBuilder::new() - .metadata_builder(|builder| builder.name("foo")) - .build() - .unwrap(); - - assert_eq!(pod.metadata.name.unwrap(), "foo"); } - #[test] - fn test_pod_builder_image_pull_secrets() { - let container = ContainerBuilder::new("container") - .expect("ContainerBuilder not created") - .image("private-comapany/product:2.4.14") - .build(); - - let pod = PodBuilder::new() - .metadata(ObjectMetaBuilder::new().name("testpod").build()) - .add_container(container) + #[rstest] + fn test_pod_builder_image_pull_secrets(mut pod_builder_with_name_and_container: PodBuilder) { + let pod = pod_builder_with_name_and_container .image_pull_secrets(vec!["company-registry-secret".to_string()].into_iter()) .build() .unwrap(); @@ -516,4 +579,88 @@ mod tests { }] ); } + + /// Test if setting a node selector generates the correct node selector labels and node affinity on the Pod. + #[rstest] + fn test_pod_builder_node_selector( + mut pod_builder_with_name_and_container: PodBuilder, + node_selector1: ( + LabelSelector, + Option>, + Option, + ), + ) { + // destruct fixture + let (node_selector, expected_labels, expected_affinity) = node_selector1; + // first test with the normal node_selector function + let pod = pod_builder_with_name_and_container + .clone() + .node_selector(node_selector.clone()) + .build() + .unwrap(); + + let spec = pod.spec.unwrap(); + assert_eq!(spec.node_selector, expected_labels); + assert_eq!(spec.affinity.unwrap().node_affinity, expected_affinity); + + // test the node_selector_opt function + let pod = pod_builder_with_name_and_container + .node_selector_opt(Some(node_selector)) + .build() + .unwrap(); + + // asserts + let spec = pod.spec.unwrap(); + assert_eq!(spec.node_selector, expected_labels); + assert_eq!(spec.affinity.unwrap().node_affinity, expected_affinity); + } + + /// Test if setting a node selector generates the correct node selector labels and node affinity on the Pod, + /// while keeping the manually set Pod affinities. Since they are mangled together, it makes sense to make sure that + /// one is not replacing the other + #[rstest] + fn test_pod_builder_node_selector_and_affinity( + mut pod_builder_with_name_and_container: PodBuilder, + node_selector1: ( + LabelSelector, + Option>, + Option, + ), + pod_affinity: PodAffinity, + pod_anti_affinity: PodAntiAffinity, + ) { + // destruct fixture + let (node_selector, expected_labels, expected_affinity) = node_selector1; + // first test with the normal functions + let pod = pod_builder_with_name_and_container + .clone() + .node_selector(node_selector.clone()) + .pod_affinity(pod_affinity.clone()) + .pod_anti_affinity(pod_anti_affinity.clone()) + .build() + .unwrap(); + + let spec = pod.spec.unwrap(); + assert_eq!(spec.node_selector, expected_labels); + let affinity = spec.affinity.unwrap(); + assert_eq!(affinity.node_affinity, expected_affinity); + assert_eq!(affinity.pod_affinity, Some(pod_affinity.clone())); + assert_eq!(affinity.pod_anti_affinity, Some(pod_anti_affinity.clone())); + + // test the *_opt functions + let pod = pod_builder_with_name_and_container + .node_selector_opt(Some(node_selector)) + .pod_affinity_opt(Some(pod_affinity.clone())) + .pod_anti_affinity_opt(Some(pod_anti_affinity.clone())) + .build() + .unwrap(); + + // asserts + let spec = pod.spec.unwrap(); + assert_eq!(spec.node_selector, expected_labels); + let affinity = spec.affinity.unwrap(); + assert_eq!(affinity.node_affinity, expected_affinity); + assert_eq!(affinity.pod_affinity, Some(pod_affinity)); + assert_eq!(affinity.pod_anti_affinity, Some(pod_anti_affinity)); + } } diff --git a/src/product_config_utils.rs b/src/product_config_utils.rs index ffb404778..5304e0315 100644 --- a/src/product_config_utils.rs +++ b/src/product_config_utils.rs @@ -625,6 +625,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (true, true, true, false) => Role { @@ -639,6 +640,7 @@ mod tests { config: build_common_config( build_test_config(GROUP_CONFIG, GROUP_ENV, GROUP_CLI), None, None, None), selector: None, + affinity: None, }}, }, (true, true, false, true) => Role { @@ -656,6 +658,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (true, true, false, false) => Role { @@ -673,6 +676,7 @@ mod tests { None, None), selector: None, + affinity: None, }}, }, (true, false, true, true) => Role { @@ -690,6 +694,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (true, false, true, false) => Role { @@ -703,6 +708,7 @@ mod tests { replicas: Some(1), config: CommonConfiguration::default(), selector: None, + affinity: None, }}, }, (true, false, false, true) => Role { @@ -721,6 +727,7 @@ mod tests { build_cli_override(GROUP_CLI_OVERRIDE) ), selector: None, + affinity: None, }}, }, (true, false, false, false) => Role { @@ -734,6 +741,7 @@ mod tests { replicas: Some(1), config: CommonConfiguration::default(), selector: None, + affinity: None, }}, }, (false, true, true, true) => Role { @@ -751,6 +759,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (false, true, true, false) => Role { @@ -768,6 +777,7 @@ mod tests { None, None), selector: None, + affinity: None, }}, }, (false, true, false, true) => Role { @@ -780,6 +790,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (false, true, false, false) => Role { @@ -792,6 +803,7 @@ mod tests { None, None), selector: None, + affinity: None, }}, }, (false, false, true, true) => Role { @@ -809,6 +821,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (false, false, true, false) => Role { @@ -822,6 +835,7 @@ mod tests { replicas: Some(1), config: CommonConfiguration::default(), selector: None, + affinity: None, }}, }, (false, false, false, true) => Role { @@ -834,6 +848,7 @@ mod tests { build_env_override(GROUP_ENV_OVERRIDE), build_cli_override(GROUP_CLI_OVERRIDE)), selector: None, + affinity: None, }}, }, (false, false, false, false) => Role { @@ -842,6 +857,7 @@ mod tests { replicas: Some(1), config: CommonConfiguration::default(), selector: None, + affinity: None, }}, }, } @@ -1045,7 +1061,8 @@ mod tests { build_config_override(file_name, "file"), build_env_override(GROUP_ENV_OVERRIDE), None), - selector: None, + selector: None, + affinity: None, }}, }; @@ -1105,6 +1122,7 @@ mod tests { None ), selector: None, + affinity: None, }, role_group_2.to_string() => RoleGroup { replicas: Some(1), @@ -1115,6 +1133,7 @@ mod tests { None ), selector: None, + affinity: None, }} }), role_2.to_string() => (vec![PropertyNameKind::Cli], Role { @@ -1133,6 +1152,7 @@ mod tests { None ), selector: None, + affinity: None, }}, }) }; @@ -1195,6 +1215,7 @@ mod tests { None ), selector: None, + affinity: None, }} } ), @@ -1210,6 +1231,7 @@ mod tests { None ), selector: None, + affinity: None, }} } ), diff --git a/src/role_utils.rs b/src/role_utils.rs index e6478577e..d527e1b33 100644 --- a/src/role_utils.rs +++ b/src/role_utils.rs @@ -93,6 +93,7 @@ use crate::{ product_config_utils::Configuration, }; use derivative::Derivative; +use k8s_openapi::api::core::v1::Affinity; use k8s_openapi::apimachinery::pkg::apis::meta::v1::LabelSelector; use kube::{runtime::reflector::ObjectRef, Resource}; use schemars::JsonSchema; @@ -165,6 +166,7 @@ impl Role { }, replicas: group.replicas, selector: group.selector, + affinity: group.affinity, }, ) }) @@ -183,6 +185,7 @@ pub struct RoleGroup { pub config: CommonConfiguration, pub replicas: Option, pub selector: Option, + pub affinity: Option, } impl RoleGroup {