rt: reduce code defined in macros (#5773)

Instead of defining code in macros, move code definition to sub modules
and use the cfg_macro to declare the module.

tokio/src/runtime/scheduler/multi_thread/handle.rs

@@ -9,6 +9,14 @@ use crate::util::RngSeedGenerator;
 
 use std::fmt;
 
+cfg_metrics! {
+    mod metrics;
+}
+
+cfg_taskdump! {
+    mod taskdump;
+}
+
 /// Handle to the multi thread scheduler
 pub(crate) struct Handle {
     /// Task spawner
@@ -53,73 +61,6 @@ impl Handle {
     }
 }
 
-cfg_metrics! {
-    use crate::runtime::{SchedulerMetrics, WorkerMetrics};
-
-    impl Handle {
-        pub(crate) fn num_workers(&self) -> usize {
-            self.shared.worker_metrics.len()
-        }
-
-        pub(crate) fn num_blocking_threads(&self) -> usize {
-            self.blocking_spawner.num_threads()
-        }
-
-        pub(crate) fn num_idle_blocking_threads(&self) -> usize {
-            self.blocking_spawner.num_idle_threads()
-        }
-
-        pub(crate) fn active_tasks_count(&self) -> usize {
-            self.shared.owned.active_tasks_count()
-        }
-
-        pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
-            &self.shared.scheduler_metrics
-        }
-
-        pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
-            &self.shared.worker_metrics[worker]
-        }
-
-        pub(crate) fn injection_queue_depth(&self) -> usize {
-            self.shared.injection_queue_depth()
-        }
-
-        pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
-            self.shared.worker_local_queue_depth(worker)
-        }
-
-        pub(crate) fn blocking_queue_depth(&self) -> usize {
-            self.blocking_spawner.queue_depth()
-        }
-    }
-}
-
-cfg_taskdump! {
-    impl Handle {
-        pub(crate) async fn dump(&self) -> crate::runtime::Dump {
-            let trace_status = &self.shared.trace_status;
-
-            // If a dump is in progress, block.
-            trace_status.start_trace_request(&self).await;
-
-            let result = loop {
-                if let Some(result) = trace_status.take_result() {
-                    break result;
-                } else {
-                    self.notify_all();
-                    trace_status.result_ready.notified().await;
-                }
-            };
-
-            // Allow other queued dumps to proceed.
-            trace_status.end_trace_request(&self).await;
-
-            result
-        }
-    }
-}
-
 impl fmt::Debug for Handle {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt.debug_struct("multi_thread::Handle { ... }").finish()

tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs

@@ -0,0 +1,41 @@
+use super::Handle;
+
+use crate::runtime::{SchedulerMetrics, WorkerMetrics};
+
+impl Handle {
+    pub(crate) fn num_workers(&self) -> usize {
+        self.shared.worker_metrics.len()
+    }
+
+    pub(crate) fn num_blocking_threads(&self) -> usize {
+        self.blocking_spawner.num_threads()
+    }
+
+    pub(crate) fn num_idle_blocking_threads(&self) -> usize {
+        self.blocking_spawner.num_idle_threads()
+    }
+
+    pub(crate) fn active_tasks_count(&self) -> usize {
+        self.shared.owned.active_tasks_count()
+    }
+
+    pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
+        &self.shared.scheduler_metrics
+    }
+
+    pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
+        &self.shared.worker_metrics[worker]
+    }
+
+    pub(crate) fn injection_queue_depth(&self) -> usize {
+        self.shared.injection_queue_depth()
+    }
+
+    pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
+        self.shared.worker_local_queue_depth(worker)
+    }
+
+    pub(crate) fn blocking_queue_depth(&self) -> usize {
+        self.blocking_spawner.queue_depth()
+    }
+}

tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs

@@ -0,0 +1,26 @@
+use super::Handle;
+
+use crate::runtime::Dump;
+
+impl Handle {
+    pub(crate) async fn dump(&self) -> Dump {
+        let trace_status = &self.shared.trace_status;
+
+        // If a dump is in progress, block.
+        trace_status.start_trace_request(&self).await;
+
+        let result = loop {
+            if let Some(result) = trace_status.take_result() {
+                break result;
+            } else {
+                self.notify_all();
+                trace_status.result_ready.notified().await;
+            }
+        };
+
+        // Allow other queued dumps to proceed.
+        trace_status.end_trace_request(&self).await;
+
+        result
+    }
+}

tokio/src/runtime/scheduler/multi_thread/mod.rs

@@ -24,9 +24,17 @@ mod worker;
 pub(crate) use worker::{Context, Launch, Shared};
 
 cfg_taskdump! {
+    mod trace;
+    use trace::TraceStatus;
+
     pub(crate) use worker::Synced;
 }
 
+cfg_not_taskdump! {
+    mod trace_mock;
+    use trace_mock::TraceStatus;
+}
+
 pub(crate) use worker::block_in_place;
 
 use crate::loom::sync::Arc;

tokio/src/runtime/scheduler/multi_thread/trace.rs

@@ -0,0 +1,61 @@
+use crate::loom::sync::atomic::{AtomicBool, Ordering};
+use crate::loom::sync::{Barrier, Mutex};
+use crate::runtime::dump::Dump;
+use crate::runtime::scheduler::multi_thread::Handle;
+use crate::sync::notify::Notify;
+
+/// Tracing status of the worker.
+pub(super) struct TraceStatus {
+    pub(super) trace_requested: AtomicBool,
+    pub(super) trace_start: Barrier,
+    pub(super) trace_end: Barrier,
+    pub(super) result_ready: Notify,
+    pub(super) trace_result: Mutex<Option<Dump>>,
+}
+
+impl TraceStatus {
+    pub(super) fn new(remotes_len: usize) -> Self {
+        Self {
+            trace_requested: AtomicBool::new(false),
+            trace_start: Barrier::new(remotes_len),
+            trace_end: Barrier::new(remotes_len),
+            result_ready: Notify::new(),
+            trace_result: Mutex::new(None),
+        }
+    }
+
+    pub(super) fn trace_requested(&self) -> bool {
+        self.trace_requested.load(Ordering::Relaxed)
+    }
+
+    pub(super) async fn start_trace_request(&self, handle: &Handle) {
+        while self
+            .trace_requested
+            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+            .is_err()
+        {
+            handle.notify_all();
+            crate::task::yield_now().await;
+        }
+    }
+
+    pub(super) fn stash_result(&self, dump: Dump) {
+        let _ = self.trace_result.lock().insert(dump);
+        self.result_ready.notify_one();
+    }
+
+    pub(super) fn take_result(&self) -> Option<Dump> {
+        self.trace_result.lock().take()
+    }
+
+    pub(super) async fn end_trace_request(&self, handle: &Handle) {
+        while self
+            .trace_requested
+            .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed)
+            .is_err()
+        {
+            handle.notify_all();
+            crate::task::yield_now().await;
+        }
+    }
+}

tokio/src/runtime/scheduler/multi_thread/trace_mock.rs

@@ -0,0 +1,11 @@
+pub(super) struct TraceStatus {}
+
+impl TraceStatus {
+    pub(super) fn new(_: usize) -> Self {
+        Self {}
+    }
+
+    pub(super) fn trace_requested(&self) -> bool {
+        false
+    }
+}

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -60,7 +60,7 @@ use crate::loom::sync::{Arc, Mutex};
 use crate::runtime;
 use crate::runtime::context;
 use crate::runtime::scheduler::multi_thread::{
-    idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, Unparker,
+    idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
 };
 use crate::runtime::scheduler::{inject, Defer, Lock};
 use crate::runtime::task::OwnedTasks;
@@ -74,8 +74,16 @@ use std::cell::RefCell;
 use std::task::Waker;
 use std::time::Duration;
 
+cfg_metrics! {
+    mod metrics;
+}
+
 cfg_taskdump! {
-    use crate::loom::sync::Barrier;
+    mod taskdump;
+}
+
+cfg_not_taskdump! {
+    mod taskdump_mock;
 }
 
 /// A scheduler worker
@@ -214,82 +222,6 @@ pub(crate) struct Context {
 /// Starts the workers
 pub(crate) struct Launch(Vec<Arc<Worker>>);
 
-cfg_not_taskdump! {
-    pub(super) struct TraceStatus {}
-
-    impl TraceStatus {
-        fn new(_: usize) -> Self {
-            Self {}
-        }
-
-        fn trace_requested(&self) -> bool {
-            false
-        }
-    }
-}
-
-cfg_taskdump! {
-    use crate::sync::notify::Notify;
-    use crate::runtime::dump::Dump;
-    use crate::loom::sync::atomic::{AtomicBool, Ordering};
-
-    /// Tracing status of the worker.
-    pub(super) struct TraceStatus {
-        pub(super) trace_requested: AtomicBool,
-        trace_start: Barrier,
-        trace_end: Barrier,
-        pub(super) result_ready: Notify,
-        pub(super) trace_result: Mutex<Option<Dump>>,
-    }
-
-    impl TraceStatus {
-        fn new(remotes_len: usize) -> Self {
-            Self {
-                trace_requested: AtomicBool::new(false),
-                trace_start: Barrier::new(remotes_len),
-                trace_end: Barrier::new(remotes_len),
-                result_ready: Notify::new(),
-                trace_result: Mutex::new(None),
-            }
-        }
-
-        fn trace_requested(&self) -> bool {
-            self.trace_requested.load(Ordering::Relaxed)
-        }
-
-        pub(super) async fn start_trace_request(&self, handle: &Handle) {
-            while self.trace_requested.compare_exchange(false,
-                true,
-                Ordering::Acquire,
-                Ordering::Relaxed).is_err()
-            {
-                handle.notify_all();
-                crate::task::yield_now().await;
-            }
-        }
-
-        fn stash_result(&self, dump: Dump) {
-            let _ = self.trace_result.lock().insert(dump);
-            self.result_ready.notify_one();
-        }
-
-        pub(super) fn take_result(&self) -> Option<Dump> {
-            self.trace_result.lock().take()
-        }
-
-        pub(super) async fn end_trace_request(&self, handle: &Handle) {
-            while self.trace_requested.compare_exchange(true,
-                false,
-                Ordering::Acquire,
-                Ordering::Relaxed).is_err()
-            {
-                handle.notify_all();
-                crate::task::yield_now().await;
-            }
-        }
-    }
-}
-
 /// Running a task may consume the core. If the core is still available when
 /// running the task completes, it is returned. Otherwise, the worker will need
 /// to stop processing.
@@ -1200,64 +1132,6 @@ impl Handle {
         }
     }
 
-    cfg_not_taskdump! {
-        fn trace_core(&self, core: Box<Core>) -> Box<Core> {
-                core
-        }
-    }
-
-    cfg_taskdump! {
-        fn trace_core(&self, mut core: Box<Core>) -> Box<Core> {
-            use crate::runtime::dump;
-            use task::trace::trace_multi_thread;
-
-            core.is_traced = false;
-
-            if core.is_shutdown {
-                return core;
-            }
-
-            // wait for other workers, or timeout without tracing
-            let timeout = Duration::from_millis(250); // a _very_ generous timeout
-            let barrier = if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) {
-                barrier
-            } else {
-                // don't attempt to trace
-                return core;
-            };
-
-            if !barrier.is_leader() {
-                // wait for leader to finish tracing
-                self.shared.trace_status.trace_end.wait();
-                return core;
-            }
-
-            // trace
-
-            let owned = &self.shared.owned;
-            let mut local = self.shared.steal_all();
-            let synced = &self.shared.synced;
-            let injection = &self.shared.inject;
-
-            // safety: `trace_multi_thread` is invoked with the same `synced` that `injection`
-            // was created with.
-            let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) }
-                .into_iter()
-                .map(dump::Task::new)
-                .collect();
-
-            let result = dump::Dump::new(traces);
-
-            // stash the result
-            self.shared.trace_status.stash_result(result);
-
-            // allow other workers to proceed
-            self.shared.trace_status.trace_end.wait();
-
-            core
-        }
-    }
-
     fn ptr_eq(&self, other: &Handle) -> bool {
         std::ptr::eq(self, other)
     }
@@ -1308,41 +1182,6 @@ fn with_current<R>(f: impl FnOnce(Option<&Context>) -> R) -> R {
     })
 }
 
-cfg_metrics! {
-    impl Shared {
-        pub(super) fn injection_queue_depth(&self) -> usize {
-            self.inject.len()
-        }
-
-        pub(super) fn worker_local_queue_depth(&self, worker: usize) -> usize {
-            self.remotes[worker].steal.len()
-        }
-    }
-}
-
-cfg_taskdump! {
-    impl Shared {
-        /// Steal all tasks from remotes into a single local queue.
-        pub(super) fn steal_all(&self) -> super::queue::Local<Arc<Handle>> {
-            let (_steal, mut local) = super::queue::local();
-
-            let worker_metrics = WorkerMetrics::new();
-            let mut stats = Stats::new(&worker_metrics);
-
-            for remote in self.remotes.iter() {
-                let steal = &remote.steal;
-                while !steal.is_empty() {
-                    if let Some(task) = steal.steal_into(&mut local, &mut stats) {
-                        local.push_back([task].into_iter());
-                    }
-                }
-            }
-
-            local
-        }
-    }
-}
-
 // `u32::abs_diff` is not available on Tokio's MSRV.
 fn abs_diff(a: u32, b: u32) -> u32 {
     if a > b {

tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs

@@ -0,0 +1,11 @@
+use super::Shared;
+
+impl Shared {
+    pub(crate) fn injection_queue_depth(&self) -> usize {
+        self.inject.len()
+    }
+
+    pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
+        self.remotes[worker].steal.len()
+    }
+}

tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs

@@ -0,0 +1,79 @@
+use super::{Core, Handle, Shared};
+
+use crate::loom::sync::Arc;
+use crate::runtime::scheduler::multi_thread::Stats;
+use crate::runtime::task::trace::trace_multi_thread;
+use crate::runtime::{dump, WorkerMetrics};
+
+use std::time::Duration;
+
+impl Handle {
+    pub(super) fn trace_core(&self, mut core: Box<Core>) -> Box<Core> {
+        core.is_traced = false;
+
+        if core.is_shutdown {
+            return core;
+        }
+
+        // wait for other workers, or timeout without tracing
+        let timeout = Duration::from_millis(250); // a _very_ generous timeout
+        let barrier =
+            if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) {
+                barrier
+            } else {
+                // don't attempt to trace
+                return core;
+            };
+
+        if !barrier.is_leader() {
+            // wait for leader to finish tracing
+            self.shared.trace_status.trace_end.wait();
+            return core;
+        }
+
+        // trace
+
+        let owned = &self.shared.owned;
+        let mut local = self.shared.steal_all();
+        let synced = &self.shared.synced;
+        let injection = &self.shared.inject;
+
+        // safety: `trace_multi_thread` is invoked with the same `synced` that `injection`
+        // was created with.
+        let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) }
+            .into_iter()
+            .map(dump::Task::new)
+            .collect();
+
+        let result = dump::Dump::new(traces);
+
+        // stash the result
+        self.shared.trace_status.stash_result(result);
+
+        // allow other workers to proceed
+        self.shared.trace_status.trace_end.wait();
+
+        core
+    }
+}
+
+impl Shared {
+    /// Steal all tasks from remotes into a single local queue.
+    pub(super) fn steal_all(&self) -> super::queue::Local<Arc<Handle>> {
+        let (_steal, mut local) = super::queue::local();
+
+        let worker_metrics = WorkerMetrics::new();
+        let mut stats = Stats::new(&worker_metrics);
+
+        for remote in self.remotes.iter() {
+            let steal = &remote.steal;
+            while !steal.is_empty() {
+                if let Some(task) = steal.steal_into(&mut local, &mut stats) {
+                    local.push_back([task].into_iter());
+                }
+            }
+        }
+
+        local
+    }
+}

tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs

@@ -0,0 +1,7 @@
+use super::{Core, Handle};
+
+impl Handle {
+    pub(super) fn trace_core(&self, core: Box<Core>) -> Box<Core> {
+        core
+    }
+}