1
0
mirror of https://github.com/fafhrd91/actix-net synced 2024-11-28 06:32:39 +01:00
actix-net/actix-rt/src/arbiter.rs

393 lines
12 KiB
Rust
Raw Normal View History

2019-09-03 00:03:03 +02:00
use std::any::{Any, TypeId};
2018-12-10 04:55:40 +01:00
use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
2018-12-10 05:30:04 +01:00
use std::{fmt, thread};
2018-12-10 04:55:40 +01:00
use futures::sync::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
use futures::sync::oneshot::{channel, Canceled, Sender};
2018-12-10 04:55:40 +01:00
use futures::{future, Async, Future, IntoFuture, Poll, Stream};
use tokio_current_thread::spawn;
use crate::builder::Builder;
use crate::system::System;
use copyless::BoxHelper;
2018-12-10 04:55:40 +01:00
thread_local!(
static ADDR: RefCell<Option<Arbiter>> = RefCell::new(None);
static RUNNING: Cell<bool> = Cell::new(false);
static Q: RefCell<Vec<Box<dyn Future<Item = (), Error = ()>>>> = RefCell::new(Vec::new());
2019-09-03 00:03:03 +02:00
static STORAGE: RefCell<HashMap<TypeId, Box<dyn Any>>> = RefCell::new(HashMap::new());
2018-12-10 04:55:40 +01:00
);
2018-12-10 07:14:29 +01:00
pub(crate) static COUNT: AtomicUsize = AtomicUsize::new(0);
2018-12-10 04:55:40 +01:00
pub(crate) enum ArbiterCommand {
Stop,
Execute(Box<dyn Future<Item = (), Error = ()> + Send>),
ExecuteFn(Box<dyn FnExec>),
2018-12-10 05:30:04 +01:00
}
impl fmt::Debug for ArbiterCommand {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
ArbiterCommand::Stop => write!(f, "ArbiterCommand::Stop"),
ArbiterCommand::Execute(_) => write!(f, "ArbiterCommand::Execute"),
ArbiterCommand::ExecuteFn(_) => write!(f, "ArbiterCommand::ExecuteFn"),
2018-12-10 05:30:04 +01:00
}
}
2018-12-10 04:55:40 +01:00
}
#[derive(Debug)]
2019-03-15 01:24:27 +01:00
/// Arbiters provide an asynchronous execution environment for actors, functions
/// and futures. When an Arbiter is created, they spawn a new OS thread, and
/// host an event loop. Some Arbiter functions execute on the current thread.
pub struct Arbiter {
sender: UnboundedSender<ArbiterCommand>,
thread_handle: Option<thread::JoinHandle<()>>,
}
impl Clone for Arbiter {
fn clone(&self) -> Self {
Self::with_sender(self.sender.clone())
}
}
2018-12-10 04:55:40 +01:00
impl Default for Arbiter {
fn default() -> Self {
Self::new()
}
}
impl Arbiter {
pub(crate) fn new_system() -> Self {
let (tx, rx) = unbounded();
let arb = Arbiter::with_sender(tx);
2018-12-10 04:55:40 +01:00
ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
RUNNING.with(|cell| cell.set(false));
2019-09-03 00:03:03 +02:00
STORAGE.with(|cell| cell.borrow_mut().clear());
2018-12-10 04:55:40 +01:00
Arbiter::spawn(ArbiterController { stop: None, rx });
arb
}
2019-03-15 01:24:27 +01:00
/// Returns the current thread's arbiter's address. If no Arbiter is present, then this
/// function will panic!
2018-12-10 05:30:04 +01:00
pub fn current() -> Arbiter {
ADDR.with(|cell| match *cell.borrow() {
Some(ref addr) => addr.clone(),
None => panic!("Arbiter is not running"),
})
}
2019-03-15 01:24:27 +01:00
/// Stop arbiter from continuing it's event loop.
2018-12-10 04:55:40 +01:00
pub fn stop(&self) {
let _ = self.sender.unbounded_send(ArbiterCommand::Stop);
2018-12-10 04:55:40 +01:00
}
/// Spawn new thread and run event loop in spawned thread.
/// Returns address of newly created arbiter.
pub fn new() -> Arbiter {
let id = COUNT.fetch_add(1, Ordering::Relaxed);
let name = format!("actix-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 || {
2018-12-10 04:55:40 +01:00
let mut rt = Builder::new().build_rt().expect("Can not create Runtime");
let arb = Arbiter::with_sender(arb_tx);
2018-12-10 04:55:40 +01:00
let (stop, stop_rx) = channel();
RUNNING.with(|cell| cell.set(true));
2019-09-03 00:03:03 +02:00
STORAGE.with(|cell| cell.borrow_mut().clear());
2018-12-10 04:55:40 +01:00
System::set_current(sys);
// start arbiter controller
rt.spawn(ArbiterController {
stop: Some(stop),
rx: arb_rx,
});
ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
// register arbiter
let _ = System::current()
.sys()
.unbounded_send(SystemCommand::RegisterArbiter(id, arb.clone()));
// run loop
let _ = match rt.block_on(stop_rx) {
Ok(code) => code,
Err(_) => 1,
};
// unregister arbiter
let _ = System::current()
.sys()
.unbounded_send(SystemCommand::UnregisterArbiter(id));
}).unwrap_or_else(|err| panic!("Cannot spawn an arbiter's thread {:?}: {:?}", &name, err));
2018-12-10 04:55:40 +01:00
Arbiter{sender: arb_tx2, thread_handle: Some(handle)}
2018-12-10 04:55:40 +01:00
}
pub(crate) fn run_system() {
RUNNING.with(|cell| cell.set(true));
Q.with(|cell| {
let mut v = cell.borrow_mut();
for fut in v.drain(..) {
spawn(fut);
}
});
}
pub(crate) fn stop_system() {
RUNNING.with(|cell| cell.set(false));
}
2019-03-15 01:24:27 +01:00
/// Spawn a future on the current thread. This does not create a new Arbiter
/// or Arbiter address, it is simply a helper for spawning futures on the current
/// thread.
2018-12-10 04:55:40 +01:00
pub fn spawn<F>(future: F)
where
F: Future<Item = (), Error = ()> + 'static,
{
RUNNING.with(move |cell| {
if cell.get() {
spawn(Box::alloc().init(future));
2018-12-10 04:55:40 +01:00
} else {
Q.with(move |cell| cell.borrow_mut().push(Box::alloc().init(future)));
2018-12-10 04:55:40 +01:00
}
});
}
2019-03-15 01:24:27 +01:00
/// Executes a future on the current thread. This does not create a new Arbiter
/// or Arbiter address, it is simply a helper for executing futures on the current
/// thread.
2018-12-10 04:55:40 +01:00
pub fn spawn_fn<F, R>(f: F)
where
F: FnOnce() -> R + 'static,
R: IntoFuture<Item = (), Error = ()> + 'static,
{
Arbiter::spawn(future::lazy(f))
}
2018-12-10 05:30:04 +01:00
2019-03-15 01:24:27 +01:00
/// Send a future to the Arbiter's thread, and spawn it.
2018-12-10 05:30:04 +01:00
pub fn send<F>(&self, future: F)
where
F: Future<Item = (), Error = ()> + Send + 'static,
{
let _ = self
.sender
2018-12-10 05:30:04 +01:00
.unbounded_send(ArbiterCommand::Execute(Box::new(future)));
}
2019-03-15 01:24:27 +01:00
/// Send a function to the Arbiter's thread, and execute it. Any result from the function
/// is discarded.
pub fn exec_fn<F>(&self, f: F)
where
F: FnOnce() + Send + 'static,
{
let _ = self
.sender
.unbounded_send(ArbiterCommand::ExecuteFn(Box::new(move || {
f();
})));
}
2019-03-15 01:24:27 +01:00
/// Send a function to the Arbiter's thread. This function will be executed asynchronously.
/// A future is created, and when resolved will contain the result of the function sent
/// to the Arbiters thread.
pub fn exec<F, R>(&self, f: F) -> impl Future<Item = R, Error = Canceled>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
{
let (tx, rx) = channel();
let _ = self
.sender
.unbounded_send(ArbiterCommand::ExecuteFn(Box::new(move || {
if !tx.is_canceled() {
let _ = tx.send(f());
}
})));
rx
}
2019-09-03 00:03:03 +02:00
/// Set item to arbiter storage
pub fn set_item<T: 'static>(item: T) {
STORAGE.with(move |cell| cell.borrow_mut().insert(TypeId::of::<T>(), Box::new(item)));
}
/// Check if arbiter storage contains item
pub fn contains_item<T: 'static>() -> bool {
STORAGE.with(move |cell| cell.borrow().get(&TypeId::of::<T>()).is_some())
}
/// Get a reference to a type previously inserted on this arbiter's storage.
///
/// Panics is item is not inserted
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)
})
}
/// Get a mutable reference to a type previously inserted on this arbiter's storage.
///
/// Panics is item is not inserted
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: UnboundedSender<ArbiterCommand>) -> Self {
Self{sender, thread_handle: None}
}
/// Wait for the event loop to stop by joining the underlying thread (if have Some).
pub fn join(&mut self) -> thread::Result<()>{
if let Some(thread_handle) = self.thread_handle.take() {
thread_handle.join()
}
else {
Ok(())
}
}
2018-12-10 04:55:40 +01:00
}
struct ArbiterController {
stop: Option<Sender<i32>>,
rx: UnboundedReceiver<ArbiterCommand>,
}
impl Drop for ArbiterController {
fn drop(&mut self) {
if thread::panicking() {
if System::current().stop_on_panic() {
2019-10-14 07:19:08 +02:00
eprintln!("Panic in Arbiter thread, shutting down system.");
2018-12-10 04:55:40 +01:00
System::current().stop_with_code(1)
2019-10-14 07:19:08 +02:00
} else {
eprintln!("Panic in Arbiter thread.");
2018-12-10 04:55:40 +01:00
}
}
}
}
impl Future for ArbiterController {
type Item = ();
type Error = ();
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
2018-12-10 05:30:04 +01:00
loop {
match self.rx.poll() {
Ok(Async::Ready(None)) | Err(_) => return Ok(Async::Ready(())),
Ok(Async::Ready(Some(item))) => match item {
ArbiterCommand::Stop => {
if let Some(stop) = self.stop.take() {
let _ = stop.send(0);
};
return Ok(Async::Ready(()));
}
ArbiterCommand::Execute(fut) => {
spawn(fut);
}
ArbiterCommand::ExecuteFn(f) => {
f.call_box();
}
2018-12-10 05:30:04 +01:00
},
Ok(Async::NotReady) => return Ok(Async::NotReady),
}
2018-12-10 04:55:40 +01:00
}
}
}
#[derive(Debug)]
pub(crate) enum SystemCommand {
Exit(i32),
RegisterArbiter(usize, Arbiter),
UnregisterArbiter(usize),
}
#[derive(Debug)]
pub(crate) struct SystemArbiter {
stop: Option<Sender<i32>>,
commands: UnboundedReceiver<SystemCommand>,
arbiters: HashMap<usize, Arbiter>,
}
impl SystemArbiter {
pub(crate) fn new(stop: Sender<i32>, commands: UnboundedReceiver<SystemCommand>) -> Self {
SystemArbiter {
commands,
stop: Some(stop),
arbiters: HashMap::new(),
}
}
}
impl Future for SystemArbiter {
type Item = ();
type Error = ();
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
2018-12-10 06:51:35 +01:00
loop {
match self.commands.poll() {
Ok(Async::Ready(None)) | Err(_) => return Ok(Async::Ready(())),
Ok(Async::Ready(Some(cmd))) => match cmd {
SystemCommand::Exit(code) => {
// stop arbiters
for arb in self.arbiters.values() {
arb.stop();
}
// stop event loop
if let Some(stop) = self.stop.take() {
let _ = stop.send(code);
}
}
SystemCommand::RegisterArbiter(name, hnd) => {
self.arbiters.insert(name, hnd);
2018-12-10 04:55:40 +01:00
}
2018-12-10 06:51:35 +01:00
SystemCommand::UnregisterArbiter(name) => {
self.arbiters.remove(&name);
2018-12-10 04:55:40 +01:00
}
2018-12-10 06:51:35 +01:00
},
Ok(Async::NotReady) => return Ok(Async::NotReady),
}
2018-12-10 04:55:40 +01:00
}
}
}
pub trait FnExec: Send + 'static {
fn call_box(self: Box<Self>);
}
2018-12-10 04:55:40 +01:00
impl<F> FnExec for F
where
F: FnOnce() + Send + 'static,
{
#[allow(clippy::boxed_local)]
fn call_box(self: Box<Self>) {
(*self)()
}
}