Improve code around SGX waitqueue

Followed up of d36e390d8176babedcf326581959958d447170cd See rust-lang/rust#109732 (comment) for more details. Co-authored-by: Jethro Beekman <[email protected]>
tcdi · May 11, 2023 · a15064d · a15064d
1 parent c296c9c
commit a15064d
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/std/src/sys/sgx/waitqueue/mod.rs b/std/src/sys/sgx/waitqueue/mod.rs
@@ -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)
@@ -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 {