nordisk/holger

server/lib/tests/export_stream_memory.rs

//! ★ **RA1-13: the bundle export door streams, it does not hold the bundle twice.**
//!
//! `BundleService.Export` answers an airgap bundle candidate (an OCI bundle is gigabytes).
//! The door used to make EVERY chunk up front (a copy of the whole bundle, collected into a
//! Vec) while still holding the member map: the bundle twice in memory before a byte left.
//! `grpc::doors::export_stream` owns the members and makes each chunk only when the stream is
//! polled, and frees a member once it is sent.
//!
//! This binary counts every live heap byte (its own global allocator) and streams a 192 MiB
//! candidate: the extra memory while it streams stays within a few chunks, and every byte
//! still arrives, in order, member by member, then the summary.

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

use holger_server_lib::grpc::doors::{export_stream, EXPORT_CHUNK_BYTES};
use holger_server_lib::grpc::holger_doors_proto as pb;

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

unsafe impl GlobalAlloc for Counting {
    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::SeqCst) + layout.size();
            PEAK.fetch_max(now, Ordering::SeqCst);
        }
        p
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) };
        LIVE.fetch_sub(layout.size(), Ordering::SeqCst);
    }
}

#[global_allocator]
static A: Counting = Counting;

const MIB: usize = 1024 * 1024;

#[test]
fn a_large_bundle_streams_in_bounded_memory() {
    // A 192 MiB candidate: three 64 MiB payloads and a small evidence member.
    let mut members = std::collections::BTreeMap::new();
    for (i, name) in ["payload/a.tar", "payload/b.tar", "payload/c.tar"]
        .iter()
        .enumerate()
    {
        members.insert(name.to_string(), vec![i as u8 + 1; 64 * MIB]);
    }
    members.insert("evidence/sbom.json".to_string(), b"{}".to_vec());
    let total: usize = members.values().map(Vec::len).sum();
    let summary = pb::ExportBundleSummary {
        bundle_id: "big".into(),
        members: 4,
        ..Default::default()
    };

    let held = LIVE.load(Ordering::SeqCst);
    PEAK.store(held, Ordering::SeqCst);
    let mut got = 0usize;
    let mut order: Vec<String> = Vec::new();
    let mut last_summary = None;
    let mut after_a = usize::MAX;
    for c in export_stream(members, EXPORT_CHUNK_BYTES, summary.clone()) {
        assert!(
            c.data.len() <= EXPORT_CHUNK_BYTES,
            "a chunk is at most one chunk"
        );
        if let Some(first) = c.data.first().filter(|_| c.path.starts_with("payload/")) {
            assert!(
                c.data.iter().all(|b| b == first),
                "a chunk is its member's bytes"
            );
        }
        got += c.data.len();
        if order.last() != Some(&c.path) {
            if c.path == "payload/b.tar" {
                after_a = LIVE.load(Ordering::SeqCst);
            }
            order.push(c.path.clone());
        }
        last_summary = c.summary.clone();
        // the chunk is dropped here, as tonic drops it once it is on the wire
    }
    let extra = PEAK.load(Ordering::SeqCst).saturating_sub(held);
    assert_eq!(got, total, "every byte arrives");
    assert_eq!(
        order,
        [
            "evidence/sbom.json",
            "payload/a.tar",
            "payload/b.tar",
            "payload/c.tar",
            ""
        ],
        "member by member, then the summary"
    );
    assert_eq!(last_summary, Some(summary), "the summary comes last");
    assert!(
        extra <= 8 * EXPORT_CHUNK_BYTES,
        "streaming a {} MiB bundle took {} MiB more than the bundle itself: it must stay within a few chunks",
        total / MIB,
        extra / MIB
    );
    assert!(
        after_a < held - 32 * MIB,
        "a member is freed once it is sent: {} MiB live at payload/b.tar, {} MiB before",
        after_a / MIB,
        held / MIB
    );
}