1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
//! [![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(Box::pin(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(Box::pin(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)
//! # Release Process
//! 1. Edit `Cargo.toml` and bump version number.
//! 1. Run `./release.sh`
//!
//! # Changelog
//! - 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
#![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.
///
/// This runtime does not unsafely convert your futures to [`Pin`][std::pin::Pin].
/// You must wrap your future in [`Box::pin`](std::box::Box::pin).
///
/// 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(Box::pin(async move {
///     an_async_fn().await.unwrap();
/// }));
/// # }
/// ```
pub fn spawn(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.
///
/// ```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(Box::pin(async {
///     prepare_request().await?;
///     execute_request().await
/// }))?;
/// # Ok(())
/// # }
/// ```
pub fn block_on<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;