vyre-conform 0.1.0

//! Counting global allocator for SQLite-style OOM injection.
#![allow(unsafe_code)]
#![allow(clippy::panic)]

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
use std::sync::atomic::{AtomicBool, Ordering};

const FAIL_DISABLED: usize = 0;
const OOM_SENTINEL: &str = "vyre-conform oom injection";

static ENABLED: AtomicBool = AtomicBool::new(false);

thread_local! {
    static COUNTER: Cell<usize> = const { Cell::new(0) };
    static FAIL_AT: Cell<usize> = const { Cell::new(FAIL_DISABLED) };
    static TRIGGERED: Cell<usize> = const { Cell::new(0) };
}

/// Global allocator that counts allocations and can fail one selected call.
pub struct OomAllocator;

/// Arm OOM accounting for this thread.
///
/// `fail_at = 0` records allocation count without injecting failure.
#[inline]
pub fn arm_thread(fail_at: usize) {
    COUNTER.with(|counter| counter.set(0));
    FAIL_AT.with(|cell| cell.set(fail_at));
    TRIGGERED.with(|cell| cell.set(0));
    ENABLED.store(true, Ordering::SeqCst);
}

/// Disarm OOM accounting and failure injection.
#[inline]
pub fn disarm_thread() {
    ENABLED.store(false, Ordering::SeqCst);
    FAIL_AT.with(|cell| cell.set(FAIL_DISABLED));
}

/// Reset thread-local counters after a probe has reported them.
#[inline]
pub fn clear_thread() {
    COUNTER.with(|counter| counter.set(0));
    FAIL_AT.with(|cell| cell.set(FAIL_DISABLED));
    TRIGGERED.with(|cell| cell.set(0));
}

/// Number of allocations observed on this thread since the last arm.
#[inline]
pub fn allocation_count() -> usize {
    COUNTER.with(Cell::get)
}

/// Allocation index that triggered failure, or zero if none triggered.
#[inline]
pub fn triggered_at() -> usize {
    TRIGGERED.with(Cell::get)
}

/// True if a panic payload came from the OOM allocator.
#[inline]
pub fn is_oom_payload(payload: &(dyn std::any::Any + Send)) -> bool {
    payload
        .downcast_ref::<&'static str>()
        .map(|message| *message == OOM_SENTINEL)
        .unwrap_or(false)
}

// SAFETY: This implementation forwards every real allocation and deallocation
// to `System` with the exact pointer/layout pairs supplied by the standard
// library. While injection is enabled it may panic before calling `System`,
// which means no pointer has been created and no deallocation obligation exists.
unsafe impl GlobalAlloc for OomAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if !ENABLED.load(Ordering::SeqCst) {
            return System.alloc(layout);
        }

        let should_fail = COUNTER.with(|counter| {
            let next = counter.get().saturating_add(1);
            counter.set(next);
            FAIL_AT.with(|fail_at| fail_at.get() != FAIL_DISABLED && next == fail_at.get())
        });

        if should_fail {
            TRIGGERED.with(|triggered| triggered.set(allocation_count()));
            ENABLED.store(false, Ordering::SeqCst);
            panic!("{OOM_SENTINEL}");
        }

        System.alloc(layout)
    }

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

#[global_allocator]
/// Process-wide allocator used by the OOM injection harness.
pub static GLOBAL_ALLOCATOR: OomAllocator = OomAllocator;