harmony/harmony-k8s/src/pod.rs

use std::time::Duration;

use k8s_openapi::api::core::v1::Pod;
use kube::{
    Error,
    api::{Api, AttachParams, ListParams},
    error::DiscoveryError,
    runtime::reflector::Lookup,
};
use log::debug;
use tokio::io::AsyncReadExt;
use tokio::time::sleep;

use crate::client::K8sClient;

impl K8sClient {
    pub async fn get_pod(&self, name: &str, namespace: Option<&str>) -> Result<Option<Pod>, Error> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };
        api.get_opt(name).await
    }

    pub async fn wait_for_pod_ready(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
    ) -> Result<(), Error> {
        let mut elapsed = 0u64;
        let interval = 5u64;
        let timeout_secs = 120u64;
        loop {
            if let Some(p) = self.get_pod(pod_name, namespace).await? {
                if let Some(phase) = p.status.and_then(|s| s.phase) {
                    if phase.to_lowercase() == "running" {
                        return Ok(());
                    }
                }
            }
            if elapsed >= timeout_secs {
                return Err(Error::Discovery(DiscoveryError::MissingResource(format!(
                    "Pod '{}' in '{}' did not become ready within {timeout_secs}s",
                    pod_name,
                    namespace.unwrap_or("<default>"),
                ))));
            }
            sleep(Duration::from_secs(interval)).await;
            elapsed += interval;
        }
    }

    /// Polls a pod until it reaches `Succeeded` or `Failed`, then returns its
    /// logs.  Used internally by node operations.
    pub(crate) async fn wait_for_pod_completion(
        &self,
        name: &str,
        namespace: &str,
    ) -> Result<String, Error> {
        let api: Api<Pod> = Api::namespaced(self.client.clone(), namespace);
        let poll_interval = Duration::from_secs(2);
        for _ in 0..60 {
            sleep(poll_interval).await;
            let p = api.get(name).await?;
            match p.status.and_then(|s| s.phase).as_deref() {
                Some("Succeeded") => {
                    let logs = api
                        .logs(name, &Default::default())
                        .await
                        .unwrap_or_default();
                    debug!("Pod {namespace}/{name} succeeded. Logs: {logs}");
                    return Ok(logs);
                }
                Some("Failed") => {
                    let logs = api
                        .logs(name, &Default::default())
                        .await
                        .unwrap_or_default();
                    debug!("Pod {namespace}/{name} failed. Logs: {logs}");
                    return Err(Error::Discovery(DiscoveryError::MissingResource(format!(
                        "Pod '{name}' failed.\n{logs}"
                    ))));
                }
                _ => {}
            }
        }
        Err(Error::Discovery(DiscoveryError::MissingResource(format!(
            "Timed out waiting for pod '{name}'"
        ))))
    }

    /// Execute a command in the first pod matching `{label}={name}`.
    pub async fn exec_app_capture_output(
        &self,
        name: String,
        label: String,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<String, String> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };
        let pod_list = api
            .list(&ListParams::default().labels(&format!("{label}={name}")))
            .await
            .expect("Failed to list pods");

        let pod_name = pod_list
            .items
            .first()
            .expect("No matching pod")
            .name()
            .expect("Pod has no name")
            .into_owned();

        match api
            .exec(
                &pod_name,
                command,
                &AttachParams::default().stdout(true).stderr(true),
            )
            .await
        {
            Err(e) => Err(e.to_string()),
            Ok(mut process) => {
                let status = process
                    .take_status()
                    .expect("No status handle")
                    .await
                    .expect("Status channel closed");

                if let Some(s) = status.status {
                    let mut buf = String::new();
                    if let Some(mut stdout) = process.stdout() {
                        stdout
                            .read_to_string(&mut buf)
                            .await
                            .map_err(|e| format!("Failed to read stdout: {e}"))?;
                    }
                    debug!("exec status: {} - {:?}", s, status.details);
                    if s == "Success" { Ok(buf) } else { Err(s) }
                } else {
                    Err("No inner status from pod exec".to_string())
                }
            }
        }
    }

    /// Execute a command in the first pod matching
    /// `app.kubernetes.io/name={name}`.
    pub async fn exec_app(
        &self,
        name: String,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<(), String> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };
        let pod_list = api
            .list(&ListParams::default().labels(&format!("app.kubernetes.io/name={name}")))
            .await
            .expect("Failed to list pods");

        let pod_name = pod_list
            .items
            .first()
            .expect("No matching pod")
            .name()
            .expect("Pod has no name")
            .into_owned();

        match api.exec(&pod_name, command, &AttachParams::default()).await {
            Err(e) => Err(e.to_string()),
            Ok(mut process) => {
                let status = process
                    .take_status()
                    .expect("No status handle")
                    .await
                    .expect("Status channel closed");

                if let Some(s) = status.status {
                    debug!("exec status: {} - {:?}", s, status.details);
                    if s == "Success" { Ok(()) } else { Err(s) }
                } else {
                    Err("No inner status from pod exec".to_string())
                }
            }
        }
    }

    /// Execute a command in a specific pod by name, capturing stdout.
    ///
    /// Returns the captured stdout on success. On failure, the error string
    /// includes stderr output from the remote command.
    pub async fn exec_pod_capture_output(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<String, String> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };

        match api
            .exec(
                pod_name,
                command,
                &AttachParams::default().stdout(true).stderr(true),
            )
            .await
        {
            Err(e) => Err(e.to_string()),
            Ok(mut process) => {
                let status = process
                    .take_status()
                    .expect("No status handle")
                    .await
                    .expect("Status channel closed");

                let mut stdout_buf = String::new();
                if let Some(mut stdout) = process.stdout() {
                    stdout
                        .read_to_string(&mut stdout_buf)
                        .await
                        .map_err(|e| format!("Failed to read stdout: {e}"))?;
                }

                let mut stderr_buf = String::new();
                if let Some(mut stderr) = process.stderr() {
                    stderr
                        .read_to_string(&mut stderr_buf)
                        .await
                        .map_err(|e| format!("Failed to read stderr: {e}"))?;
                }

                if let Some(s) = status.status {
                    debug!("exec_pod status: {} - {:?}", s, status.details);
                    if s == "Success" {
                        Ok(stdout_buf)
                    } else {
                        Err(format!("{stderr_buf}"))
                    }
                } else {
                    Err("No inner status from pod exec".to_string())
                }
            }
        }
    }

    /// Execute a command in a specific pod by name (no output capture).
    pub async fn exec_pod(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<(), String> {
        self.exec_pod_capture_output(pod_name, namespace, command)
            .await
            .map(|_| ())
    }
}