From ee1c940709ad88f83daea29612a924524e37d173 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Mon, 27 Feb 2023 10:21:42 -0800
Subject: [PATCH] time: Improve `Instant::now()` perf with test-util (#5513)

The test-util feature flag is only intended to be used with tests.
However, it is possible to enable it in release mode accidentally. This
patch reduces the overhead of `Instant::now()` when the `test-util`
feature flag is enabled but `time::pause()` is not called.

The optimization is implemented by adding a static atomic flag that
tracks if `time::pause()` has ever been called. In `Instant::now()`, the
atomic flag is first checked before the thread-local and mutex are
accessed.
---
 benches/Cargo.toml                  |  8 ++++++++
 benches/time_now.rs                 | 25 +++++++++++++++++++++++++
 tokio/src/time/clock.rs             | 18 ++++++++++++++++++
 tokio/tests/rt_time_start_paused.rs | 14 ++++++++++++++
 4 files changed, 65 insertions(+)
 create mode 100644 benches/time_now.rs
 create mode 100644 tokio/tests/rt_time_start_paused.rs

diff --git a/benches/Cargo.toml b/benches/Cargo.toml
index 10c32de0f..7b5d2e105 100644
--- a/benches/Cargo.toml
+++ b/benches/Cargo.toml
@@ -4,6 +4,9 @@ version = "0.0.0"
 publish = false
 edition = "2018"
 
+[features]
+test-util = ["tokio/test-util"]
+
 [dependencies]
 tokio = { version = "1.5.0", path = "../tokio", features = ["full"] }
 bencher = "0.1.5"
@@ -67,3 +70,8 @@ harness = false
 name = "copy"
 path = "copy.rs"
 harness = false
+
+[[bench]]
+name = "time_now"
+path = "time_now.rs"
+harness = false
diff --git a/benches/time_now.rs b/benches/time_now.rs
new file mode 100644
index 000000000..814628569
--- /dev/null
+++ b/benches/time_now.rs
@@ -0,0 +1,25 @@
+//! Benchmark spawning a task onto the basic and threaded Tokio executors.
+//! This essentially measure the time to enqueue a task in the local and remote
+//! case.
+
+#[macro_use]
+extern crate bencher;
+
+use bencher::{black_box, Bencher};
+
+fn time_now_current_thread(bench: &mut Bencher) {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    bench.iter(|| {
+        rt.block_on(async {
+            black_box(tokio::time::Instant::now());
+        })
+    })
+}
+
+bencher::benchmark_group!(time_now, time_now_current_thread,);
+
+bencher::benchmark_main!(time_now);
diff --git a/tokio/src/time/clock.rs b/tokio/src/time/clock.rs
index 1e273554e..091cf4b19 100644
--- a/tokio/src/time/clock.rs
+++ b/tokio/src/time/clock.rs
@@ -30,6 +30,8 @@ cfg_not_test_util! {
 cfg_test_util! {
     use crate::time::{Duration, Instant};
     use crate::loom::sync::Mutex;
+    use crate::loom::sync::atomic::Ordering;
+    use std::sync::atomic::AtomicBool as StdAtomicBool;
 
     cfg_rt! {
         #[track_caller]
@@ -65,6 +67,15 @@ cfg_test_util! {
         inner: Mutex<Inner>,
     }
 
+    // Used to track if the clock was ever paused. This is an optimization to
+    // avoid touching the mutex if `test-util` was accidentally enabled in
+    // release mode.
+    //
+    // A static is used so we can avoid accessing the thread-local as well. The
+    // `std` AtomicBool is used directly because loom does not support static
+    // atomics.
+    static DID_PAUSE_CLOCK: StdAtomicBool = StdAtomicBool::new(false);
+
     #[derive(Debug)]
     struct Inner {
         /// True if the ability to pause time is enabled.
@@ -199,6 +210,10 @@ cfg_test_util! {
 
     /// Returns the current instant, factoring in frozen time.
     pub(crate) fn now() -> Instant {
+        if !DID_PAUSE_CLOCK.load(Ordering::Acquire) {
+            return Instant::from_std(std::time::Instant::now());
+        }
+
         with_clock(|maybe_clock| {
             Ok(if let Some(clock) = maybe_clock {
                 clock.now()
@@ -241,6 +256,9 @@ cfg_test_util! {
                         This is the default Runtime used by `#[tokio::test].");
             }
 
+            // Track that we paused the clock
+            DID_PAUSE_CLOCK.store(true, Ordering::Release);
+
             let elapsed = match inner.unfrozen.as_ref() {
                 Some(v) => v.elapsed(),
                 None => return Err("time is already frozen")
diff --git a/tokio/tests/rt_time_start_paused.rs b/tokio/tests/rt_time_start_paused.rs
new file mode 100644
index 000000000..283f4748a
--- /dev/null
+++ b/tokio/tests/rt_time_start_paused.rs
@@ -0,0 +1,14 @@
+#![cfg(all(feature = "full"))]
+
+use tokio::time::{Duration, Instant};
+
+#[tokio::test(start_paused = true)]
+async fn test_start_paused() {
+    let now = Instant::now();
+
+    // Pause a few times w/ std sleep and ensure `now` stays the same
+    for _ in 0..5 {
+        std::thread::sleep(Duration::from_millis(1));
+        assert_eq!(now, Instant::now());
+    }
+}