aion-server 0.29.0

//! Requirement 2c: what commit-before-broadcast COSTS at token-chunk rate, on
//! the real Haematite store.
//!
//! The ordering is not in question — `live.rs::Recorder::record` appends and
//! then broadcasts, and `lifecycle_tests.rs` pins that a refused append
//! broadcasts nothing. What was unmeasured is the PRICE: one durable commit per
//! streamed chunk, on a surface that streams at whatever rate a model emits
//! tokens.
//!
//! This measures it and PRINTS it, with its conditions. It is a measurement
//! rather than a threshold on purpose: a pass/fail number here would be a
//! performance budget nobody has agreed, and a budget invented by a test is
//! exactly the kind of arbitrary limit this codebase refuses. What it asserts is
//! only that the run happened at all — every chunk committed, every chunk
//! arrived, and the indices are dense — so the numbers it prints are numbers
//! about a complete run.
//!
//! # What is measured
//!
//! - **sustained chunks/second**: `CHUNKS` frames recorded back to back through
//!   the real `Recorder`, wall-clock end to end.
//! - **the p99 commit→broadcast gap**: for each frame, the time from the instant
//!   the recorder returns (the commit is durable) to the instant a subscribed
//!   watcher has the frame in hand. This is the delay the ordering buys, and it
//!   is the number a coalescing window would be argued from.
//!
//! # No coalescing window is added
//!
//! Deliberately, and this test does not measure one. A batch size or a flush
//! interval is a behavioural value with no operator behind it yet; the
//! instruction was to measure first. The number this prints is what that
//! argument would have to start from.
//!
//! # Conditions are PRINTED, not assumed
//!
//! A throughput number without its conditions is decoration, so the output
//! carries the store mode, the host's core count, and the build profile. Read it
//! off the run rather than off this file.

use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use aion_core::{AssistantSessionEvent, AssistantSessionId, ContentType, Payload};
use aion_store::StoreError;
use aion_store::assistant::{AssistantSessionRecord, AssistantSessionStore};
use aion_store_haematite::HaematiteStore;
use chrono::Utc;

/// How many chunks the sustained measurement records.
///
/// Large enough that the per-chunk cost dominates the fixture's own setup, and
/// small enough that the cell stays a measurement rather than a soak. Not a
/// budget and not a limit on anything the server does.
const CHUNKS: usize = 2_000;

/// A chunk the size of a plausible model token burst.
const CHUNK_TEXT: &str = "the quick brown fox jumps over the lazy dog";

fn unique_temp_dir(name: &str) -> std::path::PathBuf {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |duration| duration.as_nanos());
    std::env::temp_dir().join(format!(
        "aion-assistant-commit-broadcast-{name}-{}-{nanos}",
        std::process::id()
    ))
}

/// The `p`-th percentile of `samples`, nearest-rank. `samples` is sorted here.
///
/// Integer arithmetic throughout: the answer is one of the samples actually
/// taken rather than an interpolation between two of them, and there is no
/// floating-point rounding standing between the measurement and the number
/// printed.
fn percentile(samples: &mut [Duration], per_cent: usize) -> Duration {
    if samples.is_empty() {
        return Duration::ZERO;
    }
    samples.sort_unstable();
    // Nearest-rank: ceil(p/100 × N), as integer division with the ceiling done
    // by adding 99 before dividing.
    let rank = samples.len().saturating_mul(per_cent).saturating_add(99) / 100;
    samples[rank.clamp(1, samples.len()) - 1]
}

/// Chunks per second, to one decimal place, without a float in the arithmetic.
///
/// Returned as a rendered string because that is all it is for: a number in a
/// printed line. Doing the rounding here in integers keeps the value exact.
fn rendered_rate(chunks: usize, elapsed: Duration) -> String {
    let micros = elapsed.as_micros();
    if micros == 0 {
        return String::from("immeasurable (the run took under a microsecond)");
    }
    let Ok(chunks) = u128::try_from(chunks) else {
        return String::from("immeasurable (the chunk count does not fit)");
    };
    // ×10 for the one decimal place, ×1_000_000 to turn microseconds into
    // seconds; both before the division so nothing is lost to truncation.
    let tenths = chunks.saturating_mul(10_000_000) / micros;
    format!("{}.{}", tenths / 10, tenths % 10)
}

/// The sustained rate, and the gap the ordering buys, on real Haematite.
///
/// Gated at runtime rather than with `#[ignore]`: it is a measurement, it runs
/// wherever the battery runs it, and it prints what it found.
#[tokio::test(flavor = "multi_thread")]
async fn commit_before_broadcast_at_token_chunk_rate() -> Result<(), StoreError> {
    let store = Arc::new(HaematiteStore::create(
        unique_temp_dir("single-node"),
        haematite::NodeCacheBudget::Unlimited,
    )?);
    let session_id = AssistantSessionId::new_v4();
    let now = Utc::now();
    store
        .put_assistant_session(AssistantSessionRecord {
            session_id,
            subject: String::from("operator"),
            harness: String::from("claude"),
            account: None,
            title: Some(String::from("the measurement")),
            created_at: now,
            updated_at: now,
            turns: 0,
            mcp_token_digest: None,
            commands: Vec::new(),
        })
        .await?;

    // A broadcast channel the same shape the registry mints, with a watcher
    // taken BEFORE the first commit — the gap is only measurable from a
    // subscriber that was already listening.
    let (sender, mut watcher) =
        tokio::sync::broadcast::channel::<(u64, Instant)>(CHUNKS.next_power_of_two());

    let mut gaps: Vec<Duration> = Vec::with_capacity(CHUNKS);
    let started = Instant::now();
    for nth in 0..CHUNKS {
        let event = AssistantSessionEvent::Delta {
            turn_id: String::from("t-1"),
            text: format!("{CHUNK_TEXT} {nth}"),
        };
        let bytes = serde_json::to_vec(&event)
            .map_err(|error| StoreError::Serialization(error.to_string()))?;
        // The commit. Exactly what `Recorder::record` does first.
        let index = store
            .append_assistant_transcript_event(
                &session_id,
                Utc::now(),
                Payload::new(ContentType::Json, bytes),
            )
            .await?;
        let committed = Instant::now();
        // The broadcast. Exactly what `Recorder::record` does second.
        drop(sender.send((index, committed)));
        // And the watcher's own receipt, which is where the gap ends.
        let (received, committed) = watcher.recv().await.map_err(|error| {
            StoreError::Backend(format!(
                "a broadcast frame was lost, so the gap cannot be measured: {error}"
            ))
        })?;
        assert_eq!(
            received, index,
            "the frame a watcher receives is the frame that was committed"
        );
        gaps.push(committed.elapsed());
    }
    let elapsed = started.elapsed();

    // The run really happened: every chunk is durable, dense, and in order. The
    // numbers below are about a complete run or they are about nothing.
    let transcript = store.assistant_transcript(&session_id, None).await?;
    assert_eq!(transcript.len(), CHUNKS, "every chunk must be durable");
    let indices: Vec<u64> = transcript.iter().map(|event| event.index).collect();
    let expected: Vec<u64> = (0..CHUNKS as u64).collect();
    assert_eq!(
        indices, expected,
        "the transcript must be dense and ordered"
    );

    let rate = rendered_rate(CHUNKS, elapsed);
    let p99 = percentile(&mut gaps, 99);
    let p50 = percentile(&mut gaps, 50);
    // `tracing` is not initialized in this binary, so the measurement is printed
    // where a battery reading the test's own output will find it.
    println!(
        "assistant commit-before-broadcast measurement\n  \
         store mode:            haematite, single node, 1 shard, NodeCacheBudget::Unlimited\n  \
         build profile:         {profile}\n  \
         host parallelism:      {cores}\n  \
         chunk payload:         {payload} bytes of JSON\n  \
         chunks:                {CHUNKS}\n  \
         wall clock:            {elapsed:?}\n  \
         sustained chunks/s:    {rate}\n  \
         commit→broadcast p50:  {p50:?}\n  \
         commit→broadcast p99:  {p99:?}",
        profile = if cfg!(debug_assertions) {
            "debug (unoptimized)"
        } else {
            "release"
        },
        cores = std::thread::available_parallelism().map_or_else(
            |error| format!("unknown ({error})"),
            |count| count.to_string()
        ),
        payload = transcript
            .first()
            .map_or(0, |event| event.payload.bytes().len()),
    );
    Ok(())
}