memkit-async 0.2.0-beta.1

//! Async object pool.

use crate::backpressure::MkBackpressure;
use std::collections::VecDeque;
use std::future::Future;
use std::ops::{Deref, DerefMut};
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

/// Async-aware object pool with backpressure support.
pub struct MkAsyncPool<T> {
    inner: Arc<PoolInner<T>>,
}

struct PoolInner<T> {
    items: Mutex<VecDeque<T>>,
    capacity: usize,
    available: AtomicUsize,
    backpressure: MkBackpressure,
}

impl<T> MkAsyncPool<T> {
    /// Create a new async pool.
    pub fn new(capacity: usize, backpressure: MkBackpressure) -> Self {
        Self {
            inner: Arc::new(PoolInner {
                items: Mutex::new(VecDeque::with_capacity(capacity)),
                capacity,
                available: AtomicUsize::new(0),
                backpressure,
            }),
        }
    }

    /// Add an item to the pool.
    pub fn add(&self, item: T) -> Result<(), T> {
        let mut items = self.inner.items.lock().unwrap();
        if items.len() >= self.inner.capacity {
            return Err(item);
        }
        items.push_back(item);
        self.inner.available.fetch_add(1, Ordering::Release);
        Ok(())
    }

    /// Acquire an item from the pool.
    pub async fn acquire(&self) -> Option<MkPoolGuard<T>> {
        loop {
            // Try to get an item
            if let Some(guard) = self.try_acquire() {
                return Some(guard);
            }

            // Apply backpressure
            match self.inner.backpressure {
                MkBackpressure::Fail => return None,
                MkBackpressure::Wait => {
                    YieldOnce::new().await;
                    continue;
                }
                MkBackpressure::Timeout(duration) => {
                    // Simple timeout - just fail for now
                    let _ = duration;
                    return None;
                }
                MkBackpressure::Evict => {
                    return None;
                }
            }
        }
    }

    /// Try to acquire an item without waiting.
    pub fn try_acquire(&self) -> Option<MkPoolGuard<T>> {
        let mut items = self.inner.items.lock().unwrap();
        if let Some(item) = items.pop_front() {
            self.inner.available.fetch_sub(1, Ordering::Acquire);
            Some(MkPoolGuard {
                item: Some(item),
                pool: Arc::clone(&self.inner),
            })
        } else {
            None
        }
    }

    /// Get the pool capacity.
    pub fn capacity(&self) -> usize {
        self.inner.capacity
    }

    /// Get the number of available items.
    pub fn available(&self) -> usize {
        self.inner.available.load(Ordering::Relaxed)
    }
}

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

/// Guard for a pooled item.
pub struct MkPoolGuard<T> {
    item: Option<T>,
    pool: Arc<PoolInner<T>>,
}

impl<T> Deref for MkPoolGuard<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.item.as_ref().unwrap()
    }
}

impl<T> DerefMut for MkPoolGuard<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.item.as_mut().unwrap()
    }
}

impl<T> Drop for MkPoolGuard<T> {
    fn drop(&mut self) {
        if let Some(item) = self.item.take() {
            let mut items = self.pool.items.lock().unwrap();
            if items.len() < self.pool.capacity {
                items.push_back(item);
                self.pool.available.fetch_add(1, Ordering::Release);
            }
        }
    }
}

/// Yield once to the runtime.
struct YieldOnce(bool);

impl YieldOnce {
    fn new() -> Self {
        Self(false)
    }
}

impl Future for YieldOnce {
    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
        }
    }
}

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

    #[test]
    fn test_pool_sync() {
        let pool: MkAsyncPool<u32> = MkAsyncPool::new(3, MkBackpressure::Fail);
        
        pool.add(1).unwrap();
        pool.add(2).unwrap();
        pool.add(3).unwrap();
        assert!(pool.add(4).is_err());
        
        assert_eq!(pool.available(), 3);
        
        let guard = pool.try_acquire().unwrap();
        assert_eq!(*guard, 1);
        assert_eq!(pool.available(), 2);
        
        drop(guard);
        assert_eq!(pool.available(), 3);
    }
}