Skip to content

Commit

Permalink
make IDs always be 64-bit
Browse files Browse the repository at this point in the history
Signed-off-by: Eliza Weisman <eliza@buoyant.io>
  • Loading branch information
hawkw committed Apr 22, 2022
1 parent 47a26c5 commit 09254f9
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 12 deletions.
2 changes: 1 addition & 1 deletion tokio/Cargo.toml
Expand Up @@ -65,7 +65,7 @@ process = [
"winapi/threadpoollegacyapiset",
]
# Includes basic task execution capabilities
rt = []
rt = ["once_cell"]
rt-multi-thread = [
"num_cpus",
"rt",
Expand Down
36 changes: 30 additions & 6 deletions tokio/src/runtime/task/mod.rs
Expand Up @@ -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)]
Expand Down Expand Up @@ -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
}
}
4 changes: 2 additions & 2 deletions tokio/src/task/local.rs
Expand Up @@ -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`");
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion tokio/src/task/spawn.rs
Expand Up @@ -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)
}
}
4 changes: 2 additions & 2 deletions tokio/src/util/trace.rs
Expand Up @@ -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!(
Expand Down Expand Up @@ -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
}
Expand Down

0 comments on commit 09254f9

Please sign in to comment.