supertunnel 0.1.0

//! Asynchronous Mutexes.
//!
//! Defines polling-based synchronization primitives designed for use within
//! asynchronous state machines and protocol codecs.

use std::cell::UnsafeCell;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;
use std::task::{Context, Poll};

use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio_util::sync::PollSemaphore;

/// An asynchronous mutex designed specifically for `poll_*` based state machines.
///
/// Provides fine-grained waker management for components that cannot use standard `async fn` locks.
pub struct PollMutex<T> {
    mutex: PollSemaphore,
    data: Arc<UnsafeCell<T>>,
}

/// An RAII guard representing an acquired lock on a [`PollMutex`].
///
/// Provides exclusive access to the underlying protected data via [`Deref`] and [`DerefMut`].
/// When this guard is dropped, the lock is automatically released, making the [`PollMutex`]
/// available for other tasks.
pub struct PollMutexGuard<T> {
    // RAII: as long as we hold this permit, we are allowed exclusive access to
    // the data. On drop, the permit is released back to the semaphore.
    #[allow(unused)]
    permit: OwnedSemaphorePermit,
    data: Arc<UnsafeCell<T>>,
}

impl<T> PollMutex<T> {
    /// Creates a new [`PollMutex`], using a [`PollSemaphore`] and wrapping data in an [`UnsafeCell`].
    ///
    /// # Inputs
    ///
    /// * `data: T` - The data to be protected by the mutex.
    pub fn new(data: T) -> Self {
        Self {
            mutex: PollSemaphore::new(Arc::new(Semaphore::new(1))),
            data: Arc::new(UnsafeCell::new(data)),
        }
    }

    /// Asynchronously acquires the mutex lock.
    ///
    /// Returns a future that resolves to a [`PollMutexGuard`] when the lock is acquired.
    pub fn lock(&mut self) -> impl Future<Output = PollMutexGuard<T>> {
        std::future::poll_fn(move |cx| self.poll_lock(cx))
    }

    /// Attempts to acquire the lock immediately without waiting.
    ///
    /// Unlike `lock` or `poll_lock`, this method does not register a waker if the lock
    /// is unavailable. It is useful for opportunistic locking where the task can gracefully
    /// handle contention without needing to be woken up later.
    ///
    /// # Returns
    ///
    /// * `Some(PollMutexGuard<T>)` - If the lock was successfully acquired.
    /// * `None` - If the mutex is currently locked by another task.
    pub fn try_lock(&self) -> Option<PollMutexGuard<T>> {
        if let Ok(permit) = self.mutex.clone_inner().try_acquire_owned() {
            Some(PollMutexGuard {
                permit,
                data: self.data.clone(),
            })
        } else {
            None
        }
    }

    /// Polls for the mutex lock to become available.
    ///
    /// This is the low-level polling equivalent of `lock()`. If the lock is not available,
    /// the current task's waker is registered.
    ///
    /// When this method returns Poll::Pending, the current task is scheduled to
    /// receive a wakeup when the mutex is released. Note that on multiple calls
    /// to `poll_lock()`, only the Waker from the Context passed to the most recent
    /// call is scheduled to receive a wakeup.
    ///
    /// # Returns
    ///
    /// * `Poll<PollMutexGuard<T>>` - `Poll::Ready(guard)` if the lock is acquired, otherwise `Poll::Pending`.
    pub fn poll_lock(&mut self, cx: &mut Context) -> Poll<PollMutexGuard<T>> {
        match self.mutex.poll_acquire(cx) {
            Poll::Ready(Some(permit)) => Poll::Ready(PollMutexGuard {
                permit,
                data: self.data.clone(),
            }),
            // We never close the semaphore, and self has a ref to it so it
            // could not have closed itself on drop yet.
            Poll::Ready(None) => unreachable!(),
            Poll::Pending => Poll::Pending,
        }
    }
}

// Allow PollMutex to be sent across threads if T is Send
unsafe impl<T: Send> Send for PollMutex<T> {}
unsafe impl<T: Send> Sync for PollMutex<T> {}
unsafe impl<T: Send + Sync> Sync for PollMutexGuard<T> {}

impl<T> From<T> for PollMutex<T> {
    fn from(data: T) -> Self {
        Self::new(data)
    }
}

impl<T> Clone for PollMutex<T> {
    fn clone(&self) -> Self {
        Self {
            mutex: self.mutex.clone(),
            data: self.data.clone(),
        }
    }
}

impl<T: Default> Default for PollMutex<T> {
    fn default() -> Self {
        Self::new(T::default())
    }
}

impl<T> Deref for PollMutexGuard<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        // SAFETY: The existence of the PollMutexGuard guarantees we have
        // exclusive access, so we can safely get a shared reference to the
        // data.
        unsafe { &*self.data.get() }
    }
}

impl<T> DerefMut for PollMutexGuard<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: The existence of the PollMutexGuard guarantees we have
        // exclusive access, so we can safely get a mutable reference to the
        // data.
        unsafe { &mut *self.data.get() }
    }
}

#[cfg(test)]
mod tests {
    use crate::sync::PollMutex;

    async fn push_one(mut mutex: PollMutex<Vec<usize>>) {
        mutex.lock().await.push(1);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 1000)]
    async fn poll_mutex_parallel_lock() {
        let n = 1000;
        let mut mutex = PollMutex::new(vec![]);

        let mut handles = vec![];

        for _ in 0..n {
            handles.push(tokio::spawn(push_one(mutex.clone())));
        }

        assert_eq!(handles.len(), n);

        for handle in handles {
            let _ = handle.await;
        }

        let inner = mutex.lock().await;
        assert_eq!(inner.len(), n);
        let sum: usize = inner.iter().sum();
        assert_eq!(sum, n);
    }
}