tidaldb/tidal-stress/src/recall.rs

//! The m12p1 recall oracle + open-loop read-recall ramp.
//!
//! Every perf claim before m12p1 was latency-only: recall — the fraction of the
//! true nearest neighbors an ANN query actually returns — was unmeasured at the
//! production shape (1536-dim, 100k–1M items). This module closes that gap. It
//! is the harness the whole milestone steers against.
//!
//! # How the oracle works
//!
//! 1. The corpus is seeded with **deterministic, id-keyed** embeddings
//!    ([`embedding_for`]), so the generator can reconstruct, bit-for-bit, the
//!    exact vectors the engine indexed — no need to capture them over the wire.
//! 2. [`GroundTruth`] holds that corpus in RAM and computes **brute-force cosine
//!    top-k** for any query — the exact answer the ANN index is approximating.
//!    (Order by cosine == order by the engine's L2-on-normalized distance, so the
//!    two are directly comparable; see `storage::vector`.)
//! 3. A [`QueryPool`] of realistic queries (corpus points perturbed by noise) is
//!    precomputed once, each with its ground-truth top-k, so the hot request path
//!    does only a set-overlap to score `recall@k`.
//! 4. [`run_recall_stage`] fires `/vector_search` probes open-loop (the same
//!    coordinated-omission-corrected methodology as the signal ramp) and records,
//!    per request, the true latency AND the achieved recall.
//!
//! The verdict reports per-stage **true p99** (not a closed-loop mean) and **mean
//! recall@10**, and finds the **read-knee**: the highest sustained QPS at which
//! `p99 ≤ target AND recall@10 ≥ target` both hold.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

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

use crate::client::HttpClient;
use crate::metrics::{LatencyHistogram, StatusClass};
use crate::scheduler::Stage;

// ── Deterministic corpus + query generation ──────────────────────────────────

/// SplitMix64 — a tiny, fast, dependency-free deterministic generator. Seeding it
/// from an entity id makes the whole corpus reproducible from the id alone, which
/// is what lets the generator hold the ground truth without ever reading it back
/// from the engine. The same scheme is mirrored in the engine-side recall test.
fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// One uniform `f32` in `[-0.5, 0.5)` drawn from `state`.
fn next_f32(state: &mut u64) -> f32 {
    // Top 24 bits → a uniform mantissa in [0,1), then recentre to [-0.5, 0.5).
    let bits = (splitmix64(state) >> 40) as u32;
    (bits as f32 / 16_777_216.0) - 0.5
}

/// Corpus ids per cluster center. ~100/cluster mirrors the neighbourhood size of
/// a real embedding manifold and guarantees ≥ `k` genuine neighbours per query
/// for any `k` ≤ 100.
const CLUSTER_SIZE: u64 = 100;

/// Per-component spread of the Gaussian mixture (`vector = center + SPREAD·noise`,
/// both unit). `SPREAD = 0.5` ⇒ intra-cluster cosine ≈ 0.89, inter-cluster ≈ 0 —
/// the clear neighbour structure of real embeddings.
const CLUSTER_SPREAD: f32 = 0.5;

/// A normalized pseudo-random unit vector seeded by `seed`.
fn unit_seeded(seed: u64, dim: usize) -> Vec<f32> {
    let mut state = seed;
    let mut v: Vec<f32> = (0..dim).map(|_| next_f32(&mut state)).collect();
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > f32::EPSILON {
        for x in &mut v {
            *x /= norm;
        }
    } else if let Some(first) = v.first_mut() {
        *first = 1.0;
    }
    v
}

/// The deterministic content embedding for an item id — a **clustered**
/// (Gaussian-mixture) vector.
///
/// `id` is assigned to cluster `id / CLUSTER_SIZE`; the vector is that cluster's
/// unit center plus `CLUSTER_SPREAD ×` a per-id unit noise vector. The result is
/// returned raw (un-normalized); the engine L2-normalizes on write and the oracle
/// ranks by cosine, both scale-invariant, so the cluster structure is identical
/// on both sides.
///
/// **Why clustered, not uniform-random.** Uniform-random high-dimensional vectors
/// are pathological for recall@k: by concentration of measure every pair sits at
/// cosine ≈ 0, so beyond a tiny perturbation a query's true top-k is an arbitrary
/// draw from a thick equidistant shell — recall@k then measures impossible
/// tie-breaking, not index quality, and (measured) *falls* as the corpus grows
/// (≈0.97 at 10k → ≈0.54 at 100k at 1536-D). Real text/image embeddings live on a
/// low-dimensional manifold with clusters; this models that so recall@k is a
/// meaningful index-quality metric at scale (m12p3).
#[must_use]
pub fn embedding_for(id: u64, dim: usize) -> Vec<f32> {
    let cluster = id / CLUSTER_SIZE;
    // Distinct stream constants for the center vs the per-id noise so they are
    // independent; offset+odd-multiply keeps adjacent ids/clusters well-separated.
    let center = unit_seeded(
        cluster
            .wrapping_mul(0x9E37_79B9_7F4A_7C15)
            .wrapping_add(0x00C0_FFEE),
        dim,
    );
    let noise = unit_seeded(id.wrapping_mul(0x2545_F491_4F6C_DD1D).wrapping_add(1), dim);
    center
        .iter()
        .zip(&noise)
        .map(|(c, n)| c + CLUSTER_SPREAD * n)
        .collect()
}

// ── Brute-force ground truth ─────────────────────────────────────────────────

/// The seeded corpus held in RAM for exact nearest-neighbor computation.
///
/// Stored as one flat `Vec<f32>` (id `i` occupies `[(i-1)*dim, i*dim)`) plus a
/// per-vector L2 norm so cosine is a single dot-product + divide. At 1536-dim
/// this costs `n * dim * 4` bytes (~6 GB at 1M) — the honest price of a real
/// brute-force oracle; smaller corpora (100k ≈ 600 MB) fit comfortably.
pub struct GroundTruth {
    flat: Vec<f32>,
    norms: Vec<f32>,
    dim: usize,
    n: u64,
}

impl GroundTruth {
    /// Build the deterministic corpus for ids `1..=n` at `dim` dimensions.
    #[must_use]
    pub fn build(n: u64, dim: usize) -> Self {
        let mut flat = Vec::with_capacity((n as usize) * dim);
        let mut norms = Vec::with_capacity(n as usize);
        for id in 1..=n {
            let v = embedding_for(id, dim);
            let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
            norms.push(norm);
            flat.extend_from_slice(&v);
        }
        Self {
            flat,
            norms,
            dim,
            n,
        }
    }

    #[must_use]
    pub const fn n(&self) -> u64 {
        self.n
    }

    #[must_use]
    pub const fn dim(&self) -> usize {
        self.dim
    }

    /// The raw stored vector for `id` (1-based).
    #[must_use]
    fn vector(&self, id: u64) -> &[f32] {
        let start = ((id - 1) as usize) * self.dim;
        &self.flat[start..start + self.dim]
    }

    /// The `k` ids whose corpus vector is most cosine-similar to `query`,
    /// best-first. The exact answer the ANN index approximates.
    #[must_use]
    pub fn top_k(&self, query: &[f32], k: usize) -> Vec<u64> {
        let q_norm = query.iter().map(|x| x * x).sum::<f32>().sqrt();
        if q_norm == 0.0 || k == 0 {
            return Vec::new();
        }
        // A bounded top-k kept as a min-by-score Vec (k is tiny, ~10): cheaper and
        // allocation-lighter than a full sort of n scored pairs per query.
        let mut top: Vec<(f32, u64)> = Vec::with_capacity(k + 1);
        for id in 1..=self.n {
            let xn = self.norms[(id - 1) as usize];
            if xn == 0.0 {
                continue;
            }
            let dot: f32 = query.iter().zip(self.vector(id)).map(|(a, b)| a * b).sum();
            let score = dot / (q_norm * xn); // cosine; higher = nearer
            if top.len() < k {
                top.push((score, id));
                if top.len() == k {
                    // Smallest score first so the worst-kept is at index 0.
                    top.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
                }
            } else if score > top[0].0 {
                top[0] = (score, id);
                // Re-sink the new minimum to the front (k is tiny).
                top.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
            }
        }
        // Return best-first.
        top.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
        top.into_iter().map(|(_, id)| id).collect()
    }
}

// ── Query pool (queries + precomputed ground truth) ──────────────────────────

/// A pool of realistic query vectors, each paired with its brute-force
/// ground-truth top-k.
///
/// Precomputing the (expensive) ground truth once decouples it from the hot
/// request path, where scoring a response is then just a set overlap.
pub struct QueryPool {
    pub queries: Vec<Vec<f32>>,
    pub truth: Vec<Vec<u64>>,
    pub k: usize,
}

impl QueryPool {
    /// Build `pool_size` queries and their ground-truth top-`k`.
    ///
    /// Each query is a corpus point (spread across the catalog) perturbed by
    /// deterministic noise, so it has a well-defined nearest cluster — the shape
    /// of a real preference-vector read, and far more informative than uniformly
    /// random high-dimensional queries whose neighbors are near-equidistant.
    ///
    /// Ground truth is computed in parallel across the available cores; the work
    /// is embarrassingly parallel (one independent brute-force scan per query).
    #[must_use]
    pub fn build(gt: &GroundTruth, pool_size: usize, k: usize, noise: f32) -> Self {
        let n = gt.n().max(1);
        // Spread base points across the corpus with a large odd stride so the pool
        // is not clustered on the low ids.
        let queries: Vec<Vec<f32>> = (0..pool_size)
            .map(|i| {
                let base = ((i as u64).wrapping_mul(2_654_435_761) % n) + 1;
                let mut q = gt.vector(base).to_vec();
                let mut state = (i as u64)
                    .wrapping_mul(0x100_0000_01B3)
                    .wrapping_add(0xABCD);
                for c in &mut q {
                    *c += noise * next_f32(&mut state);
                }
                q
            })
            .collect();

        let threads = std::thread::available_parallelism()
            .map(std::num::NonZeroUsize::get)
            .unwrap_or(4)
            .min(pool_size.max(1));
        let chunk = pool_size.div_ceil(threads.max(1));
        let truth: Vec<Vec<u64>> = std::thread::scope(|s| {
            // Spawn every chunk BEFORE joining any (the collect is load-bearing —
            // joining inline would serialize the scans), then concatenate in order.
            #[allow(clippy::needless_collect)]
            let handles: Vec<_> = queries
                .chunks(chunk.max(1))
                .map(|qs| s.spawn(move || qs.iter().map(|q| gt.top_k(q, k)).collect::<Vec<_>>()))
                .collect();
            handles
                .into_iter()
                .flat_map(|h| h.join().unwrap_or_default())
                .collect()
        });

        Self { queries, truth, k }
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.queries.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.queries.is_empty()
    }
}

/// `recall@k`: the fraction of the ground-truth top-k that the engine returned.
///
/// `returned` may be shorter than `k` (a thin corpus or a degraded index); the
/// denominator is the ground-truth size, never the returned size, so a short
/// answer is correctly penalised.
#[must_use]
pub fn recall_at_k(returned: &[u64], truth: &[u64]) -> f64 {
    if truth.is_empty() {
        return 1.0; // nothing to find ⇒ vacuously perfect
    }
    let truth_set: std::collections::HashSet<u64> = truth.iter().copied().collect();
    let hits = returned.iter().filter(|id| truth_set.contains(id)).count();
    hits as f64 / truth.len() as f64
}

// ── Open-loop recall stage ───────────────────────────────────────────────────

/// One recall request's result, sent from a worker to the collector.
struct RecallOutcome {
    class: StatusClass,
    /// CO-corrected latency (from the request's intended send time).
    latency: Duration,
    /// Achieved recall@k, present only when the response carried a usable id list.
    recall: Option<f64>,
}

/// Everything one recall ramp stage produced.
pub struct RecallStageStats {
    pub hist: LatencyHistogram,
    pub ok: u64,
    pub errors: u64,
    pub total: u64,
    pub recall_sum: f64,
    pub recall_count: u64,
    pub elapsed: Duration,
    pub client_shed: u64,
    pub p99_schedule_lag: Duration,
    pub max_schedule_lag: Duration,
}

impl RecallStageStats {
    fn new() -> Self {
        Self {
            hist: LatencyHistogram::default(),
            ok: 0,
            errors: 0,
            total: 0,
            recall_sum: 0.0,
            recall_count: 0,
            elapsed: Duration::ZERO,
            client_shed: 0,
            p99_schedule_lag: Duration::ZERO,
            max_schedule_lag: Duration::ZERO,
        }
    }

    fn record(&mut self, o: &RecallOutcome) {
        self.total += 1;
        if o.class == StatusClass::Ok {
            self.ok += 1;
            self.hist.record(o.latency);
        } else {
            self.errors += 1;
        }
        if let Some(r) = o.recall {
            self.recall_sum += r;
            self.recall_count += 1;
        }
    }

    /// Achieved throughput: completed requests per second over the stage.
    #[must_use]
    pub fn achieved_rps(&self) -> f64 {
        let s = self.elapsed.as_secs_f64();
        if s <= 0.0 { 0.0 } else { self.total as f64 / s }
    }

    #[must_use]
    pub fn error_rate(&self) -> f64 {
        if self.total == 0 {
            0.0
        } else {
            self.errors as f64 / self.total as f64
        }
    }

    /// Mean recall@k across the stage's successfully-parsed responses.
    #[must_use]
    pub fn mean_recall(&self) -> f64 {
        if self.recall_count == 0 {
            0.0
        } else {
            self.recall_sum / self.recall_count as f64
        }
    }

    /// True p99 latency (bucket-estimated, but a genuine tail over an open loop —
    /// NOT a closed-loop mean).
    #[must_use]
    pub fn p99(&self) -> Duration {
        self.hist.percentile(0.99)
    }
}

/// Run one recall ramp stage open-loop and return its aggregated stats.
///
/// Mirrors [`crate::scheduler::run_stage`]'s constant-arrival-rate methodology
/// (fire at the target rate regardless of outstanding responses; measure each
/// latency from its intended send time; count — never block on — an in-flight
/// shed), specialised to the single `/vector_search` op and extended to score
/// recall against the precomputed pool.
pub async fn run_recall_stage(
    pool: Arc<QueryPool>,
    bases: Arc<Vec<String>>,
    client: Arc<HttpClient>,
    stage: &Stage,
    k: usize,
    ef_search: Option<usize>,
    max_inflight: usize,
) -> RecallStageStats {
    let (tx, mut rx) = mpsc::unbounded_channel::<RecallOutcome>();
    let collector = tokio::spawn(async move {
        let mut stats = RecallStageStats::new();
        while let Some(o) = rx.recv().await {
            stats.record(&o);
        }
        stats
    });

    let sem = Arc::new(Semaphore::new(max_inflight));
    let next_base = AtomicUsize::new(0);
    let mut shed: 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 {
        let now = Instant::now();
        if now >= deadline {
            break;
        }
        if next <= now {
            match sem.clone().try_acquire_owned() {
                Ok(permit) => {
                    let intended = next;
                    lag_hist.record(now.saturating_duration_since(intended));

                    // Pick the query (round-robin over the pool) and a target base
                    // (round-robin over the gateways — every replica holds the full
                    // corpus at the S=1 recall shape).
                    let qi = (dispatched as usize) % pool.len().max(1);
                    let base_idx = next_base.fetch_add(1, Ordering::Relaxed) % bases.len().max(1);

                    let pool = pool.clone();
                    let bases = bases.clone();
                    let client = client.clone();
                    let tx = tx.clone();
                    tokio::spawn(async move {
                        let _permit = permit;
                        let query = &pool.queries[qi];
                        let truth = &pool.truth[qi];
                        let base = &bases[base_idx];
                        let (class, ids) =
                            client.vector_search_ids(base, query, k, ef_search).await;
                        let latency = Instant::now().saturating_duration_since(intended);
                        let recall = ids.map(|ids| recall_at_k(&ids, truth));
                        let _ = tx.send(RecallOutcome {
                            class,
                            latency,
                            recall,
                        });
                    });
                }
                Err(_) => shed += 1,
            }
            next += period;
            dispatched += 1;
            if dispatched.is_multiple_of(256) {
                tokio::task::yield_now().await;
            }
        } else {
            tokio::time::sleep_until(next).await;
        }
    }
    let elapsed = start.elapsed();

    drop(tx);
    let mut stats = collector.await.unwrap_or_else(|_| RecallStageStats::new());
    stats.elapsed = elapsed;
    stats.client_shed = shed;
    stats.p99_schedule_lag = lag_hist.percentile(0.99);
    stats.max_schedule_lag = lag_hist.max();
    stats
}

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

    #[test]
    fn embedding_is_deterministic_and_sized() {
        let a = embedding_for(7, 32);
        let b = embedding_for(7, 32);
        assert_eq!(a, b, "same id+dim must reproduce the same vector");
        assert_eq!(a.len(), 32);
        assert_ne!(embedding_for(7, 32), embedding_for(8, 32), "ids differ");
        // Non-zero norm (the engine rejects zero-norm embeddings).
        assert!(a.iter().map(|x| x * x).sum::<f32>() > 0.0);
    }

    #[test]
    fn ground_truth_ranks_self_first() {
        let gt = GroundTruth::build(200, 16);
        // Querying with a corpus vector returns that id first (cosine 1.0).
        let q = embedding_for(42, 16);
        let top = gt.top_k(&q, 10);
        assert_eq!(top.len(), 10);
        assert_eq!(top[0], 42, "an item's own vector is its nearest neighbor");
    }

    #[test]
    fn recall_at_k_counts_overlap_over_truth_size() {
        let truth = vec![1, 2, 3, 4, 5];
        // 3 of 5 truth ids present (extra/ordering ignored).
        assert!((recall_at_k(&[1, 2, 3, 99, 100], &truth) - 0.6).abs() < 1e-9);
        assert!((recall_at_k(&[], &truth) - 0.0).abs() < 1e-9);
        assert!((recall_at_k(&[1, 2, 3, 4, 5], &truth) - 1.0).abs() < 1e-9);
        // Empty truth ⇒ vacuously perfect (nothing to find).
        assert!((recall_at_k(&[], &[]) - 1.0).abs() < 1e-9);
    }

    #[test]
    fn query_pool_truth_aligns_with_queries() {
        let gt = GroundTruth::build(500, 16);
        let pool = QueryPool::build(&gt, 20, 10, 0.05);
        assert_eq!(pool.queries.len(), 20);
        assert_eq!(pool.truth.len(), 20, "every query has a ground-truth list");
        for t in &pool.truth {
            assert_eq!(t.len(), 10, "ground truth is top-k");
        }
        // A perturbed corpus point's nearest neighbor should be its own base or a
        // close id — the brute force must return a non-empty, sane list.
        assert!(pool.truth.iter().all(|t| !t.is_empty()));
    }
}