mirror of
https://github.com/starship/starship.git
synced 2024-11-22 08:14:10 +01: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:
parent
c251897ae8
commit
fae92b2964
@ -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 == ¤t_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, ¤t_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(¤t_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()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user