mockforge_bench/

chunked_command.rs

//! `mockforge bench-chunked` driver: single-target, spec-driven, and
//! multi-target (`--targets-file`) chunked-encoding runs.
//!
//! [`crate::chunked_bench::run`] drives one URL. This module layers the CLI
//! modes on top of it:
//!
//! - **Single target** (`--target URL`, no `--spec`): one run against the URL.
//! - **Spec-driven** (`--spec`): one run per POST/PUT/PATCH operation, with
//!   `--target` as the base URL.
//! - **Multi-target** (`--targets-file`, issue #79): every target in the file
//!   runs the single-target or spec-driven flow in parallel (bounded by
//!   `--max-concurrency`). It uses the same file format as
//!   `mockforge bench --targets-file`, including per-target `auth`, `headers`
//!   and `spec`. Per-target artifacts go to `<output>/target_<N>/`, and a
//!   roll-up goes to `<output>/chunked-multi-target-summary.json`.

use crate::chunked_bench::{run, ChunkedBenchConfig, ChunkedBenchResult};
use crate::request_gen::RequestGenerator;
use crate::spec_parser::SpecParser;
use crate::target_parser::{parse_targets_file, TargetConfig};
use anyhow::{bail, Context};
use futures::StreamExt;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::Duration;

/// Options shared by every `bench-chunked` mode.
#[derive(Debug, Clone)]
pub struct ChunkedCommand {
    /// Single target URL. Mutually exclusive with `targets_file`.
    pub target: Option<String>,
    /// File with multiple targets (same format as `bench --targets-file`).
    pub targets_file: Option<PathBuf>,
    /// Max targets running at once in multi-target mode.
    pub max_concurrency: u32,
    pub spec: Option<PathBuf>,
    pub base_path: Option<String>,
    pub operation_id: Option<String>,
    pub method: String,
    pub concurrency: u32,
    pub duration: Duration,
    pub chunk_size_bytes: usize,
    pub total_size_bytes: usize,
    pub chunk_interval_ms: u64,
    pub headers: HashMap<String, String>,
    pub insecure: bool,
    pub validate_requests: bool,
    pub export_requests: bool,
    pub output: PathBuf,
}

/// One target's resolved run plan.
#[derive(Debug, Clone)]
struct TargetRun {
    /// `""` in single-target mode, `"[target_N] "` in multi-target mode.
    prefix: String,
    url: String,
    headers: HashMap<String, String>,
    spec: Option<PathBuf>,
    output: PathBuf,
}

/// Roll-up of one target's run(s).
#[derive(Debug, Clone, Default)]
struct TargetOutcome {
    total_requests: u64,
    successful: u64,
    failed: u64,
    bytes_sent: u64,
    status_counts: HashMap<u16, u64>,
    /// An operation errored out entirely (vs. individual requests failing).
    run_failed: bool,
}

impl TargetOutcome {
    fn absorb(&mut self, r: &ChunkedBenchResult) {
        self.total_requests += r.total_requests;
        self.successful += r.successful;
        self.failed += r.failed;
        self.bytes_sent += r.bytes_sent;
        for (code, n) in &r.status_counts {
            *self.status_counts.entry(*code).or_insert(0) += n;
        }
    }
}

impl ChunkedCommand {
    /// Run the bench. Returns `Ok(true)` when every run completed, `Ok(false)`
    /// when at least one operation or target failed to run.
    pub async fn execute(&self) -> anyhow::Result<bool> {
        match (&self.target, &self.targets_file) {
            (Some(_), Some(_)) => bail!("--target and --targets-file are mutually exclusive"),
            (None, None) => bail!("either --target or --targets-file is required"),
            (Some(target), None) => {
                let run = TargetRun {
                    prefix: String::new(),
                    url: target.clone(),
                    headers: self.headers.clone(),
                    spec: self.spec.clone(),
                    output: self.output.clone(),
                };
                Ok(!self.run_target(&run).await?.run_failed)
            }
            (None, Some(file)) => self.execute_multi_target(file).await,
        }
    }

    async fn execute_multi_target(&self, file: &Path) -> anyhow::Result<bool> {
        let targets = parse_targets_file(file)
            .map_err(|e| anyhow::anyhow!("{e}"))
            .with_context(|| format!("failed to read targets file {}", file.display()))?;
        if targets.is_empty() {
            bail!("no targets found in {}", file.display());
        }
        let runs = plan_target_runs(targets, &self.headers, self.spec.as_ref(), &self.output);
        let max_concurrency = self.max_concurrency.max(1) as usize;
        println!(
            "→ Chunked bench across {} targets from {:?} (max {} in parallel)",
            runs.len(),
            file,
            max_concurrency
        );

        let mut outcomes: Vec<(usize, TargetRun, anyhow::Result<TargetOutcome>)> =
            futures::stream::iter(runs.into_iter().enumerate())
                .map(|(i, run)| async move {
                    let res = self.run_target(&run).await;
                    (i, run, res)
                })
                .buffer_unordered(max_concurrency)
                .collect()
                .await;
        outcomes.sort_by_key(|(i, _, _)| *i);

        println!();
        println!("Multi-target chunked summary:");
        let mut all_ok = true;
        let mut rows = Vec::with_capacity(outcomes.len());
        for (i, run, res) in &outcomes {
            let name = format!("target_{}", i + 1);
            match res {
                Ok(o) => {
                    if o.run_failed {
                        all_ok = false;
                    }
                    println!(
                        "  {:<10} {}  requests={} ok={} failed={} bytes={}{}",
                        name,
                        run.url,
                        o.total_requests,
                        o.successful,
                        o.failed,
                        o.bytes_sent,
                        if o.run_failed {
                            "  (some operations failed to run)"
                        } else {
                            ""
                        }
                    );
                    rows.push(serde_json::json!({
                        "target": name,
                        "url": run.url,
                        "spec": run.spec.as_ref().map(|p| p.display().to_string()),
                        // Per-target files (exports / violations) exist only in spec mode.
                        "output_dir": run.spec.as_ref().map(|_| run.output.display().to_string()),
                        "total_requests": o.total_requests,
                        "successful": o.successful,
                        "failed": o.failed,
                        "bytes_sent": o.bytes_sent,
                        "status_counts": o.status_counts,
                        "run_failed": o.run_failed,
                    }));
                }
                Err(e) => {
                    all_ok = false;
                    println!("  {:<10} {}  ERROR: {:#}", name, run.url, e);
                    rows.push(serde_json::json!({
                        "target": name,
                        "url": run.url,
                        "error": format!("{e:#}"),
                    }));
                }
            }
        }

        std::fs::create_dir_all(&self.output)
            .with_context(|| format!("failed to create output directory {:?}", self.output))?;
        let path = self.output.join("chunked-multi-target-summary.json");
        let payload = serde_json::json!({
            "targets_file": file.display().to_string(),
            "concurrency_per_target": self.concurrency,
            "max_concurrency": max_concurrency,
            "duration_secs": self.duration.as_secs(),
            "chunk_size_bytes": self.chunk_size_bytes,
            "total_size_bytes": self.total_size_bytes,
            "chunk_interval_ms": self.chunk_interval_ms,
            "targets": rows,
        });
        std::fs::write(&path, serde_json::to_string_pretty(&payload)?)
            .with_context(|| format!("failed to write {:?}", path))?;
        println!("📝 Wrote {:?}", path);
        Ok(all_ok)
    }

    fn bench_config(
        &self,
        url: String,
        method: reqwest::Method,
        headers: &HashMap<String, String>,
    ) -> ChunkedBenchConfig {
        ChunkedBenchConfig {
            target_url: url,
            method,
            concurrency: self.concurrency,
            duration: self.duration,
            chunk_size_bytes: self.chunk_size_bytes,
            total_size_bytes: self.total_size_bytes,
            chunk_interval_ms: self.chunk_interval_ms,
            headers: headers.clone(),
            skip_tls_verify: self.insecure,
        }
    }

    async fn run_target(&self, t: &TargetRun) -> anyhow::Result<TargetOutcome> {
        match &t.spec {
            Some(spec) => self.run_spec(t, spec).await,
            None => self.run_single(t).await,
        }
    }

    async fn run_single(&self, t: &TargetRun) -> anyhow::Result<TargetOutcome> {
        let p = &t.prefix;
        // Only warn once: in multi-target mode every target would repeat it.
        if p.is_empty() {
            if self.base_path.is_some() {
                eprintln!(
                    "Note: --base-path has no effect without --spec; \
                     single-target mode uses --target verbatim"
                );
            }
            if self.validate_requests {
                eprintln!("Note: --validate-requests requires --spec (skipped)");
            }
            if self.export_requests {
                eprintln!("Note: --export-requests requires --spec (skipped)");
            }
        }
        let method = reqwest::Method::from_bytes(self.method.to_uppercase().as_bytes())
            .map_err(|_| anyhow::anyhow!("Invalid HTTP method: {}", self.method))?;
        if !p.is_empty() {
            println!("{p}→ {} {}", self.method.to_uppercase(), t.url);
        }
        let r = run(self.bench_config(t.url.clone(), method, &t.headers))
            .await
            .context("chunked bench failed")?;
        print_result(
            p,
            if p.is_empty() {
                "single-target"
            } else {
                &t.url
            },
            &r,
        );
        let mut outcome = TargetOutcome::default();
        outcome.absorb(&r);
        Ok(outcome)
    }

    async fn run_spec(&self, t: &TargetRun, spec_path: &Path) -> anyhow::Result<TargetOutcome> {
        let p = &t.prefix;
        let parser = SpecParser::from_file(spec_path)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to parse spec {:?}: {}", spec_path, e))?;
        let base_url = t.url.trim_end_matches('/').to_string();
        // CLI --base-path > spec.servers > none. Empty string from CLI
        // explicitly disables (matches `mockforge bench` semantics).
        let effective_base_path: Option<String> = match &self.base_path {
            Some(bp) if bp.is_empty() => None,
            Some(bp) => Some(bp.clone()),
            None => parser.get_base_path(),
        };
        let want_methods = ["POST", "PUT", "PATCH"];
        let ops: Vec<_> = parser
            .get_operations()
            .into_iter()
            .filter(|op| want_methods.contains(&op.method.to_uppercase().as_str()))
            .filter(|op| match &self.operation_id {
                Some(id) => op.operation_id.as_deref() == Some(id.as_str()),
                None => true,
            })
            .collect();
        if ops.is_empty() {
            bail!(
                "No POST/PUT/PATCH operations matched in {:?}{}",
                spec_path,
                self.operation_id
                    .as_deref()
                    .map(|id| format!(" (operation-id={id})"))
                    .unwrap_or_default()
            );
        }

        // Pre-flight validation. Any failure here aborts before touching
        // the network so the bench doesn't half-run.
        let mut violations: Vec<serde_json::Value> = Vec::new();
        // (label, method, url, op.path)
        let mut planned: Vec<(String, String, String, String)> = Vec::with_capacity(ops.len());
        for op in &ops {
            let label = op.display_name();
            match RequestGenerator::generate_template(op) {
                Ok(template) => {
                    let path_with_base = match &effective_base_path {
                        Some(bp) if !bp.is_empty() => {
                            format!("{}{}", bp, template.generate_path())
                        }
                        _ => template.generate_path().to_string(),
                    };
                    let url = format!("{}{}", base_url, path_with_base);
                    if reqwest::Method::from_bytes(op.method.to_uppercase().as_bytes()).is_err() {
                        violations.push(serde_json::json!({
                            "operation": label,
                            "method": op.method,
                            "path": op.path,
                            "kind": "invalid_method",
                            "detail": format!("`{}` is not a valid HTTP method", op.method),
                        }));
                    } else {
                        planned.push((label, op.method.to_uppercase(), url, op.path.clone()));
                    }
                }
                Err(e) => {
                    violations.push(serde_json::json!({
                        "operation": label,
                        "method": op.method,
                        "path": op.path,
                        "kind": "template_build_failure",
                        "detail": e.to_string(),
                    }));
                }
            }
        }

        if self.validate_requests {
            if !violations.is_empty() {
                std::fs::create_dir_all(&t.output)
                    .with_context(|| format!("Failed to create output directory {:?}", t.output))?;
                let path = t.output.join("chunked-request-violations.json");
                let payload = serde_json::json!({
                    "spec": spec_path.display().to_string(),
                    "base_url": base_url,
                    "base_path": effective_base_path,
                    "violations": violations,
                });
                std::fs::write(&path, serde_json::to_string_pretty(&payload)?)
                    .with_context(|| format!("Failed to write {:?}", path))?;
                bail!("--validate-requests: {} violation(s); see {:?}", violations.len(), path);
            }
            println!("{p}✅ --validate-requests: all {} operations OK", planned.len());
        }

        println!(
            "{p}→ {} chunked bench {} from {:?} (base {}{})",
            planned.len(),
            if planned.len() == 1 {
                "operation"
            } else {
                "operations"
            },
            spec_path,
            base_url,
            effective_base_path
                .as_deref()
                .map(|bp| format!(", base-path {}", bp))
                .unwrap_or_default()
        );

        let mut outcome = TargetOutcome::default();
        let mut export_records: Vec<serde_json::Value> = Vec::new();
        for (label, method_str, url, op_path) in &planned {
            let method = match reqwest::Method::from_bytes(method_str.as_bytes()) {
                Ok(m) => m,
                Err(_) => {
                    eprintln!("{p}✗ {}: invalid method `{}`", label, method_str);
                    outcome.run_failed = true;
                    continue;
                }
            };
            println!();
            println!("{p}→ {} {}  ({})", method_str, op_path, url);
            match run(self.bench_config(url.clone(), method, &t.headers)).await {
                Ok(r) => {
                    print_result(p, label, &r);
                    outcome.absorb(&r);
                    if self.export_requests {
                        export_records.push(export_record(
                            self, label, method_str, url, op_path, &t.headers, &r,
                        ));
                    }
                }
                Err(e) => {
                    eprintln!("{p}✗ {}: chunked bench failed: {}", label, e);
                    outcome.run_failed = true;
                }
            }
        }

        if self.export_requests {
            std::fs::create_dir_all(&t.output)
                .with_context(|| format!("Failed to create output directory {:?}", t.output))?;
            let path = t.output.join("chunked-requests.json");
            let payload = serde_json::json!({
                "spec": spec_path.display().to_string(),
                "base_url": base_url,
                "base_path": effective_base_path,
                "operations": export_records,
            });
            std::fs::write(&path, serde_json::to_string_pretty(&payload)?)
                .with_context(|| format!("Failed to write {:?}", path))?;
            println!(
                "{p}📝 --export-requests: wrote {} operations to {:?}",
                export_records.len(),
                path
            );
        }

        Ok(outcome)
    }
}

/// Resolve each targets-file entry into a run plan: normalized URL, CLI
/// headers overlaid with the target's `headers` and `auth`, the target's
/// `spec` falling back to the CLI `--spec`, and `<output>/target_<N>/`.
fn plan_target_runs(
    targets: Vec<TargetConfig>,
    cli_headers: &HashMap<String, String>,
    cli_spec: Option<&PathBuf>,
    output: &Path,
) -> Vec<TargetRun> {
    targets
        .into_iter()
        .enumerate()
        .map(|(i, mut t)| {
            t.normalize_url();
            let mut headers = cli_headers.clone();
            if let Some(h) = &t.headers {
                headers.extend(h.iter().map(|(k, v)| (k.clone(), v.clone())));
            }
            if let Some(auth) = &t.auth {
                headers.insert("Authorization".to_string(), auth.clone());
            }
            let name = format!("target_{}", i + 1);
            TargetRun {
                prefix: format!("[{name}] "),
                url: t.url,
                headers,
                spec: t.spec.or_else(|| cli_spec.cloned()),
                output: output.join(name),
            }
        })
        .collect()
}

fn export_record(
    cmd: &ChunkedCommand,
    label: &str,
    method: &str,
    url: &str,
    op_path: &str,
    headers: &HashMap<String, String>,
    r: &ChunkedBenchResult,
) -> serde_json::Value {
    serde_json::json!({
        "operation": label,
        "method": method,
        "url": url,
        "spec_path": op_path,
        "headers": headers,
        "chunk_size_bytes": cmd.chunk_size_bytes,
        "total_size_bytes": cmd.total_size_bytes,
        "chunk_interval_ms": cmd.chunk_interval_ms,
        "concurrency": cmd.concurrency,
        "duration_secs": cmd.duration.as_secs(),
        "result": {
            "total_requests": r.total_requests,
            "successful": r.successful,
            "failed": r.failed,
            "bytes_sent": r.bytes_sent,
            "elapsed_ms": r.elapsed.as_millis(),
            "req_per_sec": r.req_per_sec,
            "avg_latency_ms": r.avg_latency_ms,
            "p50_ms": r.p50_ms,
            "p95_ms": r.p95_ms,
            "p99_ms": r.p99_ms,
            "status_counts": r.status_counts,
            "error_samples": r.error_samples.iter().map(|s| serde_json::json!({
                "status": s.status,
                "server_header": s.server_header,
                "body_excerpt": s.body_excerpt,
            })).collect::<Vec<_>>(),
        },
    })
}

/// Print one result block. The block is built first and printed with one
/// call so parallel targets don't interleave their lines.
fn print_result(prefix: &str, label: &str, r: &ChunkedBenchResult) {
    let mut out = String::new();
    let mut line = |s: String| {
        out.push_str(prefix);
        out.push_str(&s);
        out.push('\n');
    };
    line(format!("Chunked bench [{}] complete in {:?}", label, r.elapsed));
    line(format!("  Total requests: {}", r.total_requests));
    line(format!("  Successful:     {}", r.successful));
    line(format!("  Failed:         {}", r.failed));
    line(format!("  Bytes sent:     {}", r.bytes_sent));
    line(format!("  Throughput:     {:.2} req/s", r.req_per_sec));
    line(format!(
        "  Latency:        avg={:.1}ms p50={}ms p95={}ms p99={}ms",
        r.avg_latency_ms, r.p50_ms, r.p95_ms, r.p99_ms
    ));
    if !r.status_counts.is_empty() {
        line("  Status codes:".to_string());
        let mut codes: Vec<_> = r.status_counts.iter().collect();
        codes.sort_by_key(|(k, _)| **k);
        for (code, n) in codes {
            line(format!("    {} = {}", code, n));
        }
    }
    // Surface the captured error responses so the user can tell whether
    // errors came from MockForge, an upstream proxy, a CDN, etc. Hint at the
    // most common cause when 5xx is involved (proxy upstream timeout on long
    // chunked uploads).
    if !r.error_samples.is_empty() {
        line("  Error response samples:".to_string());
        for s in &r.error_samples {
            let server = s.server_header.as_deref().unwrap_or("(no Server header)");
            line(format!("    [{}] Server: {}", s.status, server));
            if !s.body_excerpt.is_empty() {
                line(format!("       body: {}", s.body_excerpt.replace('\n', " ")));
            }
        }
        if r.status_counts.keys().any(|c| (500..600).contains(c)) {
            line(
                "  Hint: 5xx responses with `Server:` revealing a proxy/LB usually \
                 mean the proxy timed out reading from upstream. Each chunked request \
                 takes >= (total_size_bytes / chunk_size_bytes) * chunk_interval_ms; \
                 if that exceeds the proxy's upstream timeout, errors are inevitable."
                    .to_string(),
            );
        }
    }
    print!("{out}");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn plan_target_runs_merges_headers_auth_spec_and_output() {
        let mut target_headers = HashMap::new();
        target_headers.insert("X-Env".to_string(), "b".to_string());
        let targets = vec![
            TargetConfig::from_url("10.0.0.1:8080".to_string()),
            TargetConfig {
                url: "https://b.example".to_string(),
                auth: Some("Bearer t".to_string()),
                headers: Some(target_headers),
                spec: Some(PathBuf::from("b.yaml")),
            },
        ];
        let mut cli_headers = HashMap::new();
        cli_headers.insert("X-Env".to_string(), "cli".to_string());
        let cli_spec = PathBuf::from("cli.yaml");

        let runs = plan_target_runs(targets, &cli_headers, Some(&cli_spec), Path::new("out"));

        assert_eq!(runs.len(), 2);
        assert_eq!(runs[0].url, "http://10.0.0.1:8080");
        assert_eq!(runs[0].headers.get("X-Env").map(String::as_str), Some("cli"));
        assert_eq!(runs[0].spec.as_deref(), Some(Path::new("cli.yaml")));
        assert_eq!(runs[0].output, Path::new("out").join("target_1"));
        assert_eq!(runs[0].prefix, "[target_1] ");

        assert_eq!(runs[1].headers.get("X-Env").map(String::as_str), Some("b"));
        assert_eq!(runs[1].headers.get("Authorization").map(String::as_str), Some("Bearer t"));
        assert_eq!(runs[1].spec.as_deref(), Some(Path::new("b.yaml")));
        assert_eq!(runs[1].output, Path::new("out").join("target_2"));
    }

    fn cmd(target: Option<&str>, targets_file: Option<PathBuf>) -> ChunkedCommand {
        ChunkedCommand {
            target: target.map(str::to_string),
            targets_file,
            max_concurrency: 10,
            spec: None,
            base_path: None,
            operation_id: None,
            method: "POST".to_string(),
            concurrency: 1,
            duration: Duration::from_millis(10),
            chunk_size_bytes: 1024,
            total_size_bytes: 4096,
            chunk_interval_ms: 0,
            headers: HashMap::new(),
            insecure: false,
            validate_requests: false,
            export_requests: false,
            output: PathBuf::from("bench-results"),
        }
    }

    #[tokio::test]
    async fn rejects_target_and_targets_file_together() {
        let c = cmd(Some("http://127.0.0.1:1"), Some(PathBuf::from("t.txt")));
        assert!(c.execute().await.is_err());
    }

    #[tokio::test]
    async fn rejects_neither_target_nor_targets_file() {
        assert!(cmd(None, None).execute().await.is_err());
    }
}