From 09254f9157c3dbcb4afaf54b8b1c30eb637604b4 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Fri, 22 Apr 2022 12:35:23 -0700
Subject: [PATCH] make IDs always be 64-bit

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 tokio/Cargo.toml              |  2 +-
 tokio/src/runtime/task/mod.rs | 36 +++++++++++++++++++++++++++++------
 tokio/src/task/local.rs       |  4 ++--
 tokio/src/task/spawn.rs       |  2 +-
 tokio/src/util/trace.rs       |  4 ++--
 5 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml
index 69ec3197e46..6bda46ef662 100644
--- a/tokio/Cargo.toml
+++ b/tokio/Cargo.toml
@@ -65,7 +65,7 @@ process = [
   "winapi/threadpoollegacyapiset",
 ]
 # Includes basic task execution capabilities
-rt = []
+rt = ["once_cell"]
 rt-multi-thread = [
   "num_cpus",
   "rt",
diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs
index a47268399a5..ac1ee157ea2 100644
--- a/tokio/src/runtime/task/mod.rs
+++ b/tokio/src/runtime/task/mod.rs
@@ -203,7 +203,7 @@ use std::{fmt, mem};
 #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
 // TODO(eliza): there's almost certainly no reason not to make this `Copy` as well...
 #[derive(Clone, Debug, Hash, Eq, PartialEq)]
-pub struct Id(usize);
+pub struct Id(u64);
 
 /// An owned handle to the task, tracked by ref count.
 #[repr(transparent)]
@@ -480,13 +480,37 @@ impl fmt::Display for Id {
 }
 
 impl Id {
-    pub(crate) fn next() -> Self {
-        use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
-        static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
-        Self(NEXT_ID.fetch_add(1, Relaxed))
+    // When 64-bit atomics are available, use a static `AtomicU64` counter to
+    // generate task IDs.
+    //
+    // Note(eliza): we _could_ just use `crate::loom::AtomicU64`, which switches
+    // between an atomic and mutex-based implementation here, rather than having
+    // two separate functions for targets with and without 64-bit atomics.
+    // However, because we can't use the mutex-based implementation in a static
+    // initializer directly, the 32-bit impl also has to use a `OnceCell`, and I
+    // thought it was nicer to avoid the `OnceCell` overhead on 64-bit
+    // platforms...
+    cfg_has_atomic_u64! {
+        pub(crate) fn next() -> Self {
+            use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
+            static NEXT_ID: AtomicU64 = AtomicU64::new(1);
+            Self(NEXT_ID.fetch_add(1, Relaxed))
+        }
+    }
+
+    cfg_not_has_atomic_u64! {
+        pub(crate) fn next() -> Self {
+            use once_cell::sync::Lazy;
+            use crate::loom::Mutex;
+            static NEXT_ID: Lazy<Mutex<u64>> = Lazy(|| Mutex::new(1));
+            let mut lock = NEXT_ID.lock();
+            let id = *lock;
+            lock += 1;
+            Self(id)
+        }
     }
 
-    pub(crate) fn as_usize(&self) -> usize {
+    pub(crate) fn as_u64(&self) -> u64 {
         self.0
     }
 }
diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs
index adfdca39233..32e376872f4 100644
--- a/tokio/src/task/local.rs
+++ b/tokio/src/task/local.rs
@@ -302,7 +302,7 @@ cfg_rt! {
           F::Output: 'static
     {
         let id = crate::runtime::task::Id::next();
-        let future = crate::util::trace::task(future, "local", name, id.as_usize());
+        let future = crate::util::trace::task(future, "local", name, id.as_u64());
         CURRENT.with(|maybe_cx| {
             let cx = maybe_cx
                 .expect("`spawn_local` called from outside of a `task::LocalSet`");
@@ -387,7 +387,7 @@ impl LocalSet {
         F::Output: 'static,
     {
         let id = crate::runtime::task::Id::next();
-        let future = crate::util::trace::task(future, "local", None, id.as_usize());
+        let future = crate::util::trace::task(future, "local", None, id.as_u64());
 
         let (handle, notified) = self
             .context
diff --git a/tokio/src/task/spawn.rs b/tokio/src/task/spawn.rs
index a86274a93ad..5a60f9d66e6 100644
--- a/tokio/src/task/spawn.rs
+++ b/tokio/src/task/spawn.rs
@@ -145,7 +145,7 @@ cfg_rt! {
         use crate::runtime::{task, context};
         let id = task::Id::next();
         let spawn_handle = context::spawn_handle().expect(CONTEXT_MISSING_ERROR);
-        let task = crate::util::trace::task(future, "task", name, id.as_usize());
+        let task = crate::util::trace::task(future, "task", name, id.as_u64());
         spawn_handle.spawn(task, id)
     }
 }
diff --git a/tokio/src/util/trace.rs b/tokio/src/util/trace.rs
index b2b78076004..76e8a6cbf55 100644
--- a/tokio/src/util/trace.rs
+++ b/tokio/src/util/trace.rs
@@ -10,7 +10,7 @@ cfg_trace! {
 
         #[inline]
         #[track_caller]
-        pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>, id: usize) -> Instrumented<F> {
+        pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>, id: u64) -> Instrumented<F> {
             use tracing::instrument::Instrument;
             let location = std::panic::Location::caller();
             let span = tracing::trace_span!(
@@ -92,7 +92,7 @@ cfg_time! {
 cfg_not_trace! {
     cfg_rt! {
         #[inline]
-        pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>, _: usize) -> F {
+        pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>, _: u64) -> F {
             // nop
             task
         }