Aura/crates/aura-engine/tests/streaming_reduction_memory.rs

//! BLOCKER #138 capstone: the R-reduction's peak retained memory is
//! O(trades)+O(1), independent of cycle count. Drives `GatedRecorder` (gate =
//! CLOSED) through its real `mpsc` channel with a FIXED small number of closed
//! trades but a VARIABLE large number of hold cycles, generated LAZILY (one row
//! at a time, never all at once), and measures peak live bytes with a counting
//! global allocator. A retain-everything sink's peak would scale with the cycle
//! count; the folding sink's does not. Sole test in this binary, so the global
//! peak counter is not polluted by other tests running in parallel.

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc;

use aura_core::{AnyColumn, Ctx, Node, Scalar, ScalarKind, Timestamp};
use aura_std::{GatedRecorder, PM_RECORD_KINDS, PM_WIDTH};

struct CountingAlloc;
static LIVE: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for CountingAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(layout) };
        if !p.is_null() {
            let now = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
            PEAK.fetch_max(now, Ordering::Relaxed);
        }
        p
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static A: CountingAlloc = CountingAlloc;

fn reset_peak() {
    PEAK.store(LIVE.load(Ordering::Relaxed), Ordering::Relaxed);
}
fn peak() -> usize {
    PEAK.load(Ordering::Relaxed)
}

const CLOSED: usize = 0;
const REALIZED_R: usize = 1;
const K_TRADES: usize = 8;

/// Peak live bytes to drive `GatedRecorder` over `K_TRADES` closed trades, each
/// followed by `holds` hold cycles, sourced lazily (one row built at a time),
/// then fold the emitted rows. With a folding sink the peak is O(trades).
fn peak_bytes(holds: usize) -> usize {
    reset_peak();
    let (tx, rx) = mpsc::channel();
    let mut g = GatedRecorder::new(&PM_RECORD_KINDS, CLOSED, tx);
    let mut cols: Vec<AnyColumn> =
        PM_RECORD_KINDS.iter().map(|&k| AnyColumn::with_capacity(k, 1)).collect();
    let mut ts = 0i64;
    let push = |g: &mut GatedRecorder, cols: &mut Vec<AnyColumn>, ts: &mut i64, closed: bool, r: f64| {
        // kind-correct zeros: the PM record is heterogeneous (PM_RECORD_KINDS has
        // Bool/I64/Timestamp slots), so each unset slot defaults to its column's
        // kind — a blanket f64 zero would be rejected by the kind-typed push.
        let mut row: Vec<Scalar> = PM_RECORD_KINDS
            .iter()
            .map(|&k| match k {
                ScalarKind::F64 => Scalar::f64(0.0),
                ScalarKind::I64 => Scalar::i64(0),
                ScalarKind::Bool => Scalar::bool(false),
                ScalarKind::Timestamp => Scalar::ts(Timestamp(0)),
            })
            .collect();
        debug_assert_eq!(row.len(), PM_WIDTH);
        row[CLOSED] = Scalar::bool(closed);
        row[REALIZED_R] = Scalar::f64(r);
        for (i, s) in row.iter().enumerate() {
            cols[i].push(*s).unwrap();
        }
        g.eval(Ctx::new(cols, Timestamp(*ts)));
        *ts += 1;
    };
    for k in 0..K_TRADES {
        let r = if k % 2 == 0 { 1.0 } else { -0.5 };
        push(&mut g, &mut cols, &mut ts, true, r);
        for _ in 0..holds {
            push(&mut g, &mut cols, &mut ts, false, 0.0);
        }
    }
    g.finalize();
    let rows: Vec<(Timestamp, Vec<Scalar>)> = rx.try_iter().collect();
    // the K_TRADES gated (closed) rows, plus the one genuine final row the
    // GatedRecorder flushes on finalize (the last cycle is a hold, so it was not
    // gated) — O(trades)+O(1), independent of `holds`, which is the whole point.
    assert_eq!(rows.len(), K_TRADES + 1, "only the closed trades + the final row are retained");
    let p = peak();
    drop(rows);
    p
}

#[test]
fn r_reduction_peak_memory_is_independent_of_cycle_count() {
    const SMALL_HOLDS: usize = 64; //   8 * (1+64)    =     520 cycles
    const LARGE_HOLDS: usize = 50_000; // 8 * (1+50000) = 400_008 cycles

    let small = peak_bytes(SMALL_HOLDS);
    let large = peak_bytes(LARGE_HOLDS);

    // cycle count grew ~770x; an O(trades)+O(1) sink holds its peak essentially
    // flat. Generous 4x headroom so only genuine O(cycles) retention trips it.
    assert!(
        large <= small.saturating_mul(4),
        "R-reduction peak scales with cycle count (O(cycles) retention): \
         small {} cycles -> {} bytes; large {} cycles -> {} bytes ({:.0}x). \
         The folding sink must retain O(trades)+O(1).",
        K_TRADES * (1 + SMALL_HOLDS),
        small,
        K_TRADES * (1 + LARGE_HOLDS),
        large,
        large as f64 / small.max(1) as f64,
    );
}