diff --git a/rust-lib/flowy-sys/src/module/module.rs b/rust-lib/flowy-sys/src/module/module.rs index 78069e3a80..6182ced955 100644 --- a/rust-lib/flowy-sys/src/module/module.rs +++ b/rust-lib/flowy-sys/src/module/module.rs @@ -37,7 +37,7 @@ pub type CommandServiceFactory = BoxServiceFactory<(), ServiceRequest, ServiceRe pub struct Module { name: String, data: DataContainer, - factory_map: HashMap, + service_map: HashMap, req_tx: UnboundedSender, req_rx: UnboundedReceiver, resp_tx: UnboundedSender, @@ -49,7 +49,7 @@ impl Module { Self { name: "".to_owned(), data: DataContainer::new(), - factory_map: HashMap::new(), + service_map: HashMap::new(), req_tx, req_rx, resp_tx, @@ -74,11 +74,11 @@ impl Module { R: Future + 'static, R::Output: Responder + 'static, { - self.factory_map.insert(command, factory(HandlerService::new(handler))); + self.service_map.insert(command, factory(HandlerService::new(handler))); self } - pub fn can_handle(&self, cmd: &Command) -> bool { self.factory_map.contains_key(cmd) } + pub fn can_handle(&self, cmd: &Command) -> bool { self.service_map.contains_key(cmd) } pub fn req_tx(&self) -> UnboundedSender { self.req_tx.clone() } @@ -90,6 +90,13 @@ impl Module { }, } } + + pub fn service_sender_map(&self) -> HashMap> { + self.service_map + .keys() + .map(|key| (key.clone(), self.req_tx())) + .collect::>() + } } impl Future for Module { @@ -98,7 +105,7 @@ impl Future for Module { loop { match ready!(Pin::new(&mut self.req_rx).poll_recv(cx)) { None => return Poll::Ready(()), - Some(request) => match self.factory_map.get(request.get_cmd()) { + Some(request) => match self.service_map.get(request.get_cmd()) { Some(factory) => { let fut = ModuleServiceFuture { request, diff --git a/rust-lib/flowy-sys/src/rt/system.rs b/rust-lib/flowy-sys/src/rt/system.rs index 8cbe170b05..0b36333a5d 100644 --- a/rust-lib/flowy-sys/src/rt/system.rs +++ b/rust-lib/flowy-sys/src/rt/system.rs @@ -1,12 +1,13 @@ use crate::{ - module::{Command, Module}, + module::{Command, CommandServiceFactory, Module}, request::FlowyRequest, response::FlowyResponse, rt::Runtime, + service::BoxServiceFactory, }; use futures_core::{future::LocalBoxFuture, ready, task::Context}; use futures_util::{future, pin_mut}; -use std::{cell::RefCell, future::Future, io, sync::Arc}; +use std::{cell::RefCell, collections::HashMap, future::Future, io, sync::Arc}; use tokio::{ macros::support::{Pin, Poll}, sync::{ @@ -21,7 +22,7 @@ thread_local!( pub struct FlowySystem { resp_tx: UnboundedSender, - modules: Vec, + sender_map: HashMap>, } impl FlowySystem { @@ -37,31 +38,26 @@ impl FlowySystem { let mut system = Self { resp_tx: resp_tx.clone(), - modules: vec![], + sender_map: HashMap::default(), }; let factory = module_factory(resp_tx.clone()); factory.into_iter().for_each(|m| { + system.sender_map.extend(m.service_sender_map()); runtime.spawn(m); - // system.add_module(m); }); FlowySystem::set_current(system); - let runner = SystemRunner { rt: runtime, stop_rx }; runner } pub fn handle_command(&self, cmd: Command, request: FlowyRequest) { - self.modules.iter().for_each(|m| { - if m.can_handle(&cmd) { - m.handle(request.clone()); - } - }) + if let Some(sender) = self.sender_map.get(&cmd) { + sender.send(request); + } } - pub fn add_module(&mut self, module: Module) { self.modules.push(module); } - #[doc(hidden)] pub fn set_current(sys: FlowySystem) { CURRENT.with(|cell| { @@ -91,7 +87,7 @@ impl Future for SystemController { match ready!(Pin::new(&mut self.resp_rx).poll_recv(cx)) { None => return Poll::Ready(()), Some(resp) => { - // FFI + // FF println!("Receive response: {:?}", resp); }, }