callback-result 0.2.5

use callback_result::{CallbackWaiter, WaiterError};
use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicIsize, Ordering};
use std::time::Duration;

struct CountingAllocator;

static LIVE_BYTES: AtomicIsize = AtomicIsize::new(0);

unsafe impl GlobalAlloc for CountingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { System.alloc(layout) };
        if !ptr.is_null() {
            LIVE_BYTES.fetch_add(layout.size() as isize, Ordering::SeqCst);
        }
        ptr
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { System.alloc_zeroed(layout) };
        if !ptr.is_null() {
            LIVE_BYTES.fetch_add(layout.size() as isize, Ordering::SeqCst);
        }
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        LIVE_BYTES.fetch_sub(layout.size() as isize, Ordering::SeqCst);
        unsafe { System.dealloc(ptr, layout) };
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let new_ptr = unsafe { System.realloc(ptr, layout, new_size) };
        if !new_ptr.is_null() {
            LIVE_BYTES.fetch_add(new_size as isize - layout.size() as isize, Ordering::SeqCst);
        }
        new_ptr
    }
}

#[global_allocator]
static ALLOCATOR: CountingAllocator = CountingAllocator;

#[test]
fn dropped_unique_callback_futures_do_not_accumulate_waiter_entries() {
    let waiter = CallbackWaiter::<u128, u8>::new();

    drop(waiter.create_result_future(0).unwrap());
    let baseline = LIVE_BYTES.load(Ordering::SeqCst);

    for callback_id in 1..=4096 {
        if callback_id % 2 == 0 {
            drop(waiter.create_result_future(callback_id).unwrap());
        } else {
            drop(
                waiter
                    .create_timeout_result_future(callback_id, Duration::from_secs(60))
                    .unwrap(),
            );
        }
    }

    let retained = LIVE_BYTES.load(Ordering::SeqCst) - baseline;
    assert!(
        retained < 64 * 1024,
        "dropped callback futures retained {retained} live bytes"
    );

    assert_eq!(waiter.set_result(4096, 7), Err(WaiterError::NoWaiter));
    assert!(waiter.create_result_future(4096).is_ok());
}

#[tokio::test]
async fn completed_callback_is_delivered_and_removed() {
    let waiter = CallbackWaiter::<u128, u8>::new();
    let future = waiter.create_result_future(7).unwrap();

    waiter.set_result(7, 9).unwrap();
    assert_eq!(future.await.unwrap(), 9);
    assert_eq!(waiter.set_result(7, 10), Err(WaiterError::NoWaiter));
}

#[tokio::test]
async fn timed_out_callback_rejects_late_result() {
    let waiter = CallbackWaiter::<u128, u8>::new();
    let future = waiter
        .create_timeout_result_future(11, Duration::from_millis(1))
        .unwrap();

    assert_eq!(future.await, Err(WaiterError::Timeout));
    assert_eq!(waiter.set_result(11, 3), Err(WaiterError::NoWaiter));
}