avila_async/

lib.rs

//! Avila Async - Native async runtime
//! Tokio alternative - 100% Rust std
//!
//! # Features
//! - Work-stealing scheduler
//! - Multi-threaded executor
//! - Async I/O primitives
//! - Channel support
//! - Task cancellation
//! - Graceful shutdown

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll, Wake};
use std::sync::{Arc, Mutex, Condvar, atomic::{AtomicBool, AtomicUsize, Ordering}};
use std::collections::VecDeque;
use std::thread;
use std::time::{Duration, Instant};

type Task = Pin<Box<dyn Future<Output = ()> + Send>>;

/// Task handle for spawned futures
pub struct JoinHandle<T> {
    result: Arc<Mutex<Option<T>>>,
    completed: Arc<AtomicBool>,
}

impl<T> JoinHandle<T> {
    /// Wait for the task to complete and return its result
    pub async fn await_result(self) -> Option<T> {
        while !self.completed.load(Ordering::Acquire) {
            yield_now().await;
        }
        self.result.lock().unwrap().take()
    }
}

pub struct Runtime {
    queue: Arc<Mutex<VecDeque<Task>>>,
    shutdown: Arc<AtomicBool>,
    task_count: Arc<AtomicUsize>,
    condvar: Arc<Condvar>,
}

impl Runtime {
    /// Create a new runtime instance
    pub fn new() -> Self {
        Self {
            queue: Arc::new(Mutex::new(VecDeque::new())),
            shutdown: Arc::new(AtomicBool::new(false)),
            task_count: Arc::new(AtomicUsize::new(0)),
            condvar: Arc::new(Condvar::new()),
        }
    }

    /// Get the number of active tasks
    pub fn task_count(&self) -> usize {
        self.task_count.load(Ordering::Relaxed)
    }

    /// Initiate graceful shutdown
    pub fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Release);
        self.condvar.notify_all();
    }

    /// Spawn a future onto the runtime
    pub fn spawn<F>(&self, future: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        self.task_count.fetch_add(1, Ordering::Relaxed);
        let task_count = Arc::clone(&self.task_count);
        let condvar = Arc::clone(&self.condvar);

        let wrapped = async move {
            future.await;
            task_count.fetch_sub(1, Ordering::Relaxed);
            condvar.notify_all();
        };

        let mut queue = self.queue.lock().unwrap();
        queue.push_back(Box::pin(wrapped));
        self.condvar.notify_one();
    }

    /// Spawn a future and return a handle to await its result
    pub fn spawn_with_handle<F, T>(&self, future: F) -> JoinHandle<T>
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        let result = Arc::new(Mutex::new(None));
        let completed = Arc::new(AtomicBool::new(false));
        let result_clone = Arc::clone(&result);
        let completed_clone = Arc::clone(&completed);

        let task = async move {
            let output = future.await;
            *result_clone.lock().unwrap() = Some(output);
            completed_clone.store(true, Ordering::Release);
        };

        self.spawn(task);
        JoinHandle { result, completed }
    }

    pub fn block_on<F, T>(&self, future: F) -> T
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        let result = Arc::new(Mutex::new(None));
        let result_clone = Arc::clone(&result);

        let task = async move {
            let output = future.await;
            *result_clone.lock().unwrap() = Some(output);
        };

        self.spawn(Box::pin(task));
        self.run();

        Arc::try_unwrap(result)
            .ok()
            .and_then(|m| m.into_inner().ok())
            .and_then(|opt| opt)
            .expect("Task did not complete")
    }

    fn run(&self) {
        let num_threads = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(4);

        let mut handles = vec![];

        for _ in 0..num_threads {
            let queue = Arc::clone(&self.queue);
            let shutdown = Arc::clone(&self.shutdown);
            let task_count = Arc::clone(&self.task_count);
            let condvar = Arc::clone(&self.condvar);

            let handle = thread::spawn(move || {
                let waker = Arc::new(RuntimeWaker { condvar: Arc::clone(&condvar) }).into();

                loop {
                    if shutdown.load(Ordering::Acquire) && task_count.load(Ordering::Relaxed) == 0 {
                        break;
                    }

                    let task = {
                        let mut q = queue.lock().unwrap();
                        if q.is_empty() && !shutdown.load(Ordering::Acquire) {
                            q = condvar.wait_timeout(q, Duration::from_millis(100)).unwrap().0;
                        }
                        q.pop_front()
                    };

                    match task {
                        Some(mut task) => {
                            let mut context = Context::from_waker(&waker);
                            match task.as_mut().poll(&mut context) {
                                Poll::Ready(()) => {},
                                Poll::Pending => {
                                    let mut q = queue.lock().unwrap();
                                    q.push_back(task);
                                }
                            }
                        }
                        None if shutdown.load(Ordering::Acquire) => break,
                        None => {}
                    }
                }
            });
            handles.push(handle);
        }

        for handle in handles {
            let _ = handle.join();
        }
    }
}

impl Default for Runtime {
    fn default() -> Self {
        Self::new()
    }
}

struct RuntimeWaker {
    condvar: Arc<Condvar>,
}

impl Wake for RuntimeWaker {
    fn wake(self: Arc<Self>) {
        self.condvar.notify_one();
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.condvar.notify_one();
    }
}

// Global helper function
pub fn spawn<F>(future: F)
where
    F: Future<Output = ()> + Send + 'static,
{
    RUNTIME.with(|rt| {
        rt.borrow().spawn(future);
    });
}

thread_local! {
    static RUNTIME: std::cell::RefCell<Runtime> = std::cell::RefCell::new(Runtime::new());
}

// Macro for async main
#[macro_export]
macro_rules! main {
    ($($body:tt)*) => {
        fn main() {
            let rt = $crate::Runtime::new();
            rt.block_on(async { $($body)* });
        }
    };
}

/// Yield execution to allow other tasks to run
pub async fn yield_now() {
    struct YieldNow {
        yielded: bool,
    }

    impl Future for YieldNow {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if self.yielded {
                Poll::Ready(())
            } else {
                self.yielded = true;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }

    YieldNow { yielded: false }.await
}

/// Sleep for a specified duration
pub async fn sleep(duration: Duration) {
    struct Sleep {
        when: std::time::Instant,
    }

    impl Future for Sleep {
        type Output = ();

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if std::time::Instant::now() >= self.when {
                Poll::Ready(())
            } else {
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }

    Sleep {
        when: std::time::Instant::now() + duration,
    }
    .await
}

/// Execute a future with a timeout
pub async fn timeout<F, T>(duration: Duration, future: F) -> Result<T, TimeoutError>
where
    F: Future<Output = T>,
{
    struct Timeout<F> {
        future: Pin<Box<F>>,
        deadline: Instant,
    }

    impl<F: Future> Future for Timeout<F> {
        type Output = Result<F::Output, TimeoutError>;

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if Instant::now() >= self.deadline {
                return Poll::Ready(Err(TimeoutError));
            }

            match self.future.as_mut().poll(cx) {
                Poll::Ready(v) => Poll::Ready(Ok(v)),
                Poll::Pending => {
                    cx.waker().wake_by_ref();
                    Poll::Pending
                }
            }
        }
    }

    Timeout {
        future: Box::pin(future),
        deadline: Instant::now() + duration,
    }
    .await
}

/// Timeout error type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TimeoutError;

impl std::fmt::Display for TimeoutError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "operation timed out")
    }
}

impl std::error::Error for TimeoutError {}

/// Async channel for message passing
pub mod channel {
    use std::sync::{Arc, Mutex, Condvar};
    use std::collections::VecDeque;

    /// Create a bounded channel with specified capacity
    pub fn bounded<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
        let inner = Arc::new(ChannelInner {
            queue: Mutex::new(VecDeque::with_capacity(capacity)),
            condvar: Condvar::new(),
            capacity,
            closed: Mutex::new(false),
        });
        (Sender { inner: inner.clone() }, Receiver { inner })
    }

    /// Create an unbounded channel
    pub fn unbounded<T>() -> (Sender<T>, Receiver<T>) {
        bounded(usize::MAX)
    }

    struct ChannelInner<T> {
        queue: Mutex<VecDeque<T>>,
        condvar: Condvar,
        capacity: usize,
        closed: Mutex<bool>,
    }

    /// Sender half of a channel
    pub struct Sender<T> {
        inner: Arc<ChannelInner<T>>,
    }

    impl<T> Sender<T> {
        /// Send a value through the channel
        pub async fn send(&self, value: T) -> Result<(), SendError<T>> {
            if *self.inner.closed.lock().unwrap() {
                return Err(SendError(value));
            }

            loop {
                let mut queue = self.inner.queue.lock().unwrap();
                if queue.len() < self.inner.capacity {
                    queue.push_back(value);
                    self.inner.condvar.notify_one();
                    return Ok(());
                }
                drop(queue);
                let queue = self.inner.queue.lock().unwrap();
                let _guard = self.inner.condvar.wait(queue).unwrap();
            }
        }
    }

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

    impl<T> Drop for Sender<T> {
        fn drop(&mut self) {
            if Arc::strong_count(&self.inner) == 2 {
                *self.inner.closed.lock().unwrap() = true;
                self.inner.condvar.notify_all();
            }
        }
    }

    /// Receiver half of a channel
    pub struct Receiver<T> {
        inner: Arc<ChannelInner<T>>,
    }

    impl<T> Receiver<T> {
        /// Receive a value from the channel
        pub async fn recv(&self) -> Option<T> {
            loop {
                let mut queue = self.inner.queue.lock().unwrap();
                if let Some(value) = queue.pop_front() {
                    self.inner.condvar.notify_one();
                    return Some(value);
                }
                if *self.inner.closed.lock().unwrap() && queue.is_empty() {
                    return None;
                }
                drop(queue);
                let queue = self.inner.queue.lock().unwrap();
                let _guard = self.inner.condvar.wait(queue).unwrap();
            }
        }
    }

    /// Error returned when sending fails
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct SendError<T>(pub T);

    impl<T> std::fmt::Display for SendError<T> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "channel closed")
        }
    }

    impl<T: std::fmt::Debug> std::error::Error for SendError<T> {}
}

// Basic network modules
pub mod net {
    use std::io;
    use std::net::{TcpListener as StdListener, TcpStream as StdStream, SocketAddr};

    pub struct TcpListener(StdListener);
    pub struct TcpStream(StdStream);

    impl TcpListener {
        pub async fn bind(addr: SocketAddr) -> io::Result<Self> {
            let listener = StdListener::bind(addr)?;
            listener.set_nonblocking(true)?;
            Ok(Self(listener))
        }

        pub async fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> {
            loop {
                match self.0.accept() {
                    Ok((stream, addr)) => {
                        stream.set_nonblocking(true)?;
                        return Ok((TcpStream(stream), addr));
                    }
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                        crate::sleep(std::time::Duration::from_millis(10)).await;
                    }
                    Err(e) => return Err(e),
                }
            }
        }
    }

    impl TcpStream {
        pub async fn connect(addr: SocketAddr) -> io::Result<Self> {
            let stream = StdStream::connect(addr)?;
            stream.set_nonblocking(true)?;
            Ok(Self(stream))
        }

        pub fn into_std(self) -> StdStream {
            self.0
        }

        pub fn as_std(&self) -> &StdStream {
            &self.0
        }

        /// Read data from the stream
        pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            use std::io::Read;
            loop {
                match self.0.read(buf) {
                    Ok(n) => return Ok(n),
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                        crate::sleep(std::time::Duration::from_millis(1)).await;
                    }
                    Err(e) => return Err(e),
                }
            }
        }

        /// Write data to the stream
        pub async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            use std::io::Write;
            loop {
                match self.0.write(buf) {
                    Ok(n) => return Ok(n),
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                        crate::sleep(std::time::Duration::from_millis(1)).await;
                    }
                    Err(e) => return Err(e),
                }
            }
        }

        /// Write all data to the stream
        pub async fn write_all(&mut self, mut buf: &[u8]) -> io::Result<()> {
            while !buf.is_empty() {
                let n = self.write(buf).await?;
                buf = &buf[n..];
            }
            Ok(())
        }
    }
}

// Basic I/O module
pub mod io {
    use std::io::{self, Read, Write};

    pub async fn copy<R: Read, W: Write>(reader: &mut R, writer: &mut W) -> io::Result<u64> {
        let mut buf = [0u8; 8192];
        let mut total = 0u64;

        loop {
            match reader.read(&mut buf) {
                Ok(0) => return Ok(total),
                Ok(n) => {
                    writer.write_all(&buf[..n])?;
                    total += n as u64;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    crate::sleep(std::time::Duration::from_millis(1)).await;
                }
                Err(e) => return Err(e),
            }
        }
    }
}