Merge pull request #86 from silent-rs/features/set_scheduler_global

feat(scheduler): 添加调度器中间件并优化调度器初始化

Author: hubertshelley

silent/src/lib.rs

@@ -46,4 +46,4 @@ pub use handler::HandlerWrapper;
 pub use headers;
 pub use hyper::{Method, StatusCode, header};
 #[cfg(feature = "scheduler")]
-pub use scheduler::{ProcessTime, Scheduler, SchedulerExt, Task};
+pub use scheduler::{ProcessTime, SCHEDULER, Scheduler, SchedulerExt, Task};

silent/src/prelude.rs

@@ -19,7 +19,7 @@ pub use crate::middleware::MiddleWareHandler;
 pub use crate::route::handler_append::{HandlerAppend, HandlerGetter};
 pub use crate::route::{RootRoute, Route, RouteService, RouterAdapt};
 #[cfg(feature = "scheduler")]
-pub use crate::scheduler::{SchedulerExt, Task};
+pub use crate::scheduler::{SCHEDULER, SchedulerExt, Task};
 #[cfg(feature = "security")]
 pub use crate::security::{argon2, pbkdf2};
 #[cfg(feature = "server")]

silent/src/route/root.rs

@@ -1,5 +1,3 @@
-#[cfg(feature = "scheduler")]
-use crate::Scheduler;
 #[cfg(feature = "cookie")]
 use crate::cookie::middleware::CookieMiddleware;
 use crate::middlewares::RequestTimeLogger;
@@ -17,8 +15,6 @@ use async_session::SessionStore;
 use async_trait::async_trait;
 use std::fmt;
 use std::sync::Arc;
-#[cfg(feature = "scheduler")]
-use tokio::sync::Mutex;
 
 #[derive(Clone, Default)]
 pub struct RootRoute {
@@ -27,8 +23,6 @@ pub struct RootRoute {
     #[cfg(feature = "session")]
     pub(crate) session_set: bool,
     pub(crate) configs: Option<Configs>,
-    #[cfg(feature = "scheduler")]
-    pub(crate) scheduler: Arc<Mutex<Scheduler>>,
 }
 
 impl fmt::Debug for RootRoute {
@@ -51,8 +45,6 @@ impl RootRoute {
             #[cfg(feature = "session")]
             session_set: false,
             configs: None,
-            #[cfg(feature = "scheduler")]
-            scheduler: Arc::new(Mutex::new(Scheduler::new())),
         }
     }
 
@@ -110,8 +102,6 @@ impl Handler for RootRoute {
         tracing::debug!("{:?}", req);
         let configs = self.configs.clone().unwrap_or_default();
         req.configs = configs.clone();
-        #[cfg(feature = "scheduler")]
-        req.extensions_mut().insert(self.scheduler.clone());
 
         let mut root_middlewares = vec![];
         for middleware in self.middlewares.iter().cloned() {

silent/src/scheduler/middleware.rs

@@ -0,0 +1,20 @@
+use crate::{Handler, MiddleWareHandler, Next, Request, Response, Result};
+use async_trait::async_trait;
+
+#[derive(Debug, Default)]
+pub struct SchedulerMiddleware {}
+
+impl SchedulerMiddleware {
+    pub(crate) fn new() -> Self {
+        Self::default()
+    }
+}
+
+#[async_trait]
+impl MiddleWareHandler for SchedulerMiddleware {
+    async fn handle(&self, mut req: Request, next: &Next) -> Result<Response> {
+        let scheduler = super::SCHEDULER.clone();
+        req.extensions_mut().insert(scheduler);
+        next.call(req).await
+    }
+}

silent/src/scheduler/mod.rs

@@ -1,10 +1,11 @@
+pub(crate) mod middleware;
 mod process_time;
 mod storage;
 mod task;
 pub mod traits;
 
 use anyhow::{Result, anyhow};
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 use std::thread;
 use tokio::sync::Mutex;
 use tracing::{debug, error, info};
@@ -13,6 +14,9 @@ pub use process_time::ProcessTime;
 pub use task::Task;
 pub use traits::SchedulerExt;
 
+pub static SCHEDULER: LazyLock<Arc<Mutex<Scheduler>>> =
+    LazyLock::new(|| Arc::new(Mutex::new(Scheduler::new())));
+
 #[derive(Debug, Clone)]
 pub struct Scheduler {
     tasks: Vec<Task>,

silent/src/service/mod.rs

@@ -6,6 +6,8 @@ use crate::Configs;
 use crate::Scheduler;
 use crate::core::listener::Listener;
 use crate::route::RouteService;
+#[cfg(feature = "scheduler")]
+use crate::scheduler::{SCHEDULER, middleware::SchedulerMiddleware};
 use crate::service::serve::Serve;
 use std::net::SocketAddr;
 use tokio::net::{TcpListener, UnixListener};
@@ -117,9 +119,10 @@ impl Server {
         #[cfg(feature = "cookie")]
         root_route.check_cookie();
         #[cfg(feature = "scheduler")]
-        let scheduler = root_route.scheduler.clone();
+        root_route.hook_first(SchedulerMiddleware::new());
         #[cfg(feature = "scheduler")]
         tokio::spawn(async move {
+            let scheduler = SCHEDULER.clone();
             Scheduler::schedule(scheduler).await;
         });
         let mut join_set = JoinSet::new();