armdb 0.11.2

sharded bitcask key-value storage optimized for NVMe
Documentation
//! Allocation gate for the encrypted read path.
//!
//! `PageCipher::decrypt_page` runs once per 4 KiB page of every encrypted read,
//! so an allocation there is paid on every block and every value fetched from an
//! encrypted collection (F10 of libfunc/armour#83). This binary installs its own
//! counting global allocator — it cannot live in the lib test binary, which would
//! then measure all ~1300 unit tests through it.
//!
//! Run: `cargo nextest run -p armdb --features encryption --test crypto_alloc_gate`

#![cfg(feature = "encryption")]

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

use armdb::PageCipher;

struct CountingAlloc;

static ALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for CountingAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOC_COUNT.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 {
        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
        unsafe { System.alloc_zeroed(layout) }
    }

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

#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;

const PAGE_SIZE: usize = 4096;
const PAGES: u64 = 64;

/// F10: decrypting a page must not allocate — it is handed a page-sized buffer
/// and a tag, and `ring` can authenticate them in place.
#[test]
fn decrypt_page_allocates_nothing() {
    let cipher = PageCipher::new(&[9u8; 32]).expect("cipher");
    let file_id = 11u32;

    // Prepare ciphertext + tags outside the measured window.
    let mut pages: Vec<[u8; PAGE_SIZE]> = Vec::with_capacity(PAGES as usize);
    let mut tags = Vec::with_capacity(PAGES as usize);
    for p in 0..PAGES {
        let mut page = [0u8; PAGE_SIZE];
        for (i, b) in page.iter_mut().enumerate() {
            *b = (i as u64 + p) as u8;
        }
        let tag = cipher.encrypt_page(file_id, p, &mut page).expect("encrypt");
        pages.push(page);
        tags.push(tag);
    }

    let before = ALLOC_COUNT.load(Ordering::Relaxed);
    for (p, (page, tag)) in pages.iter_mut().zip(tags.iter()).enumerate() {
        cipher
            .decrypt_page(file_id, p as u64, page.as_mut_slice(), tag)
            .expect("decrypt");
    }
    let allocs = ALLOC_COUNT.load(Ordering::Relaxed) - before;

    assert_eq!(
        allocs, 0,
        "{PAGES} page decryptions allocated {allocs} times — the encrypted read \
         path pays this per 4 KiB page"
    );

    // The round trip still has to be correct, not just cheap.
    for (p, page) in pages.iter().enumerate() {
        for (i, b) in page.iter().enumerate() {
            assert_eq!(*b, (i as u64 + p as u64) as u8, "page {p} byte {i}");
        }
    }
}