sync: add OnceCell::wait_initialized method
#7161
Conversation
Such a method was mentioned in tokio-rs#7151 as nice to have, so here it is. The previous `OnceCell` implementation was using a semaphore, but if it was convenient to synchronize initialization, it is not suited to wait for it. That's why the synchronization has completely been rewritten to handle it. There is now 3 possible states: *unset*, *locked*, and *set*. The cell is initialized to *unset*. When it is initialized, the cell is first locked, i.e. its state set to *locked*, then the value is written, and finally the state is set to *set*. As the other synchronization primitives, it uses a mutex-protected linked list of waiters, which is drained when initialization succeeds. If initialization fails, the state is rolled back to *unset* and one awaiting initializer is wakened up. The new implementation takes one less word in memory, and it has been optimized for ZST case. In fact, another goal of this new implementation is to serve as base for the `Event` primitive of tokio-rs#7151, so basically a `OnceCell<()>`. By the way, a proper tracing implementation has been added to `OnceCell`.
Co-authored-by: Alice Ryhl <aliceryhl@google.com>
|
cc8c284 fixes a bug of a waiter not being able to be enqueued a second time. There was no loom test catching the bug, but it found it while writing another one about happens-before between initialization attempts. |
tokio/src/sync/once_cell.rs
Outdated
| if node.waker.take().is_some() { | ||
| // SAFETY: The node is contained in the list. | ||
| unsafe { waiters.remove(NonNull::from(node)) }; | ||
| } |
There was a problem hiding this comment.
You don't need this condition, you can just call remove unconditionally. It's okay to call remove if the node is not in any list afair.
tokio/src/sync/once_cell.rs
Outdated
| } else { | ||
| // If there is no waker, it means the node is not queued, | ||
| // so it is pushed in the list. | ||
| if node.waker.replace(waker).is_none() { |
There was a problem hiding this comment.
Here (and above) you run the destructor of the waker with the lock held.
tokio/src/sync/once_cell.rs
Outdated
| /// Notify one initializer task. | ||
| fn notify_unlocked(&self) { | ||
| let mut waiters = self.waiters.lock(); | ||
| if let Some(mut waiter) = waiters.drain_filter(|w| w.is_initializer).next() { | ||
| // SAFETY: we hold the lock, so we can access the waker. | ||
| let waker = unsafe { waiter.as_mut().waker.take() }; | ||
| drop(waiters); | ||
| waker.unwrap().wake(); |
There was a problem hiding this comment.
Hmm .. this could be slow if there are many waiting for completion. It's O(n) every time we fail, which could become O(n^2) with many failures.
There was a problem hiding this comment.
Indeed, but it only happens when get_or[_try]_init and wait_initialized are mixed together. The current cell doesn't provide wait_initialized, so existing workflows using it should not be impacted.
I would expect wait_initialized to be more used in combination with set (as in the Event primitive from #7151), so this kind of workflow would also not be impacted.
To address the issue of the "mixed" workflow, we could consider the current solution as an acceptable tradeoff, at least to begin with, for the reasons above (no regression). Otherwise, i see several solutions:
- add a second linked list (behind the same lock), simple but would increase the size of the struct by 2 words, so one more than the current struct, which seems still acceptable
- use a single linked list, but keep an additional node reference, let's name it "middle node", which would be the first "initializer" node;
- initializer node are pushed to the front, setting the middle node if it is unset
wait_initializednode are pushed to the back- awaking a initializer removes the middle node from the list, the node on its left becoming the new middle node
- awaking all nodes just drain the list
this implementation is obviously more complex, but reduce the additional size to 1 word only (same as current)
- another slightly more complex implementation would be to use the state as tagged pointer to store the "middle node" above, so there would no need to increase the struct size. I admit I like this solution, but that will be for another time.
By the way, I've noticed that my use of drain_filter is bad, because it breaks the fairness of the initialization locking, as it starts from the front, while I'm using push_front (as there is no push_back). If we keep a single linked list, I will then write a new method to pop back with a predicate.
There was a problem hiding this comment.
Hmm. If we were okay with breaking the fairness, we could push initializers to the front and waiters to the back. This would mean that we could just pop the first element of the list to find an initializer. But we probably should not break fairness. Can you think of any similar tricks that don't break fairness?
Also, it sounds like we're missing a test for the fairness.
There was a problem hiding this comment.
Yes, we should preserve fairness. But I don't have any other tricks than the ones I described above, with a "middle node". Do you think I should try it? Or go with two lists? Or keep it like this for now, just fixing fairness?
There was a problem hiding this comment.
To be honest ... I think we should just introduce a new separate type rather than try to extend OnceCell. By having one type that only provides set and wait, and another type that only provides get_or[_try]_init, this problem just disappears.
Or does your use-case actually need a single type with both features?
Motivation
Such a method was mentioned in #7151 as nice to have, so here it is.
Solution
The previous
OnceCellimplementation was using a semaphore, but if it was convenient to synchronize initialization, it is not suited to wait for it. That's why the synchronization has completely been rewritten to handle it.There is now 3 possible states: unset, locked, and set. The cell is initialized to unset. When it is initialized, the cell is first locked, i.e. its state set to locked, then the value is written, and finally the state is set to set. As the other synchronization primitives, it uses a mutex-protected linked list of waiters, which is drained when initialization succeeds. If initialization fails, the state is rolled back to unset and one awaiting initializer is wakened up.
The new implementation takes one less word in memory, and it has been optimized for ZST case. In fact, another goal of this new implementation is to serve as base for the
Eventprimitive of #7151, so basically aOnceCell<()>.By the way, a proper tracing implementation has been added to
OnceCell.