Skip to content

Commit

Permalink
Merge pull request #25 from Angelin01/tolerations
Browse files Browse the repository at this point in the history 
Add support for Tolerations
  • Loading branch information
@Angelin01
Angelin01 authored May 6, 2024
2 parents a311340 + 149775f commit 965a4f0
Show file tree
Hide file tree
Showing 4 changed files with 289 additions and 28 deletions.
34 changes: 29 additions & 5 deletions src/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use anyhow::{Error, Result};
use axum_server::tls_rustls::RustlsConfig;
use figment::{error, Figment, providers::{Env, Format, Yaml}};
use figment::providers::Serialized;
use k8s_openapi::api::core::v1::Toleration;
use serde::{Deserialize, Serialize};

use crate::error::ConfigError;
Expand Down Expand Up @@ -87,7 +88,7 @@ impl Default for ServerConfig {
pub struct GroupConfig {
pub node_selector: Option<HashMap<String, String>>,
pub affinity: Option<Vec<String>>,
pub tolerations: Option<Vec<String>>,
pub tolerations: Option<Vec<Toleration>>,
#[serde(default)]
pub on_conflict: Conflict,
}
Expand All @@ -98,6 +99,7 @@ mod tests {

use figment::Jail;
use indoc::indoc;
use k8s_openapi::api::core::v1::Toleration;

use super::{Config, Conflict, DEFAULT_CONFIG_FILE, ENV_CONFIG_FILE, GroupConfig};

Expand All @@ -112,12 +114,20 @@ mod tests {
b: "2"
c: "3"
bar:
tolerations: ["1", "2"]
tolerations:
- key: foo
operator: Equals
value: bar
effect: NoSchedule
bazz:
affinity: []
all:
nodeSelector: {"a": "1", "b": "2", "c": "3"}
tolerations: ["1", "2"]
tolerations:
- key: foo
operator: Equals
value: bar
effect: NoSchedule
affinity: []
onConflict: Override
"# })?;
Expand All @@ -138,7 +148,14 @@ mod tests {
groups.insert("bar".into(), GroupConfig {
node_selector: None,
affinity: None,
tolerations: Some(vec!["1".into(), "2".into()]),
tolerations: Some(vec![Toleration {
effect: Some("NoSchedule".into()),
key: Some("foo".into()),
operator: Some("Equals".into()),
toleration_seconds: None,
value: Some("bar".into()),
}
]),
on_conflict: Default::default(),
});
groups.insert("bazz".into(), GroupConfig {
Expand All @@ -154,7 +171,14 @@ mod tests {
("c".into(), "3".into()),
])),
affinity: Some(vec![]),
tolerations: Some(vec!["1".into(), "2".into()]),
tolerations: Some(vec![Toleration {
effect: Some("NoSchedule".into()),
key: Some("foo".into()),
operator: Some("Equals".into()),
toleration_seconds: None,
value: Some("bar".into()),
}
]),
on_conflict: Conflict::Override,
});

Expand Down
223 changes: 222 additions & 1 deletion src/handler/mutate.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,16 @@ pub async fn mutate<S: AppState>(
}),
};

let pod_spec = request.object.as_ref()
.expect("Request object is missing")
.spec.as_ref()
.expect("Pod spec is missing");

let mut patches = Vec::new();

if let Some(node_selector_config) = &group_config.node_selector {
let node_selector_patches = patch::calculate_node_selector_patches(
request.object.as_ref().expect("Request object is missing"),
pod_spec,
node_selector_config,
&group_config.on_conflict,
);
Expand All @@ -53,6 +59,11 @@ pub async fn mutate<S: AppState>(
}
}

if let Some(tolerations_config) = &group_config.tolerations {
let toleration_patches = patch::calculate_toleration_patches(pod_spec, tolerations_config);
patches.extend(toleration_patches);
}

Ok(Json(
AdmissionResponse::from(&request)
.with_patch(json_patch::Patch(patches))
Expand All @@ -68,6 +79,7 @@ mod tests {
use axum::body::Body;
use axum::http::{Request, StatusCode};
use axum::response::Response;
use k8s_openapi::api::core::v1::Toleration;
use serde_json::json;
use tower::ServiceExt;

Expand Down Expand Up @@ -413,4 +425,213 @@ mod tests {
let result = ParsedResponse::from_response(response).await;
assert_eq!(result.status, StatusCode::INTERNAL_SERVER_ERROR);
}

#[tokio::test]
async fn when_pod_has_no_tolerations_should_insert_tolerations_and_pd_tolerations() {
let mut config = Config::default();
let group_config = GroupConfig {
node_selector: None,
affinity: None,
tolerations: Some(vec![Toleration {
key: Some("some-key".into()),
value: Some("some-value".into()),
operator: Some("Equals".into()),
effect: Some("NoSchedule".into()),
toleration_seconds: None,
}]),
on_conflict: Conflict::Reject,
};
config.groups = HashMap::from([
("bar".into(), group_config)
]);

let mut state = TestAppState::new(config);
state.kubernetes.set_namespace_group("foo", "bar");

let body = PodCreateRequestBuilder::new()
.with_namespace("foo")
.build();

let response = mutate_request(state, body).await;
let result = ParsedResponse::from_response(response).await;
assert_eq!(result.status, StatusCode::OK);
assert_eq!(result.admission_response.allowed, true);

let expected_patches = vec![
patch::add("/spec/tolerations".into(), json!([])),
patch::add("/spec/tolerations/-".into(), json!({
"key": "some-key",
"value": "some-value",
"operator": "Equals",
"effect": "NoSchedule"
})),
];

assert_eq!(result.patches, expected_patches);
}

#[tokio::test]
async fn when_pod_has_existing_tolerations_not_matching_config_should_only_insert_pd_tolerations() {
let mut config = Config::default();
let group_config = GroupConfig {
node_selector: None,
affinity: None,
tolerations: Some(vec![Toleration {
key: Some("some-key".into()),
value: Some("some-value".into()),
operator: Some("Equals".into()),
effect: Some("NoSchedule".into()),
toleration_seconds: None,
}]),
on_conflict: Conflict::Reject,
};
config.groups = HashMap::from([
("bar".into(), group_config)
]);

let mut state = TestAppState::new(config);
state.kubernetes.set_namespace_group("foo", "bar");

let body = PodCreateRequestBuilder::new()
.with_namespace("foo")
.with_toleration(Toleration {
effect: Some("NoExecute".into()),
key: Some("other".into()),
operator: Some("Exists".into()),
toleration_seconds: None,
value: None,
})
.build();

let response = mutate_request(state, body).await;
let result = ParsedResponse::from_response(response).await;
assert_eq!(result.status, StatusCode::OK);
assert_eq!(result.admission_response.allowed, true);

let expected_patches = vec![
patch::add("/spec/tolerations/-".into(), json!({
"key": "some-key",
"value": "some-value",
"operator": "Equals",
"effect": "NoSchedule"
})),
];

assert_eq!(result.patches, expected_patches);
}

#[tokio::test]
async fn when_pod_has_existing_tolerations_with_some_matching_config_should_only_insert_necessary_tolerations() {
let mut config = Config::default();
let group_config = GroupConfig {
node_selector: None,
affinity: None,
tolerations: Some(vec![
Toleration {
key: Some("some-key".into()),
value: Some("some-value".into()),
operator: Some("Equals".into()),
effect: Some("NoSchedule".into()),
toleration_seconds: None,
},
Toleration {
effect: Some("NoExecute".into()),
key: Some("other".into()),
operator: Some("Exists".into()),
toleration_seconds: None,
value: None,
},
]),
on_conflict: Conflict::Reject,
};
config.groups = HashMap::from([
("bar".into(), group_config)
]);

let mut state = TestAppState::new(config);
state.kubernetes.set_namespace_group("foo", "bar");

let body = PodCreateRequestBuilder::new()
.with_namespace("foo")
.with_toleration(Toleration {
effect: Some("NoExecute".into()),
key: Some("other".into()),
operator: Some("Exists".into()),
toleration_seconds: None,
value: None,
})
.build();

let response = mutate_request(state, body).await;
let result = ParsedResponse::from_response(response).await;
assert_eq!(result.status, StatusCode::OK);
assert_eq!(result.admission_response.allowed, true);

let expected_patches = vec![
patch::add("/spec/tolerations/-".into(), json!({
"key": "some-key",
"value": "some-value",
"operator": "Equals",
"effect": "NoSchedule"
})),
];

assert_eq!(result.patches, expected_patches);
}

#[tokio::test]
async fn when_pod_has_existing_tolerations_with_perfect_matching_config_should_should_do_nothing() {
let mut config = Config::default();
let group_config = GroupConfig {
node_selector: None,
affinity: None,
tolerations: Some(vec![
Toleration {
key: Some("some-key".into()),
value: Some("some-value".into()),
operator: Some("Equals".into()),
effect: Some("NoSchedule".into()),
toleration_seconds: None,
},
Toleration {
effect: Some("NoExecute".into()),
key: Some("other".into()),
operator: Some("Exists".into()),
toleration_seconds: None,
value: None,
},
]),
on_conflict: Conflict::Reject,
};
config.groups = HashMap::from([
("bar".into(), group_config)
]);

let mut state = TestAppState::new(config);
state.kubernetes.set_namespace_group("foo", "bar");

let body = PodCreateRequestBuilder::new()
.with_namespace("foo")
.with_toleration(Toleration {
effect: Some("NoExecute".into()),
key: Some("other".into()),
operator: Some("Exists".into()),
toleration_seconds: None,
value: None,
}).
with_toleration(Toleration {
key: Some("some-key".into()),
value: Some("some-value".into()),
operator: Some("Equals".into()),
effect: Some("NoSchedule".into()),
toleration_seconds: None,
})
.build();

let response = mutate_request(state, body).await;
let result = ParsedResponse::from_response(response).await;
assert_eq!(result.status, StatusCode::OK);
assert_eq!(result.admission_response.allowed, true);
assert!(result.patches.is_empty());
}
}
25 changes: 10 additions & 15 deletions src/test_utils/pod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,17 @@
use std::collections::BTreeMap;
use axum::body::Body;
use k8s_openapi::api::core::v1::Toleration;
use serde_json::json;

pub struct PodCreateRequestBuilder {
namespace: Option<String>,
node_selector: Option<BTreeMap<String, String>>,
tolerations: Option<Vec<Toleration>>
}

impl PodCreateRequestBuilder {
pub fn new() -> Self {
Self { namespace: None, node_selector: None }
Self { namespace: None, node_selector: None, tolerations: None }
}

pub fn with_namespace<S: AsRef<str>>(mut self, namespace: S) -> Self {
Expand All @@ -23,6 +25,12 @@ impl PodCreateRequestBuilder {
self
}

pub fn with_toleration(mut self, toleration: Toleration) -> Self {
self.tolerations.get_or_insert_with(Vec::new)
.push(toleration);
self
}

pub fn build(self) -> Body {
let data = json!({
"apiVersion": "admission.k8s.io/v1",
Expand Down Expand Up @@ -103,20 +111,7 @@ impl PodCreateRequestBuilder {
"serviceAccount": "default",
"serviceAccountName": "default",
"terminationGracePeriodSeconds": 30,
"tolerations": [
{
"effect": "NoExecute",
"key": "node.kubernetes.io/not-ready",
"operator": "Exists",
"tolerationSeconds": 300
},
{
"effect": "NoExecute",
"key": "node.kubernetes.io/unreachable",
"operator": "Exists",
"tolerationSeconds": 300
}
],
"tolerations": self.tolerations,
"volumes": []
},
"status": {}
Expand Down
Loading

0 comments on commit 965a4f0

Please sign in to comment.