tokio: use const initialized thread locals where possible (#4677)

Author: a1phyr

tokio/Cargo.toml

@@ -100,6 +100,9 @@ time = []
 # a few releases.
 stats = []
 
+[build-dependencies]
+autocfg = "1.1"
+
 [dependencies]
 tokio-macros = { version = "1.7.0", path = "../tokio-macros", optional = true }
 

tokio/build.rs

@@ -0,0 +1,22 @@
+use autocfg::AutoCfg;
+
+fn main() {
+    match AutoCfg::new() {
+        Ok(ac) => {
+            // Const-initialized thread locals were stabilized in 1.59
+            if ac.probe_rustc_version(1, 59) {
+                autocfg::emit("tokio_const_thread_local")
+            }
+        }
+
+        Err(e) => {
+            // If we couldn't detect the compiler version and features, just
+            // print a warning. This isn't a fatal error: we can still build
+            // Tokio, we just can't enable cfgs automatically.
+            println!(
+                "cargo:warning=tokio: failed to detect compiler features: {}",
+                e
+            );
+        }
+    }
+}

tokio/src/coop.rs

@@ -32,7 +32,7 @@
 use std::cell::Cell;
 
 thread_local! {
-    static CURRENT: Cell<Budget> = Cell::new(Budget::unconstrained());
+    static CURRENT: Cell<Budget> = const { Cell::new(Budget::unconstrained()) };
 }
 
 /// Opaque type tracking the amount of "work" a task may still do before

tokio/src/macros/scoped_tls.rs

@@ -10,7 +10,7 @@ macro_rules! scoped_thread_local {
         $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty>
             = $crate::macros::scoped_tls::ScopedKey {
                 inner: {
-                    thread_local!(static FOO: ::std::cell::Cell<*const ()> = {
+                    thread_local!(static FOO: ::std::cell::Cell<*const ()> = const {
                         std::cell::Cell::new(::std::ptr::null())
                     });
                     &FOO

tokio/src/macros/thread_local.rs

@@ -1,4 +1,28 @@
 #[cfg(all(loom, test))]
 macro_rules! thread_local {
+    ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty = const { $expr:expr } $(;)?) => {
+        loom::thread_local! {
+            $(#[$attrs])*
+            $vis static $name: $ty = $expr;
+        }
+    };
+
     ($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } }
 }
+
+#[cfg(all(tokio_const_thread_local, not(all(loom, test))))]
+macro_rules! thread_local {
+    ($($tts:tt)+) => { ::std::thread_local!{ $($tts)+ } }
+}
+
+#[cfg(all(not(tokio_const_thread_local), not(all(loom, test))))]
+macro_rules! thread_local {
+    ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty = const { $expr:expr } $(;)?) => {
+        ::std::thread_local! {
+            $(#[$attrs])*
+            $vis static $name: $ty = $expr;
+        }
+    };
+
+    ($($tts:tt)+) => { ::std::thread_local!{ $($tts)+ } }
+}

tokio/src/runtime/context.rs

@@ -4,7 +4,7 @@ use crate::runtime::{Handle, TryCurrentError};
 use std::cell::RefCell;
 
 thread_local! {
-    static CONTEXT: RefCell<Option<Handle>> = RefCell::new(None)
+    static CONTEXT: RefCell<Option<Handle>> = const { RefCell::new(None) }
 }
 
 pub(crate) fn try_current() -> Result<Handle, crate::runtime::TryCurrentError> {

tokio/src/runtime/enter.rs

@@ -17,7 +17,7 @@ impl EnterContext {
     }
 }
 
-thread_local!(static ENTERED: Cell<EnterContext> = Cell::new(EnterContext::NotEntered));
+thread_local!(static ENTERED: Cell<EnterContext> = const { Cell::new(EnterContext::NotEntered) });
 
 /// Represents an executor context.
 pub(crate) struct Enter {

tokio/src/task/task_local.rs

@@ -48,6 +48,22 @@ macro_rules! task_local {
 }
 
 #[doc(hidden)]
+#[cfg(tokio_const_thread_local)]
+#[macro_export]
+macro_rules! __task_local_inner {
+    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
+        $vis static $name: $crate::task::LocalKey<$t> = {
+            std::thread_local! {
+                static __KEY: std::cell::RefCell<Option<$t>> = const { std::cell::RefCell::new(None) };
+            }
+
+            $crate::task::LocalKey { inner: __KEY }
+        };
+    };
+}
+
+#[doc(hidden)]
+#[cfg(not(tokio_const_thread_local))]
 #[macro_export]
 macro_rules! __task_local_inner {
     ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {