diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index c49c08932b2..f266bea014e 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -138,6 +138,13 @@ pub(super) fn budget(f: impl FnOnce(&Cell) -> R) -> Result bool { + CONTEXT.try_with(|ctx| { + ctx.thread_id.get().is_some() + }).unwrap_or(false) + } + pub(crate) fn thread_id() -> Result { CONTEXT.try_with(|ctx| { match ctx.thread_id.get() { diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 11dfa274ee3..a7c0bd7826a 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -1112,6 +1112,14 @@ impl LocalState { #[track_caller] fn assert_called_from_owner_thread(&self) { + // Even when the thread is being destroyed, local data may still be initialized on the SGX + // platform. So calling `context::thread_id()` will always return a thread id. In line with + // the `debug_assert` below, we ignore such cases on the SGX platform. + #[cfg(target_env = "sgx")] + if !context::has_thread_id() { + return; + } + // FreeBSD has some weirdness around thread-local destruction. // TODO: remove this hack when thread id is cleaned up #[cfg(not(any(target_os = "openbsd", target_os = "freebsd")))] diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs index 5b8f47234f9..a457af0b0fe 100644 --- a/tokio/tests/task_local_set.rs +++ b/tokio/tests/task_local_set.rs @@ -307,7 +307,6 @@ fn join_local_future_elsewhere() { } // Tests for -#[cfg(not(target_env = "sgx"))] // Fails on SGX! #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn localset_in_thread_local() { @@ -575,7 +574,6 @@ async fn spawn_wakes_localset() { } #[test] -#[cfg(not(target_env = "sgx"))] // Fails on SGX! fn store_local_set_in_thread_local_with_runtime() { use tokio::runtime::Runtime;