rt: cleanup and simplify scheduler (scheduler v2.5) (#2273)

A refactor of the scheduler internals focusing on simplifying and
reducing unsafety. There are no fundamental logic changes.

* The state transitions of the core task component are refined and
reduced.
* `basic_scheduler` has most unsafety removed.
* `local_set` has most unsafety removed.
* `threaded_scheduler` limits most unsafety to its queue implementation.

2 files changed, 157 insertions, 0 deletions

diff --git a/benches/Cargo.toml b/benches/Cargo.toml
index 8c9ad1d3..f4a1d8fb 100644
--- a/benches/Cargo.toml
+++ b/benches/Cargo.toml
@@ -17,3 +17,8 @@ harness = false
 name = "mpsc"
 path = "mpsc.rs"
 harness = false
+
+[[bench]]
+name = "scheduler"
+path = "scheduler.rs"
+harness = false
diff --git a/benches/scheduler.rs b/benches/scheduler.rs
new file mode 100644
index 00000000..0562a120
--- /dev/null
+++ b/benches/scheduler.rs
@@ -0,0 +1,152 @@
+//! Benchmark implementation details of the theaded scheduler. These benches are
+//! intended to be used as a form of regression testing and not as a general
+//! purpose benchmark demonstrating real-world performance.
+
+use tokio::runtime::{self, Runtime};
+use tokio::sync::oneshot;
+
+use bencher::{benchmark_group, benchmark_main, Bencher};
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::Relaxed;
+use std::sync::{mpsc, Arc};
+
+fn spawn_many(b: &mut Bencher) {
+    const NUM_SPAWN: usize = 10_000;
+
+    let mut rt = rt();
+
+    let (tx, rx) = mpsc::sync_channel(1000);
+    let rem = Arc::new(AtomicUsize::new(0));
+
+    b.iter(|| {
+        rem.store(NUM_SPAWN, Relaxed);
+
+        rt.block_on(async {
+            for _ in 0..NUM_SPAWN {
+                let tx = tx.clone();
+                let rem = rem.clone();
+
+                tokio::spawn(async move {
+                    if 1 == rem.fetch_sub(1, Relaxed) {
+                        tx.send(()).unwrap();
+                    }
+                });
+            }
+
+            let _ = rx.recv().unwrap();
+        });
+    });
+}
+
+fn yield_many(b: &mut Bencher) {
+    const NUM_YIELD: usize = 1_000;
+    const TASKS: usize = 200;
+
+    let rt = rt();
+
+    let (tx, rx) = mpsc::sync_channel(TASKS);
+
+    b.iter(move || {
+        for _ in 0..TASKS {
+            let tx = tx.clone();
+
+            rt.spawn(async move {
+                for _ in 0..NUM_YIELD {
+                    tokio::task::yield_now().await;
+                }
+
+                tx.send(()).unwrap();
+            });
+        }
+
+        for _ in 0..TASKS {
+            let _ = rx.recv().unwrap();
+        }
+    });
+}
+
+fn ping_pong(b: &mut Bencher) {
+    const NUM_PINGS: usize = 1_000;
+
+    let mut rt = rt();
+
+    let (done_tx, done_rx) = mpsc::sync_channel(1000);
+    let rem = Arc::new(AtomicUsize::new(0));
+
+    b.iter(|| {
+        let done_tx = done_tx.clone();
+        let rem = rem.clone();
+        rem.store(NUM_PINGS, Relaxed);
+
+        rt.block_on(async {
+            tokio::spawn(async move {
+                for _ in 0..NUM_PINGS {
+                    let rem = rem.clone();
+                    let done_tx = done_tx.clone();
+
+                    tokio::spawn(async move {
+                        let (tx1, rx1) = oneshot::channel();
+                        let (tx2, rx2) = oneshot::channel();
+
+                        tokio::spawn(async move {
+                            rx1.await.unwrap();
+                            tx2.send(()).unwrap();
+                        });
+
+                        tx1.send(()).unwrap();
+                        rx2.await.unwrap();
+
+                        if 1 == rem.fetch_sub(1, Relaxed) {
+                            done_tx.send(()).unwrap();
+                        }
+                    });
+                }
+            });
+
+            done_rx.recv().unwrap();
+        });
+    });
+}
+
+fn chained_spawn(b: &mut Bencher) {
+    const ITER: usize = 1_000;
+
+    let mut rt = rt();
+
+    fn iter(done_tx: mpsc::SyncSender<()>, n: usize) {
+        if n == 0 {
+            done_tx.send(()).unwrap();
+        } else {
+            tokio::spawn(async move {
+                iter(done_tx, n - 1);
+            });
+        }
+    }
+
+    let (done_tx, done_rx) = mpsc::sync_channel(1000);
+
+    b.iter(move || {
+        let done_tx = done_tx.clone();
+
+        rt.block_on(async {
+            tokio::spawn(async move {
+                iter(done_tx, ITER);
+            });
+
+            done_rx.recv().unwrap();
+        });
+    });
+}
+
+fn rt() -> Runtime {
+    runtime::Builder::new()
+        .threaded_scheduler()
+        .core_threads(4)
+        .enable_all()
+        .build()
+        .unwrap()
+}
+
+benchmark_group!(scheduler, spawn_many, ping_pong, yield_many, chained_spawn,);
+
+benchmark_main!(scheduler);