diff --git a/tokio/src/runtime/task/id.rs b/tokio/src/runtime/task/id.rs index 82c8a7e7e90..9c3b1403ec2 100644 --- a/tokio/src/runtime/task/id.rs +++ b/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() } } diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index aa799bf2be1..8248fc0b591 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -532,6 +532,6 @@ unsafe impl sharded_list::ShardedListItem for Task { unsafe fn get_shard_id(target: NonNull) -> 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 } }