rt: add a method to retrieve task id

Author: Abutalib Aghayev

tokio/src/runtime/context.rs

@@ -1,4 +1,5 @@
 use crate::runtime::coop;
+use crate::runtime::task::Id;
 
 use std::cell::Cell;
 
@@ -17,6 +18,7 @@ struct Context {
     /// Handle to the runtime scheduler running on the current thread.
     #[cfg(feature = "rt")]
     scheduler: RefCell<Option<scheduler::Handle>>,
+    current_task_id: Cell<Option<Id>>,
 
     #[cfg(any(feature = "rt", feature = "macros"))]
     rng: FastRand,
@@ -31,6 +33,7 @@ tokio_thread_local! {
         Context {
             #[cfg(feature = "rt")]
             scheduler: RefCell::new(None),
+            current_task_id: Cell::new(None),
 
             #[cfg(any(feature = "rt", feature = "macros"))]
             rng: FastRand::new(RngSeed::new()),
@@ -85,6 +88,17 @@ cfg_rt! {
 
     pub(crate) struct DisallowBlockInPlaceGuard(bool);
 
+    pub(crate) fn set_current_task_id(id: Option<Id>) {
+        CONTEXT.with(|ctx| ctx.current_task_id.replace(id));
+    }
+
+    pub(crate) fn current_task_id() -> Id {
+        match CONTEXT.try_with(|ctx| ctx.current_task_id.get()) {
+            Ok(Some(id)) => id,
+            _ => panic!("tried to get task id from outside of the task"),
+        }
+    }
+
     pub(crate) fn try_current() -> Result<scheduler::Handle, TryCurrentError> {
         match CONTEXT.try_with(|ctx| ctx.scheduler.borrow().clone()) {
             Ok(Some(handle)) => Ok(handle),

tokio/src/runtime/task/harness.rs

@@ -1,8 +1,9 @@
 use crate::future::Future;
+use crate::runtime::context;
 use crate::runtime::task::core::{Cell, Core, CoreStage, Header, Trailer};
 use crate::runtime::task::state::Snapshot;
 use crate::runtime::task::waker::waker_ref;
-use crate::runtime::task::{JoinError, Notified, Schedule, Task};
+use crate::runtime::task::{Id, JoinError, Notified, Schedule, Task};
 
 use std::mem;
 use std::mem::ManuallyDrop;
@@ -476,6 +477,19 @@ fn poll_future<T: Future, S: Schedule>(
                 self.core.drop_future_or_output();
             }
         }
+        struct TaskIdGuard {}
+        impl TaskIdGuard {
+            fn new(id: Id) -> Self {
+                context::set_current_task_id(Some(id));
+                TaskIdGuard {}
+            }
+        }
+        impl Drop for TaskIdGuard {
+            fn drop(&mut self) {
+                context::set_current_task_id(None);
+            }
+        }
+        let _task_id_guard = TaskIdGuard::new(id);
         let guard = Guard { core };
         let res = guard.core.poll(cx);
         mem::forget(guard);

tokio/src/runtime/task/mod.rs

@@ -201,10 +201,21 @@ use std::{fmt, mem};
 /// [unstable]: crate#unstable-features
 #[cfg_attr(docsrs, doc(cfg(all(feature = "rt", tokio_unstable))))]
 #[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)]
+#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
 pub struct Id(u64);
 
+/// Returns the `Id` of the task.
+///
+/// # Panics
+///
+/// This function panics if called from outside a task.
+///
+#[allow(unreachable_pub)]
+pub fn id() -> Id {
+    use crate::runtime::context;
+    context::current_task_id()
+}
+
 /// An owned handle to the task, tracked by ref count.
 #[repr(transparent)]
 pub(crate) struct Task<S: 'static> {

tokio/src/task/mod.rs

@@ -319,6 +319,7 @@ cfg_rt! {
 
     cfg_unstable! {
         pub use crate::runtime::task::Id;
+        pub use crate::runtime::task::id;
     }
 
     cfg_trace! {

tokio/tests/task_local.rs

@@ -116,3 +116,64 @@ async fn task_local_available_on_completion_drop() {
     assert_eq!(rx.await.unwrap(), 42);
     h.await.unwrap();
 }
+
+#[tokio::test(flavor = "current_thread")]
+async fn task_id() {
+    use tokio::task;
+
+    let handle = tokio::spawn(async { println!("task id: {}", task::id()) });
+
+    handle.await.unwrap();
+}
+
+#[cfg(tokio_unstable)]
+#[tokio::test(flavor = "multi_thread")]
+async fn task_id_collision_multi_thread() {
+    use tokio::task;
+
+    let handle1 = tokio::spawn(async { task::id() });
+    let handle2 = tokio::spawn(async { task::id() });
+
+    let (id1, id2) = tokio::join!(handle1, handle2);
+    assert_ne!(id1.unwrap(), id2.unwrap());
+}
+
+#[cfg(tokio_unstable)]
+#[tokio::test(flavor = "current_thread")]
+async fn task_id_collision_current_thread() {
+    use tokio::task;
+
+    let handle1 = tokio::spawn(async { task::id() });
+    let handle2 = tokio::spawn(async { task::id() });
+
+    let (id1, id2) = tokio::join!(handle1, handle2);
+    assert_ne!(id1.unwrap(), id2.unwrap());
+}
+
+#[cfg(tokio_unstable)]
+#[tokio::test(flavor = "current_thread")]
+async fn task_ids_match_current_thread() {
+    use tokio::{sync::oneshot, task};
+
+    let (tx, rx) = oneshot::channel();
+    let handle = tokio::spawn(async {
+        let id = rx.await.unwrap();
+        assert_eq!(id, task::id());
+    });
+    tx.send(handle.id()).unwrap();
+    handle.await.unwrap();
+}
+
+#[cfg(tokio_unstable)]
+#[tokio::test(flavor = "multi_thread")]
+async fn task_ids_match_multi_thread() {
+    use tokio::{sync::oneshot, task};
+
+    let (tx, rx) = oneshot::channel();
+    let handle = tokio::spawn(async {
+        let id = rx.await.unwrap();
+        assert_eq!(id, task::id());
+    });
+    tx.send(handle.id()).unwrap();
+    handle.await.unwrap();
+}

tokio/tests/task_panic.rs

@@ -120,3 +120,16 @@ fn local_key_get_panic_caller() -> Result<(), Box<dyn Error>> {
 
     Ok(())
 }
+
+#[cfg(tokio_unstable)]
+#[test]
+fn task_id_handle_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(|| {
+        let _ = task::id();
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}