tidaldb/tidal-stress/src/scheduler.rs

//! Open-loop, constant-arrival-rate load scheduler — the methodology that makes
//! the numbers trustworthy.
//!
//! A CLOSED-loop generator (N workers each doing `send(); await; repeat`) sends
//! FEWER requests when the server slows down, so it silently under-measures both
//! load and latency — the "coordinated omission" error. This scheduler instead
//! fires requests at a fixed target rate regardless of outstanding responses,
//! and measures each request's latency from its *intended* send time, so a server
//! stall inflates the very percentiles a closed-loop test would hide.
//!
//! When the in-flight cap is reached we COUNT a client-shed rather than block —
//! blocking would re-introduce the closed-loop coupling. A non-zero shed count is
//! reported as "the generator, not the server, is the limit here".

use std::sync::Arc;
use std::time::Duration;

use tokio::sync::{Semaphore, mpsc};
use tokio::time::Instant;

use crate::client::HttpClient;
use crate::error::{Result, StressError};
use crate::metrics::{LatencyHistogram, Outcome, StageStats};
use crate::workload::Workload;

/// One step of the ramp: hold `target_rps` for `duration`.
pub struct Stage {
    pub target_rps: f64,
    pub duration: Duration,
}

/// Run one stage open-loop and return its aggregated stats.
pub async fn run_stage(
    workload: Arc<Workload>,
    client: Arc<HttpClient>,
    stage: &Stage,
    max_inflight: usize,
) -> StageStats {
    let (tx, mut rx) = mpsc::unbounded_channel::<Outcome>();
    let collector = tokio::spawn(async move {
        let mut stats = StageStats::new();
        while let Some(o) = rx.recv().await {
            stats.record(&o);
        }
        stats
    });

    let sem = Arc::new(Semaphore::new(max_inflight));
    // These are touched ONLY by this single dispatcher task — they are never moved
    // into the spawned request tasks — so they are plain locals, not atomics. The
    // local histogram captures the lag DISTRIBUTION so the generator's own catch-up
    // stalls surface as p99/max instead of being averaged away into the mean.
    let mut shed: u64 = 0;
    let mut lag_sum_ns: u64 = 0;
    let mut lag_count: u64 = 0;
    let mut lag_hist = LatencyHistogram::default();

    let start = Instant::now();
    let deadline = start + stage.duration;
    let period = Duration::from_secs_f64(1.0 / stage.target_rps.max(1e-9));
    let mut next = start;
    let mut dispatched: u64 = 0;

    loop {
        // One clock read per iteration drives both the deadline check and the
        // dispatch decision (the value was previously read twice with no work
        // between the reads).
        let now = Instant::now();
        if now >= deadline {
            break;
        }
        if next <= now {
            match sem.clone().try_acquire_owned() {
                Ok(permit) => {
                    let intended = next;
                    let lag = now.saturating_duration_since(intended);
                    lag_sum_ns += lag.as_nanos() as u64;
                    lag_count += 1;
                    lag_hist.record(lag);

                    let workload = workload.clone();
                    let client = client.clone();
                    let tx = tx.clone();
                    tokio::spawn(async move {
                        let _permit = permit;
                        // Build the plan and DROP the RNG before any await:
                        // `ThreadRng` is !Send (holds an `Rc`), so it must not be
                        // live across `.await` or the task future isn't `Send`.
                        let plan = {
                            let mut rng = rand::rng();
                            workload.next(&mut rng)
                        };
                        let op = plan.op;
                        let class = client.execute(&plan).await;
                        // CO-corrected latency: from when the request was DUE, not
                        // when it was sent — so client backlog counts against us.
                        let latency = Instant::now().saturating_duration_since(intended);
                        let _ = tx.send(Outcome { op, class, latency });
                    });
                }
                Err(_) => {
                    shed += 1;
                }
            }
            next += period;
            dispatched += 1;
            // Let the local runtime worker service spawned tasks during a
            // catch-up burst instead of monopolising it.
            if dispatched.is_multiple_of(256) {
                tokio::task::yield_now().await;
            }
        } else {
            tokio::time::sleep_until(next).await;
        }
    }
    let elapsed = start.elapsed();

    // Dropping the producer's sender lets the collector finish once every in-flight
    // task (each holding its own sender clone) completes — bounded by the client
    // request timeout, so this can never hang indefinitely.
    drop(tx);
    let mut stats = collector.await.unwrap_or_else(|_| StageStats::new());
    stats.elapsed = elapsed;
    stats.client_shed = shed;
    stats.mean_schedule_lag = if lag_count == 0 {
        Duration::ZERO
    } else {
        Duration::from_nanos(lag_sum_ns / lag_count)
    };
    stats.p99_schedule_lag = lag_hist.percentile(0.99);
    stats.max_schedule_lag = lag_hist.max();
    stats
}

// ── Ramp parsing ─────────────────────────────────────────────────────────────

/// Parse a ramp: a preset name or an explicit `rps:secs,rps:secs,...` spec.
///
/// Presets are pinned to the thepeach 100k-DAU model (see crate docs): at the
/// `peach` mix (~90% signals) a 100k-DAU TikTok-style evening peak is ≈3,900
/// total req/s, so `peach-100k` brackets and then doubles past that to answer
/// "can we handle more?". `secs` defaults from the preset.
///
/// # Errors
///
/// Returns [`StressError::Ramp`] if a custom `rps:secs,...` spec is malformed:
/// a part missing the `:` separator, an unparseable rps or secs value, or a
/// spec that yields no stages at all.
pub fn parse_ramp(spec: &str, stage_secs: u64) -> Result<Vec<Stage>> {
    let mk = |rates: &[f64]| -> Vec<Stage> {
        rates
            .iter()
            .map(|&r| Stage {
                target_rps: r,
                duration: Duration::from_secs(stage_secs),
            })
            .collect()
    };
    match spec {
        // Local validation through a port-forward (low, short).
        "smoke" => Ok(vec![
            Stage {
                target_rps: 10.0,
                duration: Duration::from_secs(10),
            },
            Stage {
                target_rps: 40.0,
                duration: Duration::from_secs(10),
            },
        ]),
        "quick" => Ok(mk(&[100.0, 500.0, 1500.0, 4000.0])),
        // The headline ramp: low → past a 100k-DAU peak → 2× beyond.
        "peach-100k" | "default" => Ok(mk(&[
            50.0, 150.0, 400.0, 800.0, 1500.0, 3000.0, 5000.0, 8000.0,
        ])),
        // Push to find the hard ceiling.
        "max" => Ok(mk(&[
            500.0, 1500.0, 3000.0, 6000.0, 10000.0, 15000.0, 20000.0,
        ])),
        _ => {
            let mut stages = Vec::new();
            for part in spec.split(',') {
                let part = part.trim();
                if part.is_empty() {
                    continue;
                }
                let (r, s) = part
                    .split_once(':')
                    .ok_or_else(|| StressError::Ramp(format!("expected rps:secs, got '{part}'")))?;
                let target_rps: f64 = r
                    .trim()
                    .parse()
                    .map_err(|_| StressError::Ramp(format!("bad rps '{r}'")))?;
                let secs: u64 = s
                    .trim()
                    .parse()
                    .map_err(|_| StressError::Ramp(format!("bad secs '{s}'")))?;
                stages.push(Stage {
                    target_rps,
                    duration: Duration::from_secs(secs),
                });
            }
            if stages.is_empty() {
                return Err(StressError::Ramp("empty ramp".into()));
            }
            Ok(stages)
        }
    }
}

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

    #[test]
    fn ramp_presets_and_specs_parse() {
        assert_eq!(parse_ramp("peach-100k", 45).expect("preset").len(), 8);
        let custom = parse_ramp("100:30,500:30,2000:60", 45).expect("spec");
        assert_eq!(custom.len(), 3);
        assert!((custom[2].target_rps - 2000.0).abs() < f64::EPSILON);
        assert_eq!(custom[2].duration.as_secs(), 60);
        assert!(parse_ramp("oops", 45).is_err());
        assert!(parse_ramp("100-30", 45).is_err());
    }
}