140 lines
3.8 KiB
Rust
Raw Normal View History

2021-06-26 23:52:03 +08:00
use crate::{
2021-06-27 15:11:41 +08:00
module::{Event, Module},
2021-06-28 14:27:16 +08:00
rt::Runtime,
stream::CommandStreamFuture,
2021-06-26 23:52:03 +08:00
};
2021-06-27 15:11:41 +08:00
use futures_core::{ready, task::Context};
2021-06-28 14:27:16 +08:00
use std::{cell::RefCell, collections::HashMap, future::Future, io, rc::Rc, sync::Arc};
2021-06-26 23:52:03 +08:00
use tokio::{
macros::support::{Pin, Poll},
sync::{
mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
oneshot,
},
};
thread_local!(
static CURRENT: RefCell<Option<Arc<FlowySystem>>> = RefCell::new(None);
);
2021-06-27 22:07:33 +08:00
#[derive(Debug)]
pub enum SystemCommand {
2021-06-27 15:11:41 +08:00
Exit(i8),
}
pub type ModuleMap = Rc<HashMap<Event, Rc<Module>>>;
2021-06-26 23:52:03 +08:00
pub struct FlowySystem {
2021-06-28 14:27:16 +08:00
sys_cmd_tx: UnboundedSender<SystemCommand>,
2021-06-26 23:52:03 +08:00
}
impl FlowySystem {
pub fn construct<F, S, T>(module_factory: F, stream_factory: S) -> SystemRunner
2021-06-26 23:52:03 +08:00
where
2021-06-28 14:27:16 +08:00
F: FnOnce() -> Vec<Module>,
S: FnOnce(ModuleMap) -> CommandStreamFuture<T>,
T: 'static,
2021-06-26 23:52:03 +08:00
{
let runtime = Runtime::new().unwrap();
2021-06-28 14:27:16 +08:00
let (sys_cmd_tx, sys_cmd_rx) = unbounded_channel::<SystemCommand>();
2021-06-26 23:52:03 +08:00
let (stop_tx, stop_rx) = oneshot::channel();
2021-06-27 15:11:41 +08:00
runtime.spawn(SystemController {
stop_tx: Some(stop_tx),
2021-06-28 14:27:16 +08:00
sys_cmd_rx,
2021-06-27 15:11:41 +08:00
});
2021-06-26 23:52:03 +08:00
2021-06-28 14:27:16 +08:00
let factory = module_factory();
let mut module_map = HashMap::new();
2021-06-26 23:52:03 +08:00
factory.into_iter().for_each(|m| {
2021-06-28 14:27:16 +08:00
let events = m.events();
let rc_module = Rc::new(m);
events.into_iter().for_each(|e| {
module_map.insert(e, rc_module.clone());
2021-06-28 14:27:16 +08:00
});
2021-06-26 23:52:03 +08:00
});
2021-06-28 16:27:46 +08:00
let system = Self { sys_cmd_tx };
let stream_fut = stream_factory(Rc::new(module_map));
runtime.spawn(stream_fut);
2021-06-26 23:52:03 +08:00
FlowySystem::set_current(system);
let runner = SystemRunner { rt: runtime, stop_rx };
runner
}
2021-06-27 15:11:41 +08:00
pub fn stop(&self) {
2021-06-28 14:27:16 +08:00
match self.sys_cmd_tx.send(SystemCommand::Exit(0)) {
2021-06-27 15:11:41 +08:00
Ok(_) => {},
Err(e) => {
log::error!("Stop system error: {}", e);
},
2021-06-27 01:24:00 +08:00
}
2021-06-26 23:52:03 +08:00
}
#[doc(hidden)]
pub fn set_current(sys: FlowySystem) {
CURRENT.with(|cell| {
*cell.borrow_mut() = Some(Arc::new(sys));
})
}
pub fn current() -> Arc<FlowySystem> {
CURRENT.with(|cell| match *cell.borrow() {
Some(ref sys) => sys.clone(),
None => panic!("System is not running"),
})
}
2021-06-27 15:11:41 +08:00
}
2021-06-26 23:52:03 +08:00
struct SystemController {
2021-06-27 15:11:41 +08:00
stop_tx: Option<oneshot::Sender<i8>>,
2021-06-28 14:27:16 +08:00
sys_cmd_rx: UnboundedReceiver<SystemCommand>,
2021-06-26 23:52:03 +08:00
}
impl Future for SystemController {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
2021-06-28 14:27:16 +08:00
match ready!(Pin::new(&mut self.sys_cmd_rx).poll_recv(cx)) {
2021-06-26 23:52:03 +08:00
None => return Poll::Ready(()),
2021-06-27 15:11:41 +08:00
Some(cmd) => match cmd {
SystemCommand::Exit(code) => {
if let Some(tx) = self.stop_tx.take() {
let _ = tx.send(code);
}
},
2021-06-26 23:52:03 +08:00
},
}
}
}
}
pub struct SystemRunner {
rt: Runtime,
2021-06-27 15:11:41 +08:00
stop_rx: oneshot::Receiver<i8>,
2021-06-26 23:52:03 +08:00
}
impl SystemRunner {
pub fn run(self) -> io::Result<()> {
let SystemRunner { rt, stop_rx } = self;
match rt.block_on(stop_rx) {
Ok(code) => {
if code != 0 {
Err(io::Error::new(
io::ErrorKind::Other,
format!("Non-zero exit code: {}", code),
))
} else {
Ok(())
}
},
Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
}
}
2021-06-28 14:27:16 +08:00
pub fn spawn<F: Future<Output = ()> + 'static>(self, future: F) -> Self {
2021-06-26 23:52:03 +08:00
self.rt.spawn(future);
self
}
}