memkit-async 0.2.0-beta.1

//! Tokio runtime integration for memkit-async.
//!
//! Provides task-local allocation context and spawn helpers.

use std::future::Future;
use std::sync::Arc;

use crate::allocator::MkAsyncFrameAlloc;

tokio::task_local! {
    /// Task-local allocator context.
    static TASK_ALLOC: Arc<MkAsyncFrameAlloc>;
}

/// Enter a task-local allocation context.
///
/// All allocations made within the future will use this allocator.
///
/// # Example
///
/// ```rust,ignore
/// use memkit_async::{MkAsyncFrameAlloc, MkAsyncFrameConfig};
/// use memkit_async::runtime::with_allocator;
///
/// #[tokio::main]
/// async fn main() {
///     let alloc = MkAsyncFrameAlloc::new(MkAsyncFrameConfig::default());
///     
///     with_allocator(alloc, async {
///         let ptr = current_allocator().alloc::<u32>().await;
///         // ...
///     }).await;
/// }
/// ```
pub async fn with_allocator<F, R>(alloc: MkAsyncFrameAlloc, f: F) -> R
where
    F: Future<Output = R>,
{
    TASK_ALLOC.scope(Arc::new(alloc), f).await
}

/// Get the current task-local allocator.
///
/// Panics if called outside of a `with_allocator` context.
pub fn current_allocator() -> Arc<MkAsyncFrameAlloc> {
    TASK_ALLOC.with(|a| Arc::clone(a))
}

/// Try to get the current task-local allocator.
///
/// Returns `None` if called outside of a `with_allocator` context.
pub fn try_current_allocator() -> Option<Arc<MkAsyncFrameAlloc>> {
    TASK_ALLOC.try_with(|a| Arc::clone(a)).ok()
}

/// Spawn a task with the current allocation context propagated.
///
/// The spawned task will have access to the same allocator as the parent.
///
/// # Example
///
/// ```rust,ignore
/// use memkit_async::runtime::{with_allocator, mk_spawn};
///
/// with_allocator(alloc, async {
///     mk_spawn(async {
///         // This task has access to the same allocator
///         let ptr = current_allocator().alloc::<u32>().await;
///     });
/// }).await;
/// ```
pub fn mk_spawn<F>(f: F) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    let alloc = try_current_allocator();
    
    tokio::spawn(async move {
        if let Some(alloc) = alloc {
            with_allocator((*alloc).clone(), f).await
        } else {
            f.await
        }
    })
}

/// Spawn a blocking task with allocation context.
pub fn mk_spawn_blocking<F, R>(f: F) -> tokio::task::JoinHandle<R>
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    tokio::task::spawn_blocking(f)
}

/// Run a future with a timeout, using Tokio's timer.
pub async fn with_timeout<F, R>(
    duration: std::time::Duration,
    f: F,
) -> Result<R, tokio::time::error::Elapsed>
where
    F: Future<Output = R>,
{
    tokio::time::timeout(duration, f).await
}

/// Yield to the Tokio runtime.
pub async fn yield_now() {
    tokio::task::yield_now().await
}

/// Sleep for the specified duration.
pub async fn sleep(duration: std::time::Duration) {
    tokio::time::sleep(duration).await
}

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

    #[tokio::test]
    async fn test_task_local_allocator() {
        let alloc = MkAsyncFrameAlloc::new(MkAsyncFrameConfig::default());
        
        with_allocator(alloc, async {
            let a = current_allocator();
            assert_eq!(a.frame(), 0);
        }).await;
    }

    #[tokio::test]
    async fn test_spawn_propagates_context() {
        let alloc = MkAsyncFrameAlloc::new(MkAsyncFrameConfig::default());
        
        with_allocator(alloc, async {
            let handle = mk_spawn(async {
                // Should have access to allocator
                let a = current_allocator();
                a.frame()
            });
            
            let frame = handle.await.unwrap();
            assert_eq!(frame, 0);
        }).await;
    }

    #[tokio::test]
    async fn test_timeout() {
        let result = with_timeout(
            std::time::Duration::from_millis(100),
            async { 42 }
        ).await;
        
        assert_eq!(result.unwrap(), 42);
    }
}