Fix pool claim race condition - #10301
Conversation
3da18aa to
4f290aa
Compare
Rich-T-kid
left a comment
Rich-T-kid
left a comment
There was a problem hiding this comment.
I haven't taken a look at the code yet but would it be possible to solve the race condition by replacing
pub fn claim(&self, pool: &dyn MemoryPool) {
*self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity()));
}with just
pub fn claim(&self, pool: &dyn MemoryPool) {
let guard = self.reservation.lock().unwrap();
drop(*guard.take());
*guard = pool.reserve(self.capacity());
}like you mentioned in the issue. Why does TrackedReservation need to be introduced?
|
Yes that is essentially what this PR does. But instead of duplicating that exact same bit of code everywhere, we lift it into a new struct that does it correctly every time & makes it easy for other things to get it right. |
@cetra3 okay nice. Ill try and take a look tomorrow |
Rich-T-kid
left a comment
Rich-T-kid
left a comment
There was a problem hiding this comment.
LGTM
left a single nit
| let mut guard = self.reservation.lock().unwrap(); | ||
| if let Some(mut reservation) = guard.take() { | ||
| // Resize the reservation | ||
| reservation.resize(new_size); | ||
|
|
||
| // Put it back | ||
| *guard = Some(reservation); | ||
| } |
| //! (pool tracker) (resizable) | ||
| //! (pool tracker) (resizable) | ||
| //! ┌──────────────────┐ fn reserve() ┌─────────────────────────┐ | ||
| //! │ trait MemoryPool │─────────────►│ trait MemoryReservation │ |
There was a problem hiding this comment.
Nit: there are a couple of spots in this PR where nothing has actually changed but a diff is still being shown. We should generally try to avoid that.
| } | ||
|
|
||
| #[test] | ||
| fn test_claim_reclaims_before_reserving() { |
this PR makes those call sites much more readable |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thanks @cetra3 and @Rich-T-kid
@waynexia can you please help review this PR as you contributed the initial memory pool infrstructure in #7303 ?
| /// This is a wrapper for the reservation so we can standardize on changing | ||
| /// and avoid race conditions in memory accounting | ||
| #[derive(Debug, Default)] | ||
| pub struct TrackedReservation { |
There was a problem hiding this comment.
Is this pub in the sense that it can be used outside this crate? If so I think we should document it more, and be clear if we need to expand the public API -- if it is meant to be internal, perhaps we can mark it as pub(crate)
There was a problem hiding this comment.
I've moved this to pub(crate). If people find it useful we can adjust but it is really an internal only thing
3 participants
Which issue does this PR close?
Rationale for this change
This adjusts the internals of the memory pool reservation so that there isn't any double accounting when replacing memory with
claimWhat changes are included in this PR?
Adds a new struct in
pool.rswhich is used to track reservations on buffers. This is essentially a newtype lift of the existing structure, but with a few specialised methods to ensure that claim accounts correctly.Are these changes tested?
Yes a new test has been added
Are there any user-facing changes?
No this is just some memory pool accounting internals