armdb 0.14.1

//! Allocation gate for direct RPC range responses.
//!
//! `dispatch_into` drives `RowWriter` through the real registered `Adapter`.
//! The response buffer is warmed before the measured dispatch, so this detects
//! per-row intermediate collections or encoders rather than its one-time growth.
//!
//! This is intentionally scoped to `RapiraCodec`: `BitcodeCodec` and
//! `PostcardCodec` allocate inside their encoders by construction.

#![cfg(all(feature = "ctx-gates", feature = "rpc"))]

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

use armdb::armour::{Db, rpc::dispatch_into};
use armdb::{CollectionMeta, Config, NoHook, RapiraCodec, ZerocopyCodec};
use armour_core::GetType;
use armour_rpc::codec::encode_request;
use armour_rpc::{OpCode, Request, RequestPayload};
use rapira::Rapira;
use zerocopy::{FromBytes, Immutable, IntoBytes};

const ROWS: u32 = 1_000;

struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOCS.fetch_add(1, Ordering::Relaxed);
        unsafe { System.alloc(layout) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        ALLOCS.fetch_add(1, Ordering::Relaxed);
        unsafe { System.alloc_zeroed(layout) }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, size: usize) -> *mut u8 {
        ALLOCS.fetch_add(1, Ordering::Relaxed);
        unsafe { System.realloc(ptr, layout, size) }
    }
}

#[global_allocator]
static A: Counting = Counting;

#[derive(Clone, Debug, GetType, Rapira)]
struct AllocationRow {
    value: u64,
}

impl CollectionMeta for AllocationRow {
    type SelfId = [u8; 8];
    const NAME: &'static str = "rpc_alloc_gate_rows";
}

fn range_request() -> Vec<u8> {
    range_request_for(AllocationRow::NAME)
}

fn range_request_for(name: &str) -> Vec<u8> {
    let mut request = Vec::new();
    encode_request(
        Request {
            op: OpCode::Range,
            hashname: xxhash_rust::xxh3::xxh3_64(name.as_bytes()),
            payload: RequestPayload::Range {
                start: Bound::Unbounded,
                end: Bound::Unbounded,
                limit: 0,
            },
        },
        &mut request,
    );
    request
}

#[repr(C)]
#[derive(Clone, Copy, Debug, FromBytes, GetType, Immutable, IntoBytes)]
struct ZerocopyAllocationRow {
    value: u64,
}

impl CollectionMeta for ZerocopyAllocationRow {
    type SelfId = [u8; 8];
    const NAME: &'static str = "rpc_alloc_gate_zerocopy_rows";
}

fn database_with_rows() -> (tempfile::TempDir, Db) {
    let dir = tempfile::tempdir().expect("database directory");
    let db = Db::open_test(dir.path()).expect("open database");
    let tree = db
        .open_typed_tree::<AllocationRow, RapiraCodec, _>(Config::test(), NoHook, &[])
        .expect("open rows collection");
    for id in 0..u64::from(ROWS) {
        tree.put(&id.to_be_bytes(), AllocationRow { value: id })
            .expect("seed row");
    }
    (dir, db)
}

fn database_with_zerocopy_rows() -> (tempfile::TempDir, Db) {
    let dir = tempfile::tempdir().expect("database directory");
    let db = Db::open_test(dir.path()).expect("open database");
    let tree = db
        .open_typed_tree::<ZerocopyAllocationRow, ZerocopyCodec, _>(Config::test(), NoHook, &[])
        .expect("open zerocopy rows collection");
    for id in 0..u64::from(ROWS) {
        tree.put(&id.to_be_bytes(), ZerocopyAllocationRow { value: id })
            .expect("seed row");
    }
    (dir, db)
}

/// A temporary `Vec<(Vec<u8>, Vec<u8>)>` collector in `Adapter::range_into`
/// makes this fail with thousands of allocation calls.
#[test]
fn range_response_allocations_do_not_scale_with_rows() {
    let (_dir, db) = database_with_rows();
    let trees = db.build_tree_map();
    let request = range_request();
    let mut response = Vec::with_capacity(1 << 20);
    response.extend_from_slice(&[0; 4]);

    dispatch_into(&trees, &request, &mut response, 1 << 20);
    response.truncate(4);

    let before = ALLOCS.load(Ordering::Relaxed);
    dispatch_into(&trees, &request, &mut response, 1 << 20);
    let allocations = ALLOCS.load(Ordering::Relaxed) - before;

    assert_eq!(response[4], 0, "range response must succeed");
    assert_eq!(
        u32::from_be_bytes(response[5..9].try_into().expect("range count")),
        ROWS,
        "range response must include every seeded row"
    );
    assert!(
        allocations < 50,
        "{ROWS} rows cost {allocations} allocations; expected O(1) direct response emission"
    );
}

/// Zerocopy's `encode_append` writes the raw layout straight into RowWriter;
/// keep it under the same O(1) per-range allocation gate as Rapira.
#[test]
fn zerocopy_range_response_allocations_do_not_scale_with_rows() {
    let (_dir, db) = database_with_zerocopy_rows();
    let trees = db.build_tree_map();
    let request = range_request_for(ZerocopyAllocationRow::NAME);
    let mut response = Vec::with_capacity(1 << 20);
    response.extend_from_slice(&[0; 4]);

    dispatch_into(&trees, &request, &mut response, 1 << 20);
    response.truncate(4);

    let before = ALLOCS.load(Ordering::Relaxed);
    dispatch_into(&trees, &request, &mut response, 1 << 20);
    let allocations = ALLOCS.load(Ordering::Relaxed) - before;

    assert_eq!(response[4], 0, "range response must succeed");
    assert_eq!(
        u32::from_be_bytes(response[5..9].try_into().expect("range count")),
        ROWS,
        "range response must include every seeded row"
    );
    assert!(
        allocations < 50,
        "{ROWS} zerocopy rows cost {allocations} allocations; expected O(1) direct response emission"
    );
}