1
0
mirror of https://github.com/Llewellynvdm/starship.git synced 2024-12-05 03:13:14 +00:00

perf(k8s): Improve performance of kubeconfig module (#6032)

* Fix config schema

* Improve performance of kubeconfig module

This module currently takes about 200 ms when using our ~10MiB
kubeconfig. This change improves its performance by:
* Only parsing the file once
* (Naively) checking if the content is yaml or json and potentially
  parse as the latter, as that seems to be much faster
This commit is contained in:
Alvaro Aleman 2024-07-16 16:26:03 -04:00 committed by GitHub
parent c251897ae8
commit fae92b2964
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,8 +1,8 @@
use yaml_rust2::YamlLoader; use serde_json::Value as JsonValue;
use yaml_rust2::{Yaml, YamlLoader};
use std::borrow::Cow; use std::borrow::Cow;
use std::env; use std::env;
use std::path;
use super::{Context, Module, ModuleConfig}; use super::{Context, Module, ModuleConfig};
@ -17,50 +17,39 @@ struct KubeCtxComponents {
cluster: Option<String>, cluster: Option<String>,
} }
fn get_current_kube_context_name(filename: path::PathBuf) -> Option<String> { fn get_current_kube_context_name<T: DataValue>(document: &T) -> Option<&str> {
let contents = utils::read_file(filename).ok()?; document
.get("current-context")
let yaml_docs = YamlLoader::load_from_str(&contents).ok()?; .and_then(DataValue::as_str)
let conf = yaml_docs.first()?;
conf["current-context"]
.as_str()
.filter(|s| !s.is_empty()) .filter(|s| !s.is_empty())
.map(String::from)
} }
fn get_kube_ctx_components( fn get_kube_ctx_components<T: DataValue>(
filename: path::PathBuf, document: &T,
current_ctx_name: &str, current_ctx_name: &str,
) -> Option<KubeCtxComponents> { ) -> Option<KubeCtxComponents> {
let contents = utils::read_file(filename).ok()?; document
.get("contexts")?
let yaml_docs = YamlLoader::load_from_str(&contents).ok()?; .as_array()?
let conf = yaml_docs.first()?;
let contexts = conf["contexts"].as_vec()?;
// Find the context with the name we're looking for
// or return None if we can't find it
let (ctx_yaml, _) = contexts
.iter() .iter()
.filter_map(|ctx| Some((ctx, ctx["name"].as_str()?))) .find(|ctx| ctx.get("name").and_then(DataValue::as_str) == Some(current_ctx_name))
.find(|(_, name)| name == &current_ctx_name)?; .map(|ctx| KubeCtxComponents {
user: ctx
let ctx_components = KubeCtxComponents { .get("context")
user: ctx_yaml["context"]["user"] .and_then(|v| v.get("user"))
.as_str() .and_then(DataValue::as_str)
.filter(|s| !s.is_empty())
.map(String::from), .map(String::from),
namespace: ctx_yaml["context"]["namespace"] namespace: ctx
.as_str() .get("context")
.filter(|s| !s.is_empty()) .and_then(|v| v.get("namespace"))
.and_then(DataValue::as_str)
.map(String::from), .map(String::from),
cluster: ctx_yaml["context"]["cluster"] cluster: ctx
.as_str() .get("context")
.filter(|s| !s.is_empty()) .and_then(|v| v.get("cluster"))
.and_then(DataValue::as_str)
.map(String::from), .map(String::from),
}; })
Some(ctx_components)
} }
fn get_aliased_name<'a>( fn get_aliased_name<'a>(
@ -97,6 +86,52 @@ fn get_aliased_name<'a>(
} }
} }
#[derive(Debug)]
enum Document {
Json(JsonValue),
Yaml(Yaml),
}
trait DataValue {
fn get(&self, key: &str) -> Option<&Self>;
fn as_str(&self) -> Option<&str>;
fn as_array(&self) -> Option<Vec<&Self>>;
}
impl DataValue for JsonValue {
fn get(&self, key: &str) -> Option<&Self> {
self.get(key)
}
fn as_str(&self) -> Option<&str> {
self.as_str()
}
fn as_array(&self) -> Option<Vec<&Self>> {
self.as_array().map(|arr| arr.iter().collect())
}
}
impl DataValue for Yaml {
fn get(&self, key: &str) -> Option<&Self> {
match self {
Yaml::Hash(map) => map.get(&Yaml::String(key.to_string())),
_ => None,
}
}
fn as_str(&self) -> Option<&str> {
self.as_str()
}
fn as_array(&self) -> Option<Vec<&Self>> {
match self {
Yaml::Array(arr) => Some(arr.iter().collect()),
_ => None,
}
}
}
pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> { pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
let mut module = context.new_module("kubernetes"); let mut module = context.new_module("kubernetes");
let config: KubernetesConfig = KubernetesConfig::try_load(module.config); let config: KubernetesConfig = KubernetesConfig::try_load(module.config);
@ -140,8 +175,13 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
.get_env("KUBECONFIG") .get_env("KUBECONFIG")
.unwrap_or(default_config_file.to_str()?.to_string()); .unwrap_or(default_config_file.to_str()?.to_string());
let current_kube_ctx_name = let raw_kubeconfigs = env::split_paths(&kube_cfg).map(|file| utils::read_file(file).ok());
env::split_paths(&kube_cfg).find_map(get_current_kube_context_name)?; let kubeconfigs = parse_kubeconfigs(raw_kubeconfigs);
let current_kube_ctx_name = kubeconfigs.iter().find_map(|v| match v {
Document::Json(json) => get_current_kube_context_name(json),
Document::Yaml(yaml) => get_current_kube_context_name(yaml),
})?;
// Even if we have multiple config files, the first key wins // Even if we have multiple config files, the first key wins
// https://kubernetes.io/docs/concepts/configuration/organize-cluster-access-kubeconfig/ // https://kubernetes.io/docs/concepts/configuration/organize-cluster-access-kubeconfig/
@ -149,9 +189,10 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
// > use only values from the first file's red-user. Even if the second file has // > use only values from the first file's red-user. Even if the second file has
// > non-conflicting entries under red-user, discard them. // > non-conflicting entries under red-user, discard them.
// for that reason, we can pick the first context with that name // for that reason, we can pick the first context with that name
let ctx_components: KubeCtxComponents = env::split_paths(&kube_cfg) let ctx_components: KubeCtxComponents = kubeconfigs.iter().find_map(|kubeconfig| match kubeconfig {
.find_map(|filename| get_kube_ctx_components(filename, &current_kube_ctx_name)) Document::Json(json) => get_kube_ctx_components(json, current_kube_ctx_name),
.unwrap_or_else(|| { Document::Yaml(yaml) => get_kube_ctx_components(yaml, current_kube_ctx_name),
}).unwrap_or_else(|| {
// TODO: figure out if returning is more sensible. But currently we have tests depending on this // TODO: figure out if returning is more sensible. But currently we have tests depending on this
log::warn!( log::warn!(
"Invalid KUBECONFIG: identified current-context `{}`, but couldn't find the context in any config file(s): `{}`.\n", "Invalid KUBECONFIG: identified current-context `{}`, but couldn't find the context in any config file(s): `{}`.\n",
@ -169,7 +210,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
.find_map(|context_config| { .find_map(|context_config| {
let context_alias = get_aliased_name( let context_alias = get_aliased_name(
Some(context_config.context_pattern), Some(context_config.context_pattern),
Some(&current_kube_ctx_name), Some(current_kube_ctx_name),
context_config.context_alias, context_config.context_alias,
)?; )?;
@ -185,7 +226,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
Some((Some(context_config), context_alias, user_alias)) Some((Some(context_config), context_alias, user_alias))
}) })
.unwrap_or_else(|| (None, current_kube_ctx_name.clone(), ctx_components.user)); .unwrap_or_else(|| (None, current_kube_ctx_name.to_string(), ctx_components.user));
// TODO: remove deprecated aliases after starship 2.0 // TODO: remove deprecated aliases after starship 2.0
let display_context = let display_context =
@ -239,6 +280,32 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
Some(module) Some(module)
} }
fn parse_kubeconfigs<I>(raw_kubeconfigs: I) -> Vec<Document>
where
I: Iterator<Item = Option<String>>,
{
raw_kubeconfigs
.filter_map(|content| match content {
Some(value) => match value.chars().next() {
// Parsing as json is about an order of magnitude faster than parsing
// as yaml, so do that if possible.
Some('{') => match serde_json::from_str(&value) {
Ok(json) => Some(Document::Json(json)),
Err(_) => parse_yaml(&value),
},
_ => parse_yaml(&value),
},
_ => None,
})
.collect()
}
fn parse_yaml(s: &str) -> Option<Document> {
YamlLoader::load_from_str(s)
.ok()
.and_then(|yaml| yaml.into_iter().next().map(Document::Yaml))
}
mod deprecated { mod deprecated {
use std::borrow::Cow; use std::borrow::Cow;
use std::collections::HashMap; use std::collections::HashMap;
@ -282,6 +349,8 @@ mod deprecated {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use crate::modules::kubernetes::parse_kubeconfigs;
use crate::modules::kubernetes::Document;
use crate::test::ModuleRenderer; use crate::test::ModuleRenderer;
use nu_ansi_term::Color; use nu_ansi_term::Color;
use std::env; use std::env;
@ -1471,4 +1540,105 @@ users: []
assert_eq!(expected, actual); assert_eq!(expected, actual);
dir.close() dir.close()
} }
#[test]
fn test_json_kubeconfig_is_parsed_as_json() -> std::io::Result<()> {
let json_kubeconfig = r#"{
"apiVersion": "v1",
"clusters": [],
"contexts": [
{
"context": {
"user": "test_user",
"namespace": "test_namespace"
},
"name": "test_context"
}
],
"current-context": "test_context",
"kind": "Config",
"preferences": {},
"users": []
}"#
.to_string();
let kubeconfigs = [Some(json_kubeconfig)];
let results = parse_kubeconfigs(kubeconfigs.iter().cloned());
let actual = results.first().unwrap();
match actual {
Document::Json(..) => {}
_ => panic!("Expected Document::Json, got {:?}", actual),
}
Ok(())
}
#[test]
fn fallback_to_yaml_parsing() -> std::io::Result<()> {
let json_kubeconfig = r#"{
"apiVersion": v1,
"clusters": [],
"contexts": [
{
"context": {
"user": test_user,
"namespace": test_namespace
},
"name": test_context
}
],
"current-context": test_context,
"kind": Config,
"preferences": {},
"users": []
}"#
.to_string();
let kubeconfigs = [Some(json_kubeconfig)];
let results = parse_kubeconfigs(kubeconfigs.iter().cloned());
let actual = results.first().unwrap();
match actual {
Document::Yaml(..) => {}
_ => panic!("Expected Document::Yaml, got {:?}", actual),
}
Ok(())
}
#[test]
fn test_parse_json_kubeconfig() -> std::io::Result<()> {
let dir = tempfile::tempdir()?;
let filename = dir.path().join("config");
let mut file = File::create(&filename)?;
file.write_all(
br#"{
"contexts": [
{
"name": "test_context",
"context": {
"user": "test_user",
"namespace": "test_namespace"
}
}
],
"current-context": "test_context",
"kind": "Config",
"apiVersion": "v1"
}
"#,
)?;
file.sync_all()?;
let actual = ModuleRenderer::new("kubernetes")
.path(dir.path())
.env("KUBECONFIG", filename.to_string_lossy().as_ref())
.config(toml::toml! {
[kubernetes]
disabled = false
format = "($user )($context )($cluster )($namespace)"
})
.collect();
let expected = Some("test_user test_context test_namespace".to_string());
assert_eq!(expected, actual);
dir.close()
}
} }