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,22 +6,25 @@ 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 use monitor::replace_monitor;
+pub use monitor::Monitor;
+pub use monitor::ABORT_ON_BLOCKING_MONITOR;
+pub use monitor::DEFAULT_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()
 });
 
@@ -31,6 +34,7 @@ pub fn 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() {

src/rt/monitor.rs

@@ -2,50 +2,126 @@ use std::sync::{Arc, Mutex};
 use std::thread;
 use std::time::Duration;
 
+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");
+
+pub(crate) static MONITOR: Lazy<Mutex<&'static Monitor>> =
+    Lazy::new(|| Mutex::new(&DEFAULT_MONITOR));
+
+/// Monitor function.
+#[derive(Debug)]
+pub struct Monitor(pub fn());
+
+impl Monitor {
+    pub(crate) fn run(&self) {
+        self.0();
+    }
+}
+
+/// Replace the monitor.
+///
+/// This will replace monitor function used by monitor thread by runtime.
+///
+/// Monitor thread is special thread that can be used to monitor wether
+/// the runtime is blocked or not.
+///
+/// Default monitor is [`DEFAULT_MONITOR`]
+///
+/// [`DEFAULT_MONITOR`]: static.DEFAULT_MONITOR.html
+pub fn replace_monitor(new: &'static Monitor) -> &'static Monitor {
+    let mut m = MONITOR.lock().unwrap();
+    let old = &m as &'static Monitor;
+    *m = new;
+    old
+}
+
+/// 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);
+
+/// Abort on blocking monitor.
+///
+/// Whenever runtime is blocked, abort the program.
+pub static ABORT_ON_BLOCKING_MONITOR: Monitor = Monitor(abort_on_blocking_monitor);
+
+async fn repeater(interval: Duration, before: impl Fn(), after: impl Fn()) {
+    loop {
+        before();
+        task::sleep(interval).await;
+        after();
+    }
+}
+
+fn repeater_blocking(interval: Duration, before: impl Fn(), after: impl Fn()) {
+    loop {
+        before();
+        thread::sleep(interval);
+        after();
+    }
+}
+
+fn default_monitor() {
+    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;
+            },
+            || {},
+        ));
+    }
+
+    task::spawn(repeater(SCALE_DOWN_INTERVAL, || rt::scale_down(), || {}));
+
+    repeater_blocking(
+        PROB_INTERVAL + Duration::from_millis(50),
+        || {
+            *running.lock().unwrap() = false;
+        },
+        || {
+            if !*running.lock().unwrap() {
+                eprintln!("WARNING: You are blocking the runtime, please use spawn_blocking");
+                rt::scale_up();
+            }
+        },
+    );
+}
+
+fn abort_on_blocking_monitor() {
+    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,15 @@ 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();
+                    })
+                })
+                .expect("cannot start a monitor thread");
+
             loop {
                 match self.reciever.recv().unwrap() {
                     Action::ScaleUp => self.start_new_thread(s),
@@ -146,7 +155,6 @@ impl Runtime {
                         // and terminate itself
                         self.injector.push(Task::Terminate)
                     }
-                    Action::Terminate => return,
                 }
             }
         })
@@ -179,9 +187,6 @@ impl Runtime {
         }
         if let Some(index) = index {
             stealers.remove(index);
-            if stealers.is_empty() {
-                self.sender.send(Action::Terminate).unwrap();
-            }
         }
     }