nordic-dev.net/embassy
2
0
Fork 0
mirror of https://github.com/embassy-rs/embassy.git synced 2024-11-22 06:42:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

embassy-sync: fix the data of LazyLock never drop

Browse Source 
Using `union` can save more space.

And the `MaybeUninit<T>` will never drop the T, when dropping the
`MaybeUninit<T>`. Fixed it.

Signed-off-by: wanglei <wllenyj@gmail.com>
This commit is contained in:
wanglei 2024-07-31 00:24:39 +08:00
parent e89ff7d129
commit 93696c912e
1 changed files with 31 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
embassy-sync/src/lazy_lock.rs
View File

@ -1,7 +1,7 @@
//! Synchronization primitive for initializing a value once, allowing others to get a reference to the value.
use core::cell::Cell;
use core::mem::MaybeUninit;
use core::cell::UnsafeCell;
use core::mem::ManuallyDrop;
use core::sync::atomic::{AtomicBool, Ordering};
/// The `LazyLock` is a synchronization primitive that allows for
@ -23,8 +23,12 @@ use core::sync::atomic::{AtomicBool, Ordering};
/// ```
pub struct LazyLock<T, F = fn() -> T> {
init: AtomicBool,
init_fn: Cell<Option<F>>,
data: Cell<MaybeUninit<T>>,
data: UnsafeCell<Data<T, F>>,
}
union Data<T, F> {
value: ManuallyDrop<T>,
f: ManuallyDrop<F>,
}
unsafe impl<T, F> Sync for LazyLock<T, F> {}
@ -34,8 +38,9 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
pub const fn new(init_fn: F) -> Self {
Self {
init: AtomicBool::new(false),
init_fn: Cell::new(Some(init_fn)),
data: Cell::new(MaybeUninit::zeroed()),
data: UnsafeCell::new(Data {
f: ManuallyDrop::new(init_fn),
}),
}
}
@ -44,7 +49,7 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
#[inline]
pub fn get(&self) -> &T {
self.ensure_init_fast();
unsafe { (*self.data.as_ptr()).assume_init_ref() }
unsafe { &(*self.data.get()).value }
}
/// Consume the `LazyLock`, returning the underlying value. The
@ -53,7 +58,10 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
#[inline]
pub fn into_inner(self) -> T {
self.ensure_init_fast();
unsafe { self.data.into_inner().assume_init() }
let this = ManuallyDrop::new(self);
let data = unsafe { core::ptr::read(&this.data) }.into_inner();
ManuallyDrop::into_inner(unsafe { data.value })
}
/// Initialize the `LazyLock` if it has not been initialized yet.
@ -75,10 +83,23 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
fn ensure_init(&self) {
critical_section::with(|_| {
if !self.init.load(Ordering::Acquire) {
let init_fn = self.init_fn.take().unwrap();
self.data.set(MaybeUninit::new(init_fn()));
let data = unsafe { &mut *self.data.get() };
let f = unsafe { ManuallyDrop::take(&mut data.f) };
let value = f();
data.value = ManuallyDrop::new(value);
self.init.store(true, Ordering::Release);
}
});
}
}
impl<T, F> Drop for LazyLock<T, F> {
fn drop(&mut self) {
if self.init.load(Ordering::Acquire) {
unsafe { ManuallyDrop::drop(&mut self.data.get_mut().value) };
} else {
unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) };
}
}
}