refactoring, now user can provide their monitor

Author: win-t

src/lib.rs

@@ -270,7 +270,7 @@ cfg_default! {
     pub mod fs;
     pub mod path;
     pub mod net;
-    pub(crate) mod rt;
+    pub mod rt;
 }
 
 cfg_unstable! {

src/rt/mod.rs

@@ -6,33 +6,27 @@ use once_cell::sync::Lazy;
 
 use crate::utils::abort_on_panic;
 
-pub use reactor::{Reactor, Watcher};
-pub use runtime::Runtime;
+pub(crate) use reactor::{Reactor, Watcher};
+pub(crate) use runtime::Runtime;
 
-mod monitor;
+pub mod monitor;
 mod reactor;
 mod runtime;
 
 /// The global runtime.
-pub static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
+pub(crate) static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
     thread::Builder::new()
         .name("async-std/runtime".to_string())
         .spawn(|| abort_on_panic(|| RUNTIME.run()))
         .expect("cannot start a runtime thread");
 
-    monitor::spawn_thread();
-
     Runtime::new()
 });
 
-/// Spawn new worker thread
-pub fn scale_up() {
+pub(crate) fn scale_up() {
     RUNTIME.scale_up();
 }
 
-/// Terminate one worker thread.
-/// The number of worker thread cannot go below the number of availabe cpus,
-/// so this function when will do nothing if that happen.
-pub fn scale_down() {
+pub(crate) fn scale_down() {
     RUNTIME.scale_down();
 }

src/rt/monitor.rs

@@ -1,51 +1,156 @@
+//! Monitor for the runtime.
+
 use std::sync::{Arc, Mutex};
 use std::thread;
-use std::time::Duration;
+use std::time::{Duration, Instant};
+
+use once_cell::sync::Lazy;
 
 use crate::rt;
 use crate::task;
-use crate::utils::abort_on_panic;
-
-pub fn spawn_thread() {
-    thread::Builder::new()
-        .name("async-std/monitor".to_string())
-        .spawn(|| {
-            const PROBING_DURATION_MS: u64 = 500;
-            const SCALING_DOWN_SEC: u64 = 1 * 60; // 1 minute
-
-            abort_on_panic(|| {
-                let running = &Arc::new(Mutex::new(false));
-
-                {
-                    let running = Arc::clone(running);
-                    task::spawn(async move {
-                        loop {
-                            *running.lock().unwrap() = true;
-                            task::sleep(Duration::from_millis(PROBING_DURATION_MS)).await;
-                        }
-                    });
-                }
-
-                {
-                    task::spawn(async {
-                        loop {
-                            task::sleep(Duration::from_secs(SCALING_DOWN_SEC)).await;
-                            rt::scale_down();
-                        }
-                    });
-                }
-
-                loop {
-                    *running.lock().unwrap() = false;
-                    thread::sleep(Duration::from_millis(PROBING_DURATION_MS * 2));
-                    if !*running.lock().unwrap() {
-                        eprintln!(
-                            "WARNING: You are blocking the runtime, please use spawn_blocking"
-                        );
-                        rt::scale_up();
-                    }
-                }
-            })
-        })
-        .expect("cannot start a monitor thread");
+
+/// Context of monitor thread
+#[derive(Debug)]
+pub struct Context(());
+
+impl Context {
+    /// Spawn a new worker thread.
+    pub fn scale_up(&self) {
+        rt::scale_up();
+    }
+
+    /// Terminate a worker thread if the number of current worker thread
+    /// is greater than number avaliable cpu
+    pub fn scale_down(&self) {
+        rt::scale_down();
+    }
+}
+
+pub(crate) static MONITOR: Lazy<Mutex<&'static Monitor>> =
+    Lazy::new(|| Mutex::new(&DEFAULT_MONITOR));
+
+/// Monitor function.
+pub struct Monitor(pub fn(&Context));
+
+impl Monitor {
+    pub(crate) fn run(&self) {
+        let scaler = Context(());
+        self.0(&scaler);
+    }
+}
+
+impl std::fmt::Debug for Monitor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("<monitor>")
+    }
+}
+
+/// Replace the monitor.
+///
+/// This will replace monitor function used by monitor thread by runtime.
+///
+/// Monitor can only be replaced before any task spawned into runtime.
+///
+/// Default monitor is [`DEFAULT_MONITOR`]
+///
+/// [`DEFAULT_MONITOR`]: static.DEFAULT_MONITOR.html
+pub fn replace(new: &'static Monitor) -> Result<&'static Monitor, &'static str> {
+    let mut m = MONITOR.try_lock().map_err(|_| "Cannot change monitor")?;
+    let old: &'static Monitor = &m;
+    *m = new;
+    Ok(old)
+}
+
+async fn repeater(interval: Duration, mut before: impl FnMut(), mut after: impl FnMut()) {
+    loop {
+        before();
+        task::sleep(interval).await;
+        after();
+    }
+}
+
+fn repeater_blocking(interval: Duration, mut before: impl FnMut(), mut after: impl FnMut()) {
+    loop {
+        before();
+        thread::sleep(interval);
+        after();
+    }
+}
+
+/// Default monitor
+///
+/// Whenever runtime is blocked, print warning and scale up the runtime.
+/// This also try to scale down when there are too many worker thread.
+pub static DEFAULT_MONITOR: Monitor = Monitor(default_monitor);
+
+fn default_monitor(scaler: &Context) {
+    const PROB_INTERVAL: Duration = Duration::from_millis(500);
+    const SCALE_DOWN_INTERVAL: Duration = Duration::from_secs(60);
+
+    let running = &Arc::new(Mutex::new(false));
+
+    {
+        let running = Arc::clone(running);
+        task::spawn(repeater(
+            PROB_INTERVAL,
+            move || {
+                *running.lock().unwrap() = true;
+            },
+            || {},
+        ));
+    }
+
+    let mut next_scalling_down = Instant::now() + SCALE_DOWN_INTERVAL;
+
+    repeater_blocking(
+        PROB_INTERVAL + Duration::from_millis(50),
+        || {
+            *running.lock().unwrap() = false;
+        },
+        move || {
+            if !*running.lock().unwrap() {
+                eprintln!("WARNING: You are blocking the runtime, please use spawn_blocking");
+                scaler.scale_up();
+            }
+
+            if next_scalling_down < Instant::now() {
+                scaler.scale_down();
+                next_scalling_down += SCALE_DOWN_INTERVAL;
+            }
+        },
+    );
+}
+
+/// Abort on blocking monitor.
+///
+/// Whenever runtime is blocked, abort the program.
+pub static ABORT_ON_BLOCKING_MONITOR: Monitor = Monitor(abort_on_blocking_monitor);
+
+fn abort_on_blocking_monitor(_scaler: &Context) {
+    const PROB_INTERVAL: Duration = Duration::from_secs(1);
+
+    let running = &Arc::new(Mutex::new(false));
+
+    {
+        let running = Arc::clone(running);
+        task::spawn(repeater(
+            PROB_INTERVAL,
+            move || {
+                *running.lock().unwrap() = true;
+            },
+            || {},
+        ));
+    }
+
+    repeater_blocking(
+        PROB_INTERVAL + Duration::from_millis(50),
+        || {
+            *running.lock().unwrap() = false;
+        },
+        || {
+            if !*running.lock().unwrap() {
+                panic!("FATAL: You are blocking the runtime, please use spawn_blocking");
+            }
+        },
+    );
 }

src/rt/runtime.rs

@@ -12,6 +12,7 @@ use crossbeam_deque::{Injector, Steal, Stealer, Worker};
 use crossbeam_utils::thread::{scope, Scope};
 use once_cell::unsync::OnceCell;
 
+use crate::rt::monitor::MONITOR;
 use crate::rt::Reactor;
 use crate::sync::Spinlock;
 use crate::task::Runnable;
@@ -40,7 +41,6 @@ enum Task {
 enum Action {
     ScaleUp,
     ScaleDown,
-    Terminate,
 }
 
 /// An async runtime.
@@ -138,6 +138,16 @@ impl Runtime {
         scope(|s| {
             (0..self.min_worker).for_each(|_| self.start_new_thread(s));
 
+            s.builder()
+                .name("async-std/monitor".to_string())
+                .spawn(move |_| {
+                    abort_on_panic(|| {
+                        MONITOR.lock().unwrap().run();
+                        panic!("Monitor function must not return");
+                    })
+                })
+                .expect("cannot start a monitor thread");
+
             loop {
                 match self.reciever.recv().unwrap() {
                     Action::ScaleUp => self.start_new_thread(s),
@@ -146,7 +156,6 @@ impl Runtime {
                         // and terminate itself
                         self.injector.push(Task::Terminate)
                     }
-                    Action::Terminate => return,
                 }
             }
         })
@@ -179,9 +188,6 @@ impl Runtime {
         }
         if let Some(index) = index {
             stealers.remove(index);
-            if stealers.is_empty() {
-                self.sender.send(Action::Terminate).unwrap();
-            }
         }
     }