//! [![crates.io version](https://img.shields.io/crates/v/safina-executor.svg)](https://crates.io/crates/safina-executor)
//! [![license: Apache 2.0](https://gitlab.com/leonhard-llc/safina-rs/-/raw/main/license-apache-2.0.svg)](http://www.apache.org/licenses/LICENSE-2.0)
//! [![unsafe forbidden](https://gitlab.com/leonhard-llc/safina-rs/-/raw/main/unsafe-forbidden-success.svg)](https://github.com/rust-secure-code/safety-dance/)
//! [![pipeline status](https://gitlab.com/leonhard-llc/safina-rs/badges/main/pipeline.svg)](https://gitlab.com/leonhard-llc/safina-rs/-/pipelines)
//!
//! This is a safe Rust async executor.
//!
//! It is part of [`safina`](https://crates.io/crates/safina), a safe async runtime.
//!
//! # Features
//! - `forbid(unsafe_code)`
//! - Depends only on `std`
//! - Good test coverage (100%)
//!
//! # Limitations
//! - Requires Rust `nightly`, for
//!   [OnceCell](https://doc.rust-lang.org/std/lazy/struct.OnceCell.html)
//!   and
//!   [Wake trait](https://doc.rust-lang.org/std/task/trait.Wake.html)
//! - Allocates memory
//! - Not optimized
//!
//! # Documentation
//! https://docs.rs/safina-executor
//!
//! # Examples
//! ```rust
//! # fn f() {
//! safina_executor::increase_threads_to(1);
//! let (sender, receiver) = std::sync::mpsc::channel();
//! safina_executor::spawn(async move {
//!     sender.send(()).unwrap();
//! });
//! receiver.recv().unwrap();
//! # }
//! ```
//!
//! ```rust
//! # async fn prepare_request() -> Result<(), std::io::Error> { Ok(()) }
//! # async fn execute_request() -> Result<(), std::io::Error> { Ok(()) }
//! # fn f() -> Result<(), std::io::Error> {
//! std::thread::spawn(safina_executor::work);
//! let result = safina_executor::block_on(async {
//!     prepare_request().await?;
//!     execute_request().await
//! })?;
//! # Ok(())
//! # }
//! ```
//!
//! # Alternatives
//! - [smol](https://crates.io/crates/smol)
//!   - popular
//!   - Contains generous amounts of `unsafe` code
//! - [async-std](https://crates.io/crates/async-std)
//!   - popular
//!   - Contains generous amounts of `unsafe` code
//! - [tokio](https://crates.io/crates/tokio)
//!   - very popular
//!   - Fast, internally complicated, and full of `unsafe`
//! - [nostd_async](https://crates.io/crates/nostd_async)
//!
//! # Changelog
//! - v0.1.2 - Let callers pass futures to `spawn` and `block_on` without
//!   using `Box::pin`.
//!   Add `spawn_unpin` and `block_on_unpin` for folks who need to avoid allocating.
//!   so callers don't have to.
//! - v0.1.1 - Fix badge and update readme
//! - v0.1.0 - Renamed from `safina`
//!
//! # TO DO
//! - DONE - Implement `spawn`
//! - DONE - Implement `block_on`
//! - DONE - Implement `increase_threads_to`
//! - DONE - Drop finished futures
//! - DONE - Handle task panic
//! - DONE - Add `stop_threads`, `allow_threads`, and `increase_threads_to`.
//! - DONE - Add tests
//! - DONE - Add docs
//! - DONE - Publish on crates.io
//! - DONE - Add an #[async_test] macro
//! - Add a stress test
//! - Add a benchmark
//! - Make a version of the crate that uses
//!   unsafe [`once_cell`](https://crates.io/crates/once_cell)
//!   and unsafe [`RawWaker`](https://doc.rust-lang.org/std/task/struct.RawWaker.html)
//!   and builds with Rust `stable`.
//! - Add an #[async_main] macro
//!
//! # Release Process
//! 1. Edit `Cargo.toml` and bump version number.
//! 1. Run `./release.sh`
#![forbid(unsafe_code)]
#![feature(once_cell)]
#![feature(wake_trait)]

use std::error::Error;
use std::fmt::{Display, Formatter};
use std::future::Future;
use std::lazy::{SyncLazy, SyncOnceCell};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::mpsc::{Receiver, Sender, SyncSender};
use std::sync::{Arc, Mutex};
use std::task::Poll;
use std::time::{Duration, Instant};

static NUM_THREADS_RUNNING: AtomicU32 = AtomicU32::new(0);
static STOP_THREADS: AtomicBool = AtomicBool::new(false);
// Why isn't Mutex::new const?  See https://github.com/rust-lang/rust/issues/66806
static NUM_THREADS_SPAWNED: SyncLazy<Mutex<u32>> = SyncLazy::new(|| Mutex::new(0));

type TaskHandle = Arc<Mutex<dyn (Future<Output = ()>) + Send + Unpin>>;

static TASK_SENDER: SyncOnceCell<Mutex<Sender<TaskHandle>>> = SyncOnceCell::new();
static TASK_RECEIVER: SyncLazy<Mutex<Receiver<TaskHandle>>> = SyncLazy::new(|| {
    let (sender, receiver) = std::sync::mpsc::channel();
    TASK_SENDER.set(Mutex::new(sender)).unwrap();
    Mutex::new(receiver)
});

// TODO(mleonhard) Make a TaskWaker and its clones single-use.
struct TaskWaker {
    task_handle: TaskHandle,
}

impl std::task::Wake for TaskWaker {
    fn wake(self: Arc<Self>) {
        TASK_SENDER
            .get()
            .unwrap()
            .lock()
            .unwrap()
            .send(self.task_handle.clone())
            .unwrap();
    }
}

/// Returned by [`stop_threads`] when one or more worker threads fail to exit.
/// This can happen when a task takes too long between awaits.
#[derive(Debug)]
pub struct ThreadStopTimeout {}
impl Display for ThreadStopTimeout {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        std::fmt::Debug::fmt(self, f)
    }
}
impl Error for ThreadStopTimeout {}

/// The number of worker threads currently running.
///
/// The runtime does not automatically start or stop worker threads.
/// You must do that with
/// [`increase_threads_to`], [`work`], [`stop_threads`], and [`allow_threads`].
pub fn num_threads_running() -> u32 {
    NUM_THREADS_RUNNING.load(Ordering::Acquire)
}

/// Spawns new worker threads until the total number of workers is `n`.
///
/// Does nothing if the number of threads is already `n` or higher.
pub fn increase_threads_to(n: u32) {
    let mut num_spawned_guard = NUM_THREADS_SPAWNED.lock().unwrap();
    *num_spawned_guard = num_spawned_guard.max(NUM_THREADS_RUNNING.load(Ordering::Acquire));
    while *num_spawned_guard < n {
        std::thread::spawn(work);
        *num_spawned_guard += 1;
    }
}

/// Signals all threads to stop.  Waits for them all to stop.
///
/// If they don't all stop, gives up after 3 seconds and returns [`ThreadStopTimeout`].
pub fn stop_threads() -> Result<(), ThreadStopTimeout> {
    let mut num_spawned_guard = NUM_THREADS_SPAWNED.lock().unwrap();
    *num_spawned_guard = 0;
    STOP_THREADS.store(true, Ordering::Release);
    let deadline = Instant::now() + Duration::from_secs(3);
    loop {
        if NUM_THREADS_RUNNING.load(Ordering::Acquire) == 0 {
            return Ok(());
        }
        if Instant::now() > deadline {
            return Err(ThreadStopTimeout {});
        }
        std::thread::yield_now();
        // std::thread::sleep(Duration::from_millis(1));
    }
}

/// Clears the signal for worker threads to stop.
/// Call this after [`stop_threads`] and before [`increase_threads_to`] or [`work`].
pub fn allow_threads() {
    STOP_THREADS.store(false, Ordering::Release);
}

/// Adds a task that will execute `fut`.
/// The task runs on any available worker thread.
///
/// If no worker threads are running, the new task will not execute.
/// The task will execute once a worker thread begins work.
/// See [`increase_threads_to`] and [`work`].
///
/// Returns immediately.
///
/// Uses
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make the future
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use [`spawn_unpin`](#method.spawn_unpin) to avoid this allocation.
///
/// Example:
/// ```rust
/// # use std::time::Duration;
/// # async fn an_async_fn() -> Result<(), std::io::Error> { Ok(()) }
/// # fn f() {
/// safina_executor::increase_threads_to(1);
/// safina_executor::spawn(async move {
///     an_async_fn().await.unwrap();
/// });
/// # }
/// ```
pub fn spawn(fut: impl (Future<Output = ()>) + Send + 'static) {
    spawn_unpin(Box::pin(fut))
}

/// Adds a task that will execute `fut`.
/// The task runs on any available worker thread.
///
/// If no worker threads are running, the new task will not execute.
/// The task will execute once a worker thread begins work.
/// See [`increase_threads_to`] and [`work`].
///
/// Returns immediately.
///
/// Note that `fut` must be
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make it Unpin.  The [`spawn`](#method.spawn) function does this for you.
/// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
/// to do it with unsafe code that does not allocate memory.
pub fn spawn_unpin(fut: impl (Future<Output = ()>) + Send + Unpin + 'static) {
    // Create channel.  This avoids racing worker thread on startup.
    SyncLazy::force(&TASK_RECEIVER);
    TASK_SENDER
        .get()
        .unwrap()
        .lock()
        .unwrap()
        .send(Arc::new(Mutex::new(fut)))
        .unwrap();
}

/// Use the current thread as a worker in the runtime.
///
/// Call [`stop_threads`] to make the thread return from this call.
pub fn work() {
    NUM_THREADS_RUNNING.fetch_add(1, Ordering::AcqRel);
    while !STOP_THREADS.load(Ordering::Acquire) {
        let result_task_handle = TASK_RECEIVER
            .lock()
            .unwrap()
            .recv_timeout(Duration::from_millis(100));
        if let Ok(task_handle) = result_task_handle {
            let waker = std::task::Waker::from(Arc::new(TaskWaker {
                task_handle: task_handle.clone(),
            }));
            let _ = std::panic::catch_unwind(|| {
                let mut cx = std::task::Context::from_waker(&waker);
                let mut fut_guard = task_handle.lock().unwrap();
                Pin::new(&mut *fut_guard).poll(&mut cx)
            });
        }
    }
    NUM_THREADS_RUNNING.fetch_sub(1, Ordering::AcqRel);
}

/// Executes the future on the current thread and returns its result.
/// Panics if the future panics.
///
/// `fut` can call [`spawn`] to create tasks.
/// Those tasks will execute on worker threads even after `fut` completes and this thread returns.
///
/// Uses
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make the future
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use [`block_on_unpin`](#method.block_on_unpin) to avoid this allocation.
///
/// ```rust
/// # use std::time::{Duration, Instant};
/// # async fn prepare_request() -> Result<(), std::io::Error> { Ok(()) }
/// # async fn execute_request() -> Result<(), std::io::Error> { Ok(()) }
/// # fn f() -> Result<(), std::io::Error> {
/// std::thread::spawn(safina_executor::work);
/// let result = safina_executor::block_on(async {
///     prepare_request().await?;
///     execute_request().await
/// })?;
/// # Ok(())
/// # }
/// ```
pub fn block_on<R>(fut: impl (Future<Output = R>) + Send + 'static) -> R {
    block_on_unpin(Box::pin(fut))
}

/// Executes the future on the current thread and returns its result.
/// Panics if the future panics.
///
/// `fut` can call [`spawn`] to create tasks.
/// Those tasks will execute on worker threads even after `fut` completes and this thread returns.
///
/// Note that `fut` must be
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make it Unpin.  The [`block_on`](#method.block_on) function does this for you.
/// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
/// to do it with unsafe code that does not allocate memory.
pub fn block_on_unpin<R>(mut fut: impl (Future<Output = R>) + Send + Unpin + 'static) -> R {
    struct BlockOnTaskWaker(SyncSender<()>);
    impl std::task::Wake for BlockOnTaskWaker {
        fn wake(self: Arc<Self>) {
            let _ = self.0.send(());
        }
    }
    let (sender, receiver) = std::sync::mpsc::sync_channel(1);
    let waker = std::task::Waker::from(Arc::new(BlockOnTaskWaker(sender)));
    let mut cx = std::task::Context::from_waker(&waker);
    loop {
        if let Poll::Ready(result) = Pin::new(&mut fut).poll(&mut cx) {
            return result;
        }
        receiver.recv().unwrap();
    }
}

#[cfg(test)]
mod tests;