From fae92b29646f73bf1eb843fde80885c58480e85b Mon Sep 17 00:00:00 2001 From: Alvaro Aleman Date: Tue, 16 Jul 2024 16:26:03 -0400 Subject: [PATCH] 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 --- src/modules/kubernetes.rs | 264 +++++++++++++++++++++++++++++++------- 1 file changed, 217 insertions(+), 47 deletions(-) diff --git a/src/modules/kubernetes.rs b/src/modules/kubernetes.rs index c0838a3e..6c0e17dc 100644 --- a/src/modules/kubernetes.rs +++ b/src/modules/kubernetes.rs @@ -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::env; -use std::path; use super::{Context, Module, ModuleConfig}; @@ -17,50 +17,39 @@ struct KubeCtxComponents { cluster: Option, } -fn get_current_kube_context_name(filename: path::PathBuf) -> Option { - let contents = utils::read_file(filename).ok()?; - - let yaml_docs = YamlLoader::load_from_str(&contents).ok()?; - let conf = yaml_docs.first()?; - conf["current-context"] - .as_str() +fn get_current_kube_context_name(document: &T) -> Option<&str> { + document + .get("current-context") + .and_then(DataValue::as_str) .filter(|s| !s.is_empty()) - .map(String::from) } -fn get_kube_ctx_components( - filename: path::PathBuf, +fn get_kube_ctx_components( + document: &T, current_ctx_name: &str, ) -> Option { - let contents = utils::read_file(filename).ok()?; - - let yaml_docs = YamlLoader::load_from_str(&contents).ok()?; - 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 + document + .get("contexts")? + .as_array()? .iter() - .filter_map(|ctx| Some((ctx, ctx["name"].as_str()?))) - .find(|(_, name)| name == ¤t_ctx_name)?; - - let ctx_components = KubeCtxComponents { - user: ctx_yaml["context"]["user"] - .as_str() - .filter(|s| !s.is_empty()) - .map(String::from), - namespace: ctx_yaml["context"]["namespace"] - .as_str() - .filter(|s| !s.is_empty()) - .map(String::from), - cluster: ctx_yaml["context"]["cluster"] - .as_str() - .filter(|s| !s.is_empty()) - .map(String::from), - }; - - Some(ctx_components) + .find(|ctx| ctx.get("name").and_then(DataValue::as_str) == Some(current_ctx_name)) + .map(|ctx| KubeCtxComponents { + user: ctx + .get("context") + .and_then(|v| v.get("user")) + .and_then(DataValue::as_str) + .map(String::from), + namespace: ctx + .get("context") + .and_then(|v| v.get("namespace")) + .and_then(DataValue::as_str) + .map(String::from), + cluster: ctx + .get("context") + .and_then(|v| v.get("cluster")) + .and_then(DataValue::as_str) + .map(String::from), + }) } 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>; +} + +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> { + 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> { + match self { + Yaml::Array(arr) => Some(arr.iter().collect()), + _ => None, + } + } +} + pub fn module<'a>(context: &'a Context) -> Option> { let mut module = context.new_module("kubernetes"); let config: KubernetesConfig = KubernetesConfig::try_load(module.config); @@ -140,8 +175,13 @@ pub fn module<'a>(context: &'a Context) -> Option> { .get_env("KUBECONFIG") .unwrap_or(default_config_file.to_str()?.to_string()); - let current_kube_ctx_name = - env::split_paths(&kube_cfg).find_map(get_current_kube_context_name)?; + let raw_kubeconfigs = env::split_paths(&kube_cfg).map(|file| utils::read_file(file).ok()); + 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 // https://kubernetes.io/docs/concepts/configuration/organize-cluster-access-kubeconfig/ @@ -149,9 +189,10 @@ pub fn module<'a>(context: &'a Context) -> Option> { // > use only values from the first file's red-user. Even if the second file has // > non-conflicting entries under red-user, discard them. // for that reason, we can pick the first context with that name - let ctx_components: KubeCtxComponents = env::split_paths(&kube_cfg) - .find_map(|filename| get_kube_ctx_components(filename, ¤t_kube_ctx_name)) - .unwrap_or_else(|| { + let ctx_components: KubeCtxComponents = kubeconfigs.iter().find_map(|kubeconfig| match kubeconfig { + Document::Json(json) => get_kube_ctx_components(json, current_kube_ctx_name), + 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 log::warn!( "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> { .find_map(|context_config| { let context_alias = get_aliased_name( Some(context_config.context_pattern), - Some(¤t_kube_ctx_name), + Some(current_kube_ctx_name), context_config.context_alias, )?; @@ -185,7 +226,7 @@ pub fn module<'a>(context: &'a Context) -> Option> { 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 let display_context = @@ -239,6 +280,32 @@ pub fn module<'a>(context: &'a Context) -> Option> { Some(module) } +fn parse_kubeconfigs(raw_kubeconfigs: I) -> Vec +where + I: Iterator>, +{ + 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 { + YamlLoader::load_from_str(s) + .ok() + .and_then(|yaml| yaml.into_iter().next().map(Document::Yaml)) +} + mod deprecated { use std::borrow::Cow; use std::collections::HashMap; @@ -282,6 +349,8 @@ mod deprecated { #[cfg(test)] mod tests { + use crate::modules::kubernetes::parse_kubeconfigs; + use crate::modules::kubernetes::Document; use crate::test::ModuleRenderer; use nu_ansi_term::Color; use std::env; @@ -1471,4 +1540,105 @@ users: [] assert_eq!(expected, actual); 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() + } }