horus/benches/memory_usage.rs

use criterion::{black_box, criterion_group, criterion_main, Criterion, BenchmarkId};
use hero_supervisor_openrpc_client::SupervisorClientBuilder;
use hero_job::Job;
use tokio::runtime::Runtime;
use std::time::Duration;
use std::collections::HashMap;
use uuid::Uuid;
use chrono::Utc;

const SUPERVISOR_URL: &str = "http://127.0.0.1:3030";
const ADMIN_SECRET: &str = "SECRET";

fn create_runtime() -> Runtime {
    Runtime::new().unwrap()
}

fn create_test_job(runner: &str, command: &str, args: Vec<String>) -> Job {
    Job {
        id: Uuid::new_v4().to_string(),
        caller_id: "benchmark".to_string(),
        context_id: "test".to_string(),
        payload: serde_json::json!({
            "command": command,
            "args": args
        }).to_string(),
        runner: runner.to_string(),
        timeout: 30,
        env_vars: HashMap::new(),
        created_at: Utc::now(),
        updated_at: Utc::now(),
        signatures: vec![],
    }
}

#[cfg(target_os = "macos")]
fn get_memory_usage() -> Option<usize> {
    use std::process::Command;
    let output = Command::new("ps")
        .args(&["-o", "rss=", "-p", &std::process::id().to_string()])
        .output()
        .ok()?;
    String::from_utf8(output.stdout)
        .ok()?
        .trim()
        .parse::<usize>()
        .ok()
        .map(|kb| kb * 1024)
}

#[cfg(target_os = "linux")]
fn get_memory_usage() -> Option<usize> {
    use std::fs;
    let status = fs::read_to_string("/proc/self/status").ok()?;
    for line in status.lines() {
        if line.starts_with("VmRSS:") {
            let kb = line.split_whitespace().nth(1)?.parse::<usize>().ok()?;
            return Some(kb * 1024);
        }
    }
    None
}

fn memory_job_creation(c: &mut Criterion) {
    let rt = create_runtime();
    let client = rt.block_on(async {
        SupervisorClientBuilder::new()
            .url(SUPERVISOR_URL)
            .secret(ADMIN_SECRET)
            .build()
            .expect("Failed to create client")
    });

    rt.block_on(async {
        let _ = client.runner_create("hero").await;
    });

    let mut group = c.benchmark_group("memory_job_creation");
    
    for num_jobs in [10, 50, 100, 200].iter() {
        group.bench_with_input(
            BenchmarkId::from_parameter(num_jobs),
            num_jobs,
            |b, &num_jobs| {
                b.iter_custom(|iters| {
                    let mut total_duration = Duration::ZERO;
                    
                    for _ in 0..iters {
                        let mem_before = get_memory_usage().unwrap_or(0);
                        
                        let start = std::time::Instant::now();
                        rt.block_on(async {
                            let mut jobs = Vec::new();
                            for i in 0..num_jobs {
                                let job = create_test_job("hero", "echo", vec![format!("mem_test_{}", i)]);
                                jobs.push(job);
                            }
                            black_box(jobs);
                        });
                        total_duration += start.elapsed();
                        
                        let mem_after = get_memory_usage().unwrap_or(0);
                        let mem_delta = mem_after.saturating_sub(mem_before);
                        
                        if mem_delta > 0 {
                            eprintln!("Memory delta for {} jobs: {} KB", num_jobs, mem_delta / 1024);
                        }
                    }
                    
                    total_duration
                });
            },
        );
    }
    
    group.finish();
}

fn memory_client_creation(c: &mut Criterion) {
    let rt = create_runtime();
    
    let mut group = c.benchmark_group("memory_client_creation");
    
    for num_clients in [1, 10, 50, 100].iter() {
        group.bench_with_input(
            BenchmarkId::from_parameter(num_clients),
            num_clients,
            |b, &num_clients| {
                b.iter_custom(|iters| {
                    let mut total_duration = Duration::ZERO;
                    
                    for _ in 0..iters {
                        let mem_before = get_memory_usage().unwrap_or(0);
                        
                        let start = std::time::Instant::now();
                        rt.block_on(async {
                            let mut clients = Vec::new();
                            for _ in 0..num_clients {
                                let client = SupervisorClientBuilder::new()
                                    .url(SUPERVISOR_URL)
                                    .secret(ADMIN_SECRET)
                                    .build()
                                    .expect("Failed to create client");
                                clients.push(client);
                            }
                            black_box(clients);
                        });
                        total_duration += start.elapsed();
                        
                        let mem_after = get_memory_usage().unwrap_or(0);
                        let mem_delta = mem_after.saturating_sub(mem_before);
                        
                        if mem_delta > 0 {
                            eprintln!("Memory delta for {} clients: {} KB", num_clients, mem_delta / 1024);
                        }
                    }
                    
                    total_duration
                });
            },
        );
    }
    
    group.finish();
}

fn memory_payload_sizes(c: &mut Criterion) {
    let mut group = c.benchmark_group("memory_payload_sizes");
    
    for size_kb in [1, 10, 100, 1000].iter() {
        group.bench_with_input(
            BenchmarkId::from_parameter(format!("{}KB", size_kb)),
            size_kb,
            |b, &size_kb| {
                b.iter_custom(|iters| {
                    let mut total_duration = Duration::ZERO;
                    
                    for _ in 0..iters {
                        let mem_before = get_memory_usage().unwrap_or(0);
                        
                        let start = std::time::Instant::now();
                        let large_data = "x".repeat(size_kb * 1024);
                        let job = create_test_job("hero", "echo", vec![large_data]);
                        black_box(job);
                        total_duration += start.elapsed();
                        
                        let mem_after = get_memory_usage().unwrap_or(0);
                        let mem_delta = mem_after.saturating_sub(mem_before);
                        
                        if mem_delta > 0 {
                            eprintln!("Memory delta for {}KB payload: {} KB", size_kb, mem_delta / 1024);
                        }
                    }
                    
                    total_duration
                });
            },
        );
    }
    
    group.finish();
}

criterion_group!(
    memory_benches,
    memory_job_creation,
    memory_client_creation,
    memory_payload_sizes,
);

criterion_main!(memory_benches);