task: use NonZeroU64 for task::Id (#6733)

tokio/src/runtime/task/id.rs

@@ -1,6 +1,6 @@
 use crate::runtime::context;
 
-use std::fmt;
+use std::{fmt, num::NonZeroU64};
 
 /// An opaque ID that uniquely identifies a task relative to all other currently
 /// running tasks.
@@ -24,7 +24,7 @@ use std::fmt;
 #[cfg_attr(docsrs, doc(cfg(all(feature = "rt", tokio_unstable))))]
 #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
 #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
-pub struct Id(pub(crate) u64);
+pub struct Id(pub(crate) NonZeroU64);
 
 /// Returns the [`Id`] of the currently running task.
 ///
@@ -78,21 +78,22 @@ impl Id {
         use crate::loom::sync::atomic::StaticAtomicU64;
 
         #[cfg(all(test, loom))]
-        {
-            crate::loom::lazy_static! {
-                static ref NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
-            }
-            Self(NEXT_ID.fetch_add(1, Relaxed))
+        crate::loom::lazy_static! {
+            static ref NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
         }
 
         #[cfg(not(all(test, loom)))]
-        {
-            static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
-            Self(NEXT_ID.fetch_add(1, Relaxed))
+        static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
+
+        loop {
+            let id = NEXT_ID.fetch_add(1, Relaxed);
+            if let Some(id) = NonZeroU64::new(id) {
+                return Self(id);
+            }
         }
     }
 
     pub(crate) fn as_u64(&self) -> u64 {
-        self.0
+        self.0.get()
     }
 }

tokio/src/runtime/task/mod.rs

@@ -532,6 +532,6 @@ unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
     unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
         // SAFETY: The caller guarantees that `target` points at a valid task.
         let task_id = unsafe { Header::get_id(target) };
-        task_id.0 as usize
+        task_id.0.get() as usize
     }
 }