From 192a45ccf11164d5c7409006b7b3aa366adb8d1c Mon Sep 17 00:00:00 2001 From: Janito Vaqueiro Ferreira Filho Date: Thu, 21 Oct 2021 08:47:04 -0300 Subject: [PATCH] Refactor rate limiting to not store `Sleep` type (#2915) In newer Tokio versions the `Sleep` type doesn't implement `Unpin`, so it's a little more complicated to use it. In this case it was easier to refactor the code to not store the `Sleep` type instead of wrapping it in a `Pin` type. --- zebra-network/src/peer_set/candidate_set.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/zebra-network/src/peer_set/candidate_set.rs b/zebra-network/src/peer_set/candidate_set.rs index b14e2b6e..9383f677 100644 --- a/zebra-network/src/peer_set/candidate_set.rs +++ b/zebra-network/src/peer_set/candidate_set.rs @@ -1,7 +1,7 @@ -use std::{cmp::min, mem, sync::Arc, time::Duration}; +use std::{cmp::min, sync::Arc}; use futures::stream::{FuturesUnordered, StreamExt}; -use tokio::time::{sleep, timeout, Instant, Sleep}; +use tokio::time::{sleep_until, timeout, Instant}; use tower::{Service, ServiceExt}; use zebra_chain::serialization::DateTime32; @@ -114,7 +114,7 @@ mod tests; pub(crate) struct CandidateSet { pub(super) address_book: Arc>, pub(super) peer_service: S, - wait_next_handshake: Sleep, + min_next_handshake: Instant, min_next_crawl: Instant, } @@ -131,7 +131,7 @@ where CandidateSet { address_book, peer_service, - wait_next_handshake: sleep(Duration::from_secs(0)), + min_next_handshake: Instant::now(), min_next_crawl: Instant::now(), } } @@ -321,9 +321,8 @@ where }; // SECURITY: rate-limit new outbound peer connections - (&mut self.wait_next_handshake).await; - let mut sleep = sleep(constants::MIN_PEER_CONNECTION_INTERVAL); - mem::swap(&mut self.wait_next_handshake, &mut sleep); + sleep_until(self.min_next_handshake).await; + self.min_next_handshake = Instant::now() + constants::MIN_PEER_CONNECTION_INTERVAL; Some(reconnect) }