Hello.
Is there a type in standard library that is similar to OnceLock type, except that instead of initializing the value once, it would allow to take the value once, thread-safely, without mutable reference and without consuming the object? A kind of thread-safe Option::take.
It's not hard to implement such a type yourself, something like this:
pub struct OnceTake<T> {
taken: AtomicBool,
value: UnsafeCell<T>,
}
impl<T> OnceTake<T> {
pub fn new(value: T) -> Self {
let value = UnsafeCell::new(value);
Self {
taken: AtomicBool::new(false),
value,
}
}
pub fn take(&self) -> Option<T> {
if self.taken.swap(true, Ordering::Acquire) {
let value = unsafe { core::mem::replace(&mut *self.value.get(), core::mem::zeroed()) };
Some(value)
} else {
None
}
}
}
But it would be great to have a standartized type, to have confidence in its safety, and to avoid reinventing the wheel.
One possible use case, is calling FnOnce inside Fn:
fn subscribe_to_event(callback: Box<dyn Fn() + Send + Sync>) {
// ...
}
fn foo(callback: Box<dyn FnOnce() + Send + Sync>) {
let once_take_callback = OnceTake::new(callback);
subscribe_to_event(Box::new(move || {
if let Some(cb) = once_take_callback.take() {
cb();
}
}));
}
For what it's worth, that example code is unsound. UnsafeCell still requires that its contents be initialized. You could put a MaybeUninit inside, or an Option.
As far as I'm aware, this doesn't exist in std. You could use something like Mutex<Option> to make it safe, but that'd be a little inefficient at runtime, and a lot of mental overhead. There may be a library for this, though I'm not aware of one.
2 Likes
I think that's just a Cell<Option> if you don't need thread safety.
fn subscribe_to_event(f: impl Fn()) {
f()
}
fn foo(callback: Box<dyn FnOnce()>) {
let mut once_take_callback = Cell::new(Some(callback));
subscribe_to_event(|| {
if let Some(cb) = once_take_callback.take() {
cb();
}
});
}
If you do need thread safety, I'd use Mutex<Option> as @Kyllingene said.
I forgot to add that subscribe_to_event requires Send + Sync callback. I modified the example to account for that. Cell wouldn't work in this case since it is !Sync.
The code in OP has a soundness hole because zeroed is not valid for every type. The classic example is NonZero<T>. Meanwhile, Cell<Option<T>> is sound:
So is UnsafeCell.
1 Like
In that case, Mutex<Option<T>> will serve just fine. But if you're looking for a more elegant, copy+paste solution, it's not too difficult:
use std::sync::atomic::{AtomicBool, Ordering};
use std::mem::MaybeUninit;
pub struct OnceTake<T> {
taken: AtomicBool,
inner: MaybeUninit<T>,
}
// SAFETY: The inner type can only be taken, and only once, through the thread-
// safe mechanism of the atomic lock.
unsafe impl<T> Send for OnceTake<T> {}
unsafe impl<T> Sync for OnceTake<T> {}
impl<T> OnceTake<T> {
pub fn new(data: T) -> Self {
Self {
taken: AtomicBool::new(false),
inner: MaybeUninit::new(data),
}
}
pub fn take(&self) -> Option<T> {
use std::ops::Not;
self.taken
.swap(true, Ordering::Acquire)
.not()
.then(|| {
// SAFETY: This is guaranteed to be the only time this is moved.
unsafe { self.inner.assume_init_read() }
})
}
}
impl<T> Drop for OnceTake<T> {
fn drop(&mut self) {
if !*self.taken.get_mut() {
// SAFETY: `self.taken == false` means that `self.inner.data` is initialized
unsafe { self.inner.assume_init_drop() };
}
}
}
This is, essentially, an atomic Option with only take implemented. It doesn't use an extra discriminant, it's cheaper than a Mutex, and it's Send + Sync.
1 Like
I wrote basically the same thing as well Rust Playground
I put the drop internals in a method so it can be used to read T.
I'm pretty sure it only needs Relaxed since the only thing that matters is that it happens once. It doesn't need to have any happens-before relationship with any other thread.
Seems right that you can unconditionally implement Sync (never mind, see comment below), but I don't think you can do the same for Send.
1 Like
Your Send and Sync impls are definitely unsound because they put no bounds on the inner type. You can use OnceTake as written in your example to send a !Send type like Rc between threads. I think the correct bounds would be
// SAFETY:
// Sending a OnceTake<T> can be used to send a T.
// Sending an owned OnceTake<T> cannot be used to send a &T. (requiring T: Sync)
unsafe impl<T: Send> Send for OnceTake<T> {}
// SAFETY:
// Being able to share a &OnceTake<T> allows T to be sent between threads via `take`.
// No Sync bound is necessary because &OnceTake<T> gives no access to &T.
unsafe impl<T: Send> Sync for OnceTake<T> {}
2 Likes
I think you are looking for a oneshot channel?
3 Likes
Yes, my bad! Thank you for the correction!