memkit_async/

allocator.rs

//! Async-aware frame allocator.

use std::alloc::Layout;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::backpressure::MkBackpressure;

/// Configuration for async frame allocator.
#[derive(Debug, Clone)]
pub struct MkAsyncFrameConfig {
    /// Size of each frame arena.
    pub arena_size: usize,
    /// Backpressure policy when memory is exhausted.
    pub backpressure: MkBackpressure,
}

impl Default for MkAsyncFrameConfig {
    fn default() -> Self {
        Self {
            arena_size: 16 * 1024 * 1024, // 16 MB
            backpressure: MkBackpressure::Wait,
        }
    }
}

/// Async-aware frame allocator.
///
/// Tracks allocations per-task rather than per-thread,
/// making it safe to use across `.await` points.
pub struct MkAsyncFrameAlloc {
    inner: Arc<AsyncAllocInner>,
}

struct AsyncAllocInner {
    config: MkAsyncFrameConfig,
    /// Memory arena
    arena: Vec<u8>,
    /// Current allocation head
    head: AtomicUsize,
    /// Frame counter
    frame: AtomicU64,
    /// Active allocations count (prevents reset while in use)
    active_count: AtomicUsize,
}

impl MkAsyncFrameAlloc {
    /// Create a new async frame allocator.
    pub fn new(config: MkAsyncFrameConfig) -> Self {
        let arena = vec![0u8; config.arena_size];
        Self {
            inner: Arc::new(AsyncAllocInner {
                config,
                arena,
                head: AtomicUsize::new(0),
                frame: AtomicU64::new(0),
                active_count: AtomicUsize::new(0),
            }),
        }
    }

    /// Begin a new async frame.
    pub async fn begin_frame(&self) -> MkAsyncFrameGuard {
        // Wait for all active allocations to be released
        WaitForEmpty::new(&self.inner.active_count).await;
        
        // Increment frame and reset head
        self.inner.frame.fetch_add(1, Ordering::SeqCst);
        self.inner.head.store(0, Ordering::SeqCst);
        
        MkAsyncFrameGuard {
            alloc: self.clone(),
        }
    }

    /// Get the current frame number.
    pub fn frame(&self) -> u64 {
        self.inner.frame.load(Ordering::SeqCst)
    }

    /// Allocate memory asynchronously.
    pub async fn alloc<T>(&self) -> Option<*mut T> {
        let layout = Layout::new::<T>();
        self.alloc_layout(layout).await.map(|p| p as *mut T)
    }

    /// Allocate with layout.
    pub async fn alloc_layout(&self, layout: Layout) -> Option<*mut u8> {
        let size = layout.size();
        let align = layout.align();
        
        loop {
            let current = self.inner.head.load(Ordering::Acquire);
            let aligned = (current + align - 1) & !(align - 1);
            
            if aligned + size > self.inner.arena.len() {
                // Out of memory - apply backpressure
                match self.inner.config.backpressure {
                    MkBackpressure::Fail => return None,
                    MkBackpressure::Wait => {
                        // Yield and retry
                        tokio_yield().await;
                        continue;
                    }
                    MkBackpressure::Timeout(duration) => {
                        // TODO: Implement timeout
                        let _ = duration;
                        return None;
                    }
                    MkBackpressure::Evict => {
                        // TODO: Implement eviction
                        return None;
                    }
                }
            }
            
            // Try to bump the head
            match self.inner.head.compare_exchange_weak(
                current,
                aligned + size,
                Ordering::AcqRel,
                Ordering::Relaxed,
            ) {
                Ok(_) => {
                    self.inner.active_count.fetch_add(1, Ordering::Relaxed);
                    let ptr = self.inner.arena.as_ptr() as *mut u8;
                    return Some(unsafe { ptr.add(aligned) });
                }
                Err(_) => continue,
            }
        }
    }

    /// Release an allocation (for tracking).
    pub fn release(&self) {
        self.inner.active_count.fetch_sub(1, Ordering::Relaxed);
    }

    /// Get current memory usage.
    pub fn used(&self) -> usize {
        self.inner.head.load(Ordering::Relaxed)
    }

    /// Get remaining capacity.
    pub fn remaining(&self) -> usize {
        self.inner.arena.len() - self.used()
    }
}

impl Clone for MkAsyncFrameAlloc {
    fn clone(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
        }
    }
}

/// Guard that represents an active frame.
pub struct MkAsyncFrameGuard {
    alloc: MkAsyncFrameAlloc,
}

impl MkAsyncFrameGuard {
    /// Get the frame number.
    pub fn frame(&self) -> u64 {
        self.alloc.frame()
    }
}

impl Drop for MkAsyncFrameGuard {
    fn drop(&mut self) {
        // Reset head when frame ends
        self.alloc.inner.head.store(0, Ordering::SeqCst);
    }
}

/// Future that waits for active count to reach zero.
struct WaitForEmpty<'a> {
    counter: &'a AtomicUsize,
}

impl<'a> WaitForEmpty<'a> {
    fn new(counter: &'a AtomicUsize) -> Self {
        Self { counter }
    }
}

impl<'a> Future for WaitForEmpty<'a> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.counter.load(Ordering::Acquire) == 0 {
            Poll::Ready(())
        } else {
            // Wake immediately to retry
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

/// Yield to the async runtime.
async fn tokio_yield() {
    struct Yield(bool);
    
    impl Future for Yield {
        type Output = ();
        
        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if self.0 {
                Poll::Ready(())
            } else {
                self.0 = true;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }
    
    Yield(false).await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_async_frame_alloc_sync() {
        let alloc = MkAsyncFrameAlloc::new(MkAsyncFrameConfig::default());
        assert_eq!(alloc.frame(), 0);
        assert_eq!(alloc.used(), 0);
    }
}