use std::any::{Any, TypeId};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
use std::{cell::RefCell, collections::HashMap, fmt, future::Future, pin::Pin, thread};
use async_channel::{unbounded, Receiver, Sender};
use async_oneshot as oneshot;
use futures_core::stream::Stream;
use crate::system::System;
thread_local!(
static ADDR: RefCell<Option<Arbiter>> = RefCell::new(None);
static STORAGE: RefCell<HashMap<TypeId, Box<dyn Any>>> = RefCell::new(HashMap::new());
);
pub(super) static COUNT: AtomicUsize = AtomicUsize::new(0);
pub(super) enum ArbiterCommand {
Stop,
Execute(Box<dyn Future<Output = ()> + Unpin + Send>),
ExecuteFn(Box<dyn FnExec>),
}
pub struct Arbiter {
sender: Sender<ArbiterCommand>,
thread_handle: Option<thread::JoinHandle<()>>,
}
impl fmt::Debug for Arbiter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Arbiter")
}
}
impl Default for Arbiter {
fn default() -> Arbiter {
Arbiter::new()
}
}
impl Clone for Arbiter {
fn clone(&self) -> Self {
Self::with_sender(self.sender.clone())
}
}
impl Arbiter {
#[allow(clippy::borrowed_box)]
pub(super) fn new_system() -> (Self, ArbiterController) {
let (tx, rx) = unbounded();
let arb = Arbiter::with_sender(tx);
ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
STORAGE.with(|cell| cell.borrow_mut().clear());
(arb, ArbiterController { stop: None, rx })
}
pub fn current() -> Arbiter {
ADDR.with(|cell| match *cell.borrow() {
Some(ref addr) => addr.clone(),
None => panic!("Arbiter is not running"),
})
}
pub fn stop(&self) {
let _ = self.sender.try_send(ArbiterCommand::Stop);
}
pub fn new() -> Arbiter {
let id = COUNT.fetch_add(1, Ordering::Relaxed);
let name = format!("ntex-rt:worker:{}", id);
let sys = System::current();
let (arb_tx, arb_rx) = unbounded();
let arb_tx2 = arb_tx.clone();
let handle = thread::Builder::new()
.name(name.clone())
.spawn(move || {
let arb = Arbiter::with_sender(arb_tx);
let (stop, stop_rx) = oneshot::oneshot();
STORAGE.with(|cell| cell.borrow_mut().clear());
System::set_current(sys);
crate::block_on(async move {
crate::spawn(ArbiterController {
stop: Some(stop),
rx: arb_rx,
});
ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
let _ = System::current()
.sys()
.try_send(SystemCommand::RegisterArbiter(id, arb));
let _ = stop_rx.await;
});
let _ = System::current()
.sys()
.try_send(SystemCommand::UnregisterArbiter(id));
})
.unwrap_or_else(|err| {
panic!("Cannot spawn an arbiter's thread {:?}: {:?}", &name, err)
});
Arbiter {
sender: arb_tx2,
thread_handle: Some(handle),
}
}
pub fn spawn<F>(&self, future: F)
where
F: Future<Output = ()> + Send + Unpin + 'static,
{
let _ = self
.sender
.try_send(ArbiterCommand::Execute(Box::new(future)));
}
pub fn exec<F, R>(&self, f: F) -> impl Future<Output = Result<R, oneshot::Closed>>
where
F: FnOnce() -> R + Send + 'static,
R: Sync + Send + 'static,
{
let (mut tx, rx) = oneshot::oneshot();
let _ = self
.sender
.try_send(ArbiterCommand::ExecuteFn(Box::new(move || {
if !tx.is_closed() {
let _ = tx.send(f());
}
})));
rx
}
pub fn exec_fn<F>(&self, f: F)
where
F: FnOnce() + Send + 'static,
{
let _ = self
.sender
.try_send(ArbiterCommand::ExecuteFn(Box::new(move || {
f();
})));
}
pub fn set_item<T: 'static>(item: T) {
STORAGE
.with(move |cell| cell.borrow_mut().insert(TypeId::of::<T>(), Box::new(item)));
}
pub fn contains_item<T: 'static>() -> bool {
STORAGE.with(move |cell| cell.borrow().get(&TypeId::of::<T>()).is_some())
}
pub fn get_item<T: 'static, F, R>(mut f: F) -> R
where
F: FnMut(&T) -> R,
{
STORAGE.with(move |cell| {
let st = cell.borrow();
let item = st
.get(&TypeId::of::<T>())
.and_then(|boxed| (&**boxed as &(dyn Any + 'static)).downcast_ref())
.unwrap();
f(item)
})
}
pub fn get_mut_item<T: 'static, F, R>(mut f: F) -> R
where
F: FnMut(&mut T) -> R,
{
STORAGE.with(move |cell| {
let mut st = cell.borrow_mut();
let item = st
.get_mut(&TypeId::of::<T>())
.and_then(|boxed| (&mut **boxed as &mut (dyn Any + 'static)).downcast_mut())
.unwrap();
f(item)
})
}
fn with_sender(sender: Sender<ArbiterCommand>) -> Self {
Self {
sender,
thread_handle: None,
}
}
pub fn join(&mut self) -> thread::Result<()> {
if let Some(thread_handle) = self.thread_handle.take() {
thread_handle.join()
} else {
Ok(())
}
}
}
pub(crate) struct ArbiterController {
stop: Option<oneshot::Sender<i32>>,
rx: Receiver<ArbiterCommand>,
}
impl Drop for ArbiterController {
fn drop(&mut self) {
if thread::panicking() {
if System::current().stop_on_panic() {
eprintln!("Panic in Arbiter thread, shutting down system.");
System::current().stop_with_code(1)
} else {
eprintln!("Panic in Arbiter thread.");
}
}
}
}
impl Future for ArbiterController {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
match Pin::new(&mut self.rx).poll_next(cx) {
Poll::Ready(None) => return Poll::Ready(()),
Poll::Ready(Some(item)) => match item {
ArbiterCommand::Stop => {
if let Some(mut stop) = self.stop.take() {
let _ = stop.send(0);
};
return Poll::Ready(());
}
ArbiterCommand::Execute(fut) => {
crate::spawn(fut);
}
ArbiterCommand::ExecuteFn(f) => {
f.call_box();
}
},
Poll::Pending => return Poll::Pending,
}
}
}
}
#[derive(Debug)]
pub(super) enum SystemCommand {
Exit(i32),
RegisterArbiter(usize, Arbiter),
UnregisterArbiter(usize),
}
#[derive(Debug)]
pub(super) struct SystemArbiter {
stop: Option<oneshot::Sender<i32>>,
commands: Receiver<SystemCommand>,
arbiters: HashMap<usize, Arbiter>,
}
impl SystemArbiter {
pub(super) fn new(
stop: oneshot::Sender<i32>,
commands: Receiver<SystemCommand>,
) -> Self {
SystemArbiter {
commands,
stop: Some(stop),
arbiters: HashMap::new(),
}
}
}
impl Future for SystemArbiter {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
match Pin::new(&mut self.commands).poll_next(cx) {
Poll::Ready(None) => return Poll::Ready(()),
Poll::Ready(Some(cmd)) => match cmd {
SystemCommand::Exit(code) => {
for arb in self.arbiters.values() {
arb.stop();
}
if let Some(mut stop) = self.stop.take() {
let _ = stop.send(code);
}
}
SystemCommand::RegisterArbiter(name, hnd) => {
self.arbiters.insert(name, hnd);
}
SystemCommand::UnregisterArbiter(name) => {
self.arbiters.remove(&name);
}
},
Poll::Pending => return Poll::Pending,
}
}
}
}
pub(super) trait FnExec: Send + 'static {
fn call_box(self: Box<Self>);
}
impl<F> FnExec for F
where
F: FnOnce() + Send + 'static,
{
#[allow(clippy::boxed_local)]
fn call_box(self: Box<Self>) {
(*self)()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_arbiter_local_storage() {
let _s = System::new("test");
Arbiter::set_item("test");
assert!(Arbiter::get_item::<&'static str, _, _>(|s| *s == "test"));
assert!(Arbiter::get_mut_item::<&'static str, _, _>(|s| *s == "test"));
assert!(Arbiter::contains_item::<&'static str>());
assert!(format!("{:?}", Arbiter::current()).contains("Arbiter"));
}
}