Skip to content

Commit

Permalink
Improve code around SGX waitqueue
Browse files Browse the repository at this point in the history
Followed up of d36e390d8176babedcf326581959958d447170cd
See rust-lang/rust#109732 (comment)
for more details.

Co-authored-by: Jethro Beekman <jethro@fortanix.com>
  • Loading branch information
Urgau and Jethro Beekman committed May 11, 2023
1 parent c296c9c commit a15064d
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions std/src/sys/sgx/waitqueue/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -202,12 +202,18 @@ impl WaitQueue {
pub fn notify_one<T>(
mut guard: SpinMutexGuard<'_, WaitVariable<T>>,
) -> Result<WaitGuard<'_, T>, SpinMutexGuard<'_, WaitVariable<T>>> {
// 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)
Expand All @@ -223,13 +229,17 @@ impl WaitQueue {
pub fn notify_all<T>(
mut guard: SpinMutexGuard<'_, WaitVariable<T>>,
) -> Result<WaitGuard<'_, T>, SpinMutexGuard<'_, WaitVariable<T>>> {
// 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() {
count += 1;
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 {
Expand Down

0 comments on commit a15064d

Please sign in to comment.