fix: restrict max active queries (#4178)

This commit is contained in:
Matthias Seitz
2023-08-14 11:45:28 +02:00
committed by GitHub
parent 4ff8bca977
commit 72236490ce
2 changed files with 7 additions and 2 deletions

View File

@@ -27,6 +27,11 @@ impl RateLimit {
RateLimit { rate, state, sleep: Box::pin(tokio::time::sleep_until(until)) }
}
/// Returns the configured limit of the [RateLimit]
pub fn limit(&self) -> u64 {
self.rate.limit()
}
/// Checks if the [RateLimit] is ready to handle a new call
pub fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<()> {
match self.state {

View File

@@ -78,8 +78,8 @@ impl<R: Resolver, K: EnrKeyUnambiguous> QueryPool<R, K> {
return Poll::Ready(event)
}
// queue in new queries
'queries: loop {
// queue in new queries if we have capacity
'queries: while self.active_queries.len() < self.rate_limit.limit() as usize {
if self.rate_limit.poll_ready(cx).is_ready() {
if let Some(query) = self.queued_queries.pop_front() {
self.rate_limit.tick();