use std::panic::{self, AssertUnwindSafe};
use std::sync::{Condvar, Mutex};
use std::collections::VecDeque;
use std::task::{Context, Poll};
use std::time::Duration;
use std::future::Future;
use std::pin::Pin;
use std::thread;
use futures::executor::block_on;
use futures::future::BoxFuture;
use futures::channel::oneshot;
use tracing::warn;
use crate::error::{Error, ErrorKind, Result};
const IDLE_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_THREADS: usize = 1024;
type Job = Box<dyn FnOnce() + Send>;
static POOL: Pool = Pool {
state: Mutex::new(State { jobs: VecDeque::new(), threads: 0, idle: 0, started: 0 }),
cond: Condvar::new(),
};
struct Pool {
state: Mutex<State>,
cond: Condvar,
}
struct State {
jobs: VecDeque<Job>,
threads: usize,
idle: usize,
started: usize,
}
pub fn spawn<T: Send + 'static>(func: impl FnOnce() -> Result<T> + Send + 'static) -> Task<T> {
let (tx, rx) = oneshot::channel();
let job: Job = Box::new(move || { let _ = tx.send(func()); });
let mut state = POOL.state.lock().unwrap();
state.jobs.push_back(job);
if state.idle < state.jobs.len() && state.threads < MAX_THREADS {
state.threads += 1;
state.started += 1;
if let Err(err) = thread::Builder::new().name(format!("pool-{}", state.started)).spawn(worker) {
state.threads -= 1;
warn!(%err, "Failed to start a thread of the pool");
}
}
drop(state);
POOL.cond.notify_one();
Task(rx)
}
pub fn run<T>(blocking: bool, request: impl Future<Output = Result<T>> + Send + 'static) -> BoxFuture<'static, Result<T>>
where
T: Send + 'static,
{
if blocking {
Box::pin(spawn(move || block_on(request)))
} else {
Box::pin(request)
}
}
#[derive(Debug)]
pub struct Task<T>(oneshot::Receiver<Result<T>>);
impl<T> Future for Task<T> {
type Output = Result<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.0).poll(cx).map(|res| res.unwrap_or_else(|_| Err(Error::new(ErrorKind::Internal)
.with_description("The job was dropped by the pool before it answered"))))
}
}
fn worker() {
let mut state = POOL.state.lock().unwrap();
loop {
if let Some(job) = state.jobs.pop_front() {
drop(state);
let _ = panic::catch_unwind(AssertUnwindSafe(job));
state = POOL.state.lock().unwrap();
continue;
}
state.idle += 1;
let (next, timeout) = POOL.cond.wait_timeout(state, IDLE_TIMEOUT).unwrap();
state = next;
state.idle -= 1;
if timeout.timed_out() && state.jobs.is_empty() {
state.threads -= 1;
return;
}
}
}