Claim API Race Condition when reclaiming existing memory #10139

Describe the bug

There appears to be a section of time that, if you are using atomics for reservations, that the claim feature can double account how much memory is "claimed" at certain time periods when reclaiming memory.

When we've been testing the new claim pool feature, we've seen the total claimed memory sometimes double that of what is actually in use.

To Reproduce

This is a race condition in the way that the claim pool currently works and is a little subtle:

pub fn claim(&self, pool: &dyn MemoryPool) {
    *self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity()));
}

If you break this down about when the reservation is live, and when the old one is dropped, it is roughly this equivalent sequence:

pub fn claim(&self, pool: &dyn MemoryPool) {
    let guard = self.reservation.lock().unwrap(); // old claim still exists here
    let new_claim = pool.reserve(self.capacity()); // BOTH old claim + new claim exist
    *guard = Some(new_claim); // old claim is dropped
}

Expected behavior

I think that the expected behaviour is that the existing claim is dropped before the new one is requested.

I.e, I think this is what I'd expect to fix it:

pub fn claim(&self, pool: &dyn MemoryPool) {
    let guard = self.reservation.lock().unwrap();
    drop(*guard.take());
    *guard = pool.reserve(self.capacity());
}

Additional context

No response