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!