make task_id_counter thread local in loom tests

tokio/src/runtime/context.rs

@@ -6,6 +6,8 @@ use std::cell::Cell;
 #[cfg(any(feature = "rt", feature = "macros"))]
 use crate::util::rand::FastRand;
 
+#[cfg(loom)]
+use crate::loom::sync::atomic::StaticAtomicU64;
 cfg_rt! {
     mod blocking;
     pub(crate) use blocking::{disallow_block_in_place, try_enter_blocking_region, BlockingRegionGuard};
@@ -72,6 +74,10 @@ struct Context {
         any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
     ))]
     trace: trace::Context,
+
+    /// In loom tests, we make task_id_counter in thread local for ensuring each loom thread is deterministic
+    #[cfg(loom)]
+    task_id_counter: StaticAtomicU64,
 }
 
 tokio_thread_local! {
@@ -117,6 +123,8 @@ tokio_thread_local! {
                 )
             ))]
             trace: trace::Context::new(),
+            #[cfg(loom)]
+            task_id_counter: StaticAtomicU64::new(1),
         }
     }
 }
@@ -135,6 +143,14 @@ pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, A
     CONTEXT.try_with(|ctx| f(&ctx.budget))
 }
 
+#[cfg(loom)]
+pub(super) fn next_task_id() -> Result<u64, AccessError> {
+    CONTEXT.try_with(|ctx| {
+        ctx.task_id_counter
+            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+    })
+}
+
 cfg_rt! {
     use crate::runtime::ThreadId;
 

tokio/src/runtime/task/id.rs

@@ -74,11 +74,17 @@ impl fmt::Display for Id {
 
 impl Id {
     pub(crate) fn next() -> Self {
-        use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64};
+        #[cfg(loom)]
+        return Self(context::next_task_id().unwrap());
 
-        static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
+        #[cfg(not(loom))]
+        {
+            use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64};
 
-        Self(NEXT_ID.fetch_add(1, Relaxed))
+            static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
+
+            Self(NEXT_ID.fetch_add(1, Relaxed))
+        }
     }
 
     pub(crate) fn as_u64(&self) -> u64 {