From f5aede9c822875345eb7eb468ecd9bc4568ed112 Mon Sep 17 00:00:00 2001 From: Urgau Date: Thu, 11 May 2023 10:53:16 +0200 Subject: Improve code around SGX waitqueue Followed up of d36e390d8176babedcf326581959958d447170cd See https://github.com/rust-lang/rust/pull/109732#issuecomment-1543574908 for more details. Co-authored-by: Jethro Beekman --- library/std/src/sys/sgx/waitqueue/mod.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) (limited to 'library/std/src/sys/sgx/waitqueue') diff --git a/library/std/src/sys/sgx/waitqueue/mod.rs b/library/std/src/sys/sgx/waitqueue/mod.rs index ca649ebd9d5..5e1d859ee99 100644 --- a/library/std/src/sys/sgx/waitqueue/mod.rs +++ b/library/std/src/sys/sgx/waitqueue/mod.rs @@ -202,12 +202,18 @@ impl WaitQueue { pub fn notify_one( mut guard: SpinMutexGuard<'_, WaitVariable>, ) -> Result, SpinMutexGuard<'_, WaitVariable>> { + // SAFETY: lifetime of the pop() return value is limited to the map + // closure (The closure return value is 'static). The underlying + // stack frame won't be freed until after the WaitGuard created below + // is dropped. unsafe { - if let Some(entry) = guard.queue.inner.pop() { + let tcs = guard.queue.inner.pop().map(|entry| -> Tcs { let mut entry_guard = entry.lock(); - let tcs = entry_guard.tcs; entry_guard.wake = true; - drop(entry_guard); + entry_guard.tcs + }); + + if let Some(tcs) = tcs { Ok(WaitGuard { mutex_guard: Some(guard), notified_tcs: NotifiedTcs::Single(tcs) }) } else { Err(guard) @@ -223,6 +229,9 @@ impl WaitQueue { pub fn notify_all( mut guard: SpinMutexGuard<'_, WaitVariable>, ) -> Result, SpinMutexGuard<'_, WaitVariable>> { + // SAFETY: lifetime of the pop() return values are limited to the + // while loop body. The underlying stack frames won't be freed until + // after the WaitGuard created below is dropped. unsafe { let mut count = 0; while let Some(entry) = guard.queue.inner.pop() { @@ -230,6 +239,7 @@ impl WaitQueue { let mut entry_guard = entry.lock(); entry_guard.wake = true; } + if let Some(count) = NonZeroUsize::new(count) { Ok(WaitGuard { mutex_guard: Some(guard), notified_tcs: NotifiedTcs::All { count } }) } else { -- cgit 1.4.1-3-g733a5