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

@@ -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<String>,
 }
 
-fn get_current_kube_context_name(filename: path::PathBuf) -> Option<String> {
-    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<T: DataValue>(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<T: DataValue>(
+    document: &T,
     current_ctx_name: &str,
 ) -> Option<KubeCtxComponents> {
-    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 == &current_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<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>> {
     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<Module<'a>> {
         .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<Module<'a>> {
     // > 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, &current_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<Module<'a>> {
         .find_map(|context_config| {
             let context_alias = get_aliased_name(
                 Some(context_config.context_pattern),
-                Some(&current_kube_ctx_name),
+                Some(current_kube_ctx_name),
                 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))
         })
-        .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<Module<'a>> {
     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 {
     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()
+    }
 }