Repository: tokio-rs/tokio
File: tokio/src/sync/set_once.rs (Rust)

Code before:
            Some(unsafe { self.value.with_mut(|ptr| ptr::read(ptr).assume_init()) })
        } else {
            None
        }
    }

    /// Waits until the value is set.
    ///
    /// If the `SetOnce` is already initialized, it will return the value
    /// immediately.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe.
    pub async fn wait(&self) -> &T {
        loop {
            if let Some(val) = self.get() {
                return val;
            }

            let notify_fut = self.notify.notified();
            {
                // Taking the lock here ensures that a concurrent call to `set`
                // will see the creation of `notify_fut` in case the check
                // fails.
                let _guard = self.notify.lock();

                if self.value_set.load(Ordering::Relaxed) {
                    // SAFETY: the state is initialized
                    return unsafe { self.get_unchecked() };
                }
            }

            // wait until the value is set
            notify_fut.await;
        }
    }
}

// Since `get` gives us access to immutable references of the SetOnce, SetOnce
// can only be Sync if T is Sync, otherwise SetOnce would allow sharing
// references of !Sync values across threads. We need T to be Send in order for
// SetOnce to by Sync because we can use `set` on `&SetOnce<T>` to send values
// (of type T) across threads.
unsafe impl<T: Sync + Send> Sync for SetOnce<T> {}

// Access to SetOnce's value is guarded by the Atomic boolean flag
// and atomic operations on `value_set`, so as long as T itself is Send
// it's safe to send it to another thread
unsafe impl<T: Send> Send for SetOnce<T> {}


Review comment (suggestion):
Awaiting `notify_fut` also involves taking the lock. Perhaps we could avoid a unlock/lock here? Granted, it's more tricky to do here than in `set`.

Code after:
            // implementation does not try to drop the value again.
            self.value_set.store(false, Ordering::Relaxed);

            // SAFETY: The SetOnce is currently initialized, we can assume the
            // value is initialized and return that, when we return the value
            // we give the drop handler to the return scope.
            Some(unsafe { self.value.with_mut(|ptr| ptr::read(ptr).assume_init()) })
        } else {
            None
        }
    }

    /// Waits until the value is set.
    ///
    /// If the `SetOnce` is already initialized, it will return the value
    /// immediately.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe.
    pub async fn wait(&self) -> &T {
        loop {
            if let Some(val) = self.get() {
                return val;
            }

            let notify_fut = self.notify.notified();
            pin!(notify_fut);

            poll_fn(|cx| {
                // Register under the notify's internal lock.
                let ret = notify_fut.as_mut().poll(cx);
                if self.value_set.load(Ordering::Relaxed) {
                    return Poll::Ready(());
                }
                ret
            })
            .await;
        }
    }
}

// Since `get` gives us access to immutable references of the SetOnce, SetOnce
// can only be Sync if T is Sync, otherwise SetOnce would allow sharing
// references of !Sync values across threads. We need T to be Send in order for
// SetOnce to by Sync because we can use `set` on `&SetOnce<T>` to send values
// (of type T) across threads.
unsafe impl<T: Sync + Send> Sync for SetOnce<T> {}

// Access to SetOnce's value is guarded by the Atomic boolean flag
// and atomic operations on `value_set`, so as long as T itself is Send