watcheragent/WatcherAgent/src/docker/stats.rs

use bollard::query_parameters::{ListContainersOptions, StatsOptions};
use bollard::Docker;
use futures_util::stream::TryStreamExt;
use serde::{Deserialize, Serialize};
use std::error::Error;

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ContainerCpuInfo {
    pub container_id: String,
    pub cpu_usage_percent: f64,
    pub system_cpu_usage: u64,
    pub container_cpu_usage: u64,
    pub online_cpus: u32,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ContainerNetworkInfo {
    pub container_id: String,
    pub rx_bytes: u64,
    pub tx_bytes: u64,
    pub rx_packets: u64,
    pub tx_packets: u64,
    pub rx_errors: u64,
    pub tx_errors: u64,
}

/// Get container statistics for all containers using an existing Docker client
pub async fn get_container_stats(
    docker: &Docker,
) -> Result<(Vec<ContainerCpuInfo>, Vec<ContainerNetworkInfo>), Box<dyn Error + Send + Sync>> {
    let containers = docker
        .list_containers(Some(ListContainersOptions {
            all: true,
            ..Default::default()
        }))
        .await?;

    let mut cpu_infos = Vec::new();
    let mut net_infos = Vec::new();

    for container in containers {
        let id = container.id.unwrap_or_default();

        // Skip if no ID
        if id.is_empty() {
            continue;
        }

        let mut stats_stream = docker.stats(
            &id,
            Some(StatsOptions {
                stream: false,
                one_shot: true,
            }),
        );

        if let Some(stats) = stats_stream.try_next().await? {
            // CPU Info
            if let (Some(cpu_stats), Some(precpu_stats)) = (&stats.cpu_stats, &stats.precpu_stats) {
                if let (Some(cpu_usage), Some(pre_cpu_usage)) =
                    (&cpu_stats.cpu_usage, &precpu_stats.cpu_usage)
                {
                    let cpu_delta = cpu_usage
                        .total_usage
                        .unwrap_or(0)
                        .saturating_sub(pre_cpu_usage.total_usage.unwrap_or(0));

                    let system_delta = cpu_stats
                        .system_cpu_usage
                        .unwrap_or(0)
                        .saturating_sub(precpu_stats.system_cpu_usage.unwrap_or(0));

                    let online_cpus = cpu_stats.online_cpus.unwrap_or(1);

                    let cpu_percent = if system_delta > 0 && online_cpus > 0 {
                        (cpu_delta as f64 / system_delta as f64) * online_cpus as f64 * 100.0
                    } else {
                        0.0
                    };

                    cpu_infos.push(ContainerCpuInfo {
                        container_id: id.clone(),
                        cpu_usage_percent: cpu_percent,
                        system_cpu_usage: cpu_stats.system_cpu_usage.unwrap_or(0),
                        container_cpu_usage: cpu_usage.total_usage.unwrap_or(0),
                        online_cpus,
                    });
                }
            }

            // Network Info
            if let Some(networks) = stats.networks {
                for (_name, net) in networks {
                    net_infos.push(ContainerNetworkInfo {
                        container_id: id.clone(),
                        rx_bytes: net.rx_bytes.unwrap(),
                        tx_bytes: net.tx_bytes.unwrap(),
                        rx_packets: net.rx_packets.unwrap(),
                        tx_packets: net.tx_packets.unwrap(),
                        rx_errors: net.rx_errors.unwrap(),
                        tx_errors: net.tx_errors.unwrap(),
                    });
                }
            }
        }
    }

    Ok((cpu_infos, net_infos))
}

/// Get container statistics for a specific container
pub async fn get_single_container_stats(
    docker: &Docker,
    container_id: &str,
) -> Result<(Option<ContainerCpuInfo>, Option<ContainerNetworkInfo>), Box<dyn Error + Send + Sync>>
{
    let mut stats_stream = docker.stats(
        container_id,
        Some(StatsOptions {
            stream: false,
            one_shot: true,
        }),
    );

    if let Some(stats) = stats_stream.try_next().await? {
        let mut cpu_info = None;
        let mut net_info = None;

        // CPU Info
        if let (Some(cpu_stats), Some(precpu_stats)) = (&stats.cpu_stats, &stats.precpu_stats) {
            if let (Some(cpu_usage), Some(pre_cpu_usage)) =
                (&cpu_stats.cpu_usage, &precpu_stats.cpu_usage)
            {
                let cpu_delta = cpu_usage
                    .total_usage
                    .unwrap_or(0)
                    .saturating_sub(pre_cpu_usage.total_usage.unwrap_or(0));

                let system_delta = cpu_stats
                    .system_cpu_usage
                    .unwrap_or(0)
                    .saturating_sub(precpu_stats.system_cpu_usage.unwrap_or(0));

                let online_cpus = cpu_stats.online_cpus.unwrap_or(1);

                let cpu_percent = if system_delta > 0 && online_cpus > 0 {
                    (cpu_delta as f64 / system_delta as f64) * online_cpus as f64 * 100.0
                } else {
                    0.0
                };

                cpu_info = Some(ContainerCpuInfo {
                    container_id: container_id.to_string(),
                    cpu_usage_percent: cpu_percent,
                    system_cpu_usage: cpu_stats.system_cpu_usage.unwrap_or(0),
                    container_cpu_usage: cpu_usage.total_usage.unwrap_or(0),
                    online_cpus,
                });
            }
        }

        // Network Info
        if let Some(networks) = stats.networks {
            // Take the first network interface (usually eth0)
            if let Some((_name, net)) = networks.into_iter().next() {
                net_info = Some(ContainerNetworkInfo {
                    container_id: container_id.to_string(),
                    rx_bytes: net.rx_bytes.unwrap(),
                    tx_bytes: net.tx_bytes.unwrap(),
                    rx_packets: net.rx_packets.unwrap(),
                    tx_packets: net.tx_packets.unwrap(),
                    rx_errors: net.rx_errors.unwrap(),
                    tx_errors: net.tx_errors.unwrap(),
                });
            }
        }

        Ok((cpu_info, net_info))
    } else {
        Ok((None, None))
    }
}

/// Get total network statistics across all containers
pub async fn get_total_network_stats(
    docker: &Docker,
) -> Result<(u64, u64), Box<dyn Error + Send + Sync>> {
    let (_, net_infos) = get_container_stats(docker).await?;

    let total_rx: u64 = net_infos.iter().map(|net| net.rx_bytes).sum();
    let total_tx: u64 = net_infos.iter().map(|net| net.tx_bytes).sum();

    Ok((total_rx, total_tx))
}

/// Get average CPU usage across all containers
pub async fn get_average_cpu_usage(docker: &Docker) -> Result<f64, Box<dyn Error + Send + Sync>> {
    let (cpu_infos, _) = get_container_stats(docker).await?;

    if cpu_infos.is_empty() {
        return Ok(0.0);
    }

    let total_cpu: f64 = cpu_infos.iter().map(|cpu| cpu.cpu_usage_percent).sum();
    Ok(total_cpu / cpu_infos.len() as f64)
}