diff --git a/libs/utils/src/leaky_bucket.rs b/libs/utils/src/leaky_bucket.rs index bcea5c84f4..a8e03e1566 100644 --- a/libs/utils/src/leaky_bucket.rs +++ b/libs/utils/src/leaky_bucket.rs @@ -142,14 +142,17 @@ struct Queue { state: Option, } -impl Queue { +impl RateLimiter { /// returns the sleep_counter start value on await. /// sleep_counter end value can be found within the enqueued. - fn wait(this: &Mutex) -> Enqueued<'_> { + fn wait(&self, count: usize) -> Enqueued<'_> { Enqueued { entry: pin_list::Node::new(), - queue: this, + limiter: self, + sleep_counter: 0, state: None, + count, + start: tokio::time::Instant::now(), } } } @@ -157,9 +160,10 @@ impl Queue { type RateLimitQueue = dyn pin_list::Types< Id = pin_list::id::Checked, // the waker that lets us wake the next in the queue - Protected = Waker, - // the token that gives us access to the rate limit state - Removed = LeakyBucketState, + Protected = (Waker, Instant, usize), + // the token that gives us access to the rate limit state. + // if None, then we were granted access by the leader + Removed = Option, // the sleep count at the start of the enqueue Unprotected = u64, >; @@ -168,8 +172,11 @@ pin_project_lite::pin_project! { struct Enqueued<'a> { #[pin] entry: Node, - state: Option<(LeakyBucketState, u64)>, - queue: &'a Mutex, + state: Option, + sleep_counter: u64, + limiter: &'a RateLimiter, + start: Instant, + count: usize, } impl<'a> PinnedDrop for Enqueued<'a> { @@ -177,31 +184,48 @@ pin_project_lite::pin_project! { let this = this.project(); #[allow(clippy::mut_mutex_lock, reason = "false positive")] - let mut q = this.queue.lock().unwrap(); + let mut q = this.limiter.queue.lock().unwrap(); - let state = if let Some(init) = this.entry.initialized_mut() { + let mut state = if let Some(init) = this.entry.initialized_mut() { let (data, _start_count) = init.reset(&mut q.queue) ; match data { // we were in the queue and are not holding any resources. - NodeData::Linked(_) => return, + NodeData::Linked(_) | NodeData::Removed(None) => return, // we were the head of the queue and were about to be the current leader - NodeData::Removed(state) => state + NodeData::Removed(Some(state)) => state } - } else if let Some((state, sleep_counter)) = this.state.take() { + } else if let Some(state) = this.state.take() { // we were holding the lock, and are now releasing it. - q.sleep_counter = sleep_counter; + q.sleep_counter = *this.sleep_counter; state } else { // we apparently didn't even get into the queue to begin with return; }; - // wake the next in the queue and give them the rate limit state - match q.queue.cursor_front_mut().remove_current(state) { - // wake up the next leader - Ok(waker) => waker.wake(), - // no tasks left in the queue. unlocked - Err(state) => q.state = Some(state), + let mut cursor = q.queue.cursor_front_mut(); + loop { + match cursor.protected() { + Some((_waker, start, count)) => { + match state.add_tokens(&this.limiter.config, *start, *count as f64) { + Ok(()) => { + cursor.remove_current(None) + .map_err(|_| {}).expect("we have just checked that the current node is in the list"); + }, + // next in the queue has to sleep + Err(_ready_at) => { + cursor.remove_current(Some(state)) + .map_err(|_| {}).expect("we have just checked that the current node is in the list"); + break; + } + } + }, + // no tasks left in the queue. unlocked + None => { + q.state = Some(state); + break; + } + } } } } @@ -215,7 +239,7 @@ impl std::future::Future for Enqueued<'_> { let mut node = this.entry; #[allow(clippy::mut_mutex_lock, reason = "false positive")] - let mut q = this.queue.lock().unwrap(); + let mut q = this.limiter.queue.lock().unwrap(); if let Some(init) = node.as_mut().initialized_mut() { // we are registered in the queue @@ -223,7 +247,8 @@ impl std::future::Future for Enqueued<'_> { match init.take_removed(&q.queue) { // if we are removed from the queue, that means we are the new leader Ok((state, start_count)) => { - *this.state = Some((state, q.sleep_counter)); + *this.state = state; + *this.sleep_counter = q.sleep_counter; Poll::Ready(start_count) } // if we are not removed from the queue, that means we are still waiting @@ -231,6 +256,7 @@ impl std::future::Future for Enqueued<'_> { Err(init) => { init.protected_mut(&mut q.queue) .unwrap() + .0 .clone_from(cx.waker()); Poll::Pending } @@ -241,11 +267,16 @@ impl std::future::Future for Enqueued<'_> { if let Some(state) = q.state.take() { // we are the first in the queue and it is not yet acquired - *this.state = Some((state, q.sleep_counter)); + *this.state = Some(state); + *this.sleep_counter = q.sleep_counter; Poll::Ready(start_count) } else { // we push ourselves to the back of the queue - q.queue.push_back(node, cx.waker().clone(), start_count); + q.queue.push_back( + node, + (cx.waker().clone(), *this.start, *this.count), + start_count, + ); Poll::Pending } } @@ -273,21 +304,21 @@ impl RateLimiter { /// returns true if we did throttle pub async fn acquire(&self, count: usize) -> bool { - let start = tokio::time::Instant::now(); - - let mut entry = std::pin::pin!(Queue::wait(&self.queue)); + let mut entry = std::pin::pin!(self.wait(count)); let start_count = entry.as_mut().await; - let (state, sleep_counter) = entry - .project() - .state - .as_mut() - .expect("token should be init if we returned from enqueued"); + let entry = entry.project(); + + let Some(state) = entry.state.as_mut() else { + // we were woken up without the state, + // thus the state leader must have allowed us to continue + return start_count < *entry.sleep_counter; + }; loop { - match state.add_tokens(&self.config, start, count as f64) { - Ok(()) => return start_count < *sleep_counter, + match state.add_tokens(&self.config, *entry.start, count as f64) { + Ok(()) => return start_count < *entry.sleep_counter, Err(ready_at) => { - *sleep_counter += 1; + *entry.sleep_counter += 1; tokio::time::sleep_until(ready_at).await; } }