Statement Distribution Per Peer Rate Limit

polkadot/node/network/statement-distribution/src/v2/mod.rs

@@ -60,6 +60,8 @@ use futures::{
 	channel::{mpsc, oneshot},
 	stream::FuturesUnordered,
 	SinkExt, StreamExt,
+	select,
+	future::FutureExt,
 };
 
 use std::{
@@ -3335,33 +3337,55 @@ pub(crate) async fn respond_task(
 	mut sender: mpsc::Sender<ResponderMessage>,
 ) {
 	let mut pending_out = FuturesUnordered::new();
+	let mut active_peers = HashSet::new();
+
 	loop {
-		// Ensure we are not handling too many requests in parallel.
-		if pending_out.len() >= MAX_PARALLEL_ATTESTED_CANDIDATE_REQUESTS as usize {
-			// Wait for one to finish:
-			pending_out.next().await;
-		}
+		select! {
+			// New request
+			request_result = receiver.recv(|| vec![COST_INVALID_REQUEST]).fuse() => {
+				let request = match request_result.into_nested() {
+					Ok(Ok(v)) => v,
+					Err(fatal) => {
+						gum::debug!(target: LOG_TARGET, error = ?fatal, "Shutting down request responder");
+						return
+					},
+					Ok(Err(jfyi)) => {
+						gum::debug!(target: LOG_TARGET, error = ?jfyi, "Decoding request failed");
+						continue
+					},
+				};
 
-		let req = match receiver.recv(|| vec![COST_INVALID_REQUEST]).await.into_nested() {
-			Ok(Ok(v)) => v,
-			Err(fatal) => {
-				gum::debug!(target: LOG_TARGET, error = ?fatal, "Shutting down request responder");
-				return
+				// If peer currently being served drop request
+				if active_peers.contains(&request.peer) {
+					continue
+				}
+
+				// If we are over parallel limit wait for one to finish
+				if pending_out.len() >= MAX_PARALLEL_ATTESTED_CANDIDATE_REQUESTS as usize {
+					let result = pending_out.select_next_some().await;
+					let (_, peer) = result;
+					active_peers.remove(&peer);
+				}
+
+				// Start serving the request
+				let (pending_sent_tx, pending_sent_rx) = oneshot::channel();
+				let peer = request.peer.clone();
+				if let Err(err) = sender
+					.feed(ResponderMessage { request: request, sent_feedback: pending_sent_tx })
+					.await
+				{
+					gum::debug!(target: LOG_TARGET, ?err, "Shutting down responder");
+					return
+				}
+				let future_with_peer = pending_sent_rx.map(move |result| (result, peer));
+				pending_out.push(future_with_peer);
+				active_peers.insert(peer);
 			},
-			Ok(Err(jfyi)) => {
-				gum::debug!(target: LOG_TARGET, error = ?jfyi, "Decoding request failed");
-				continue
+			// Request served/finished
+			result = pending_out.select_next_some() => {
+				let (_, peer) = result;
+				active_peers.remove(&peer);
 			},
-		};
-
-		let (pending_sent_tx, pending_sent_rx) = oneshot::channel();
-		if let Err(err) = sender
-			.feed(ResponderMessage { request: req, sent_feedback: pending_sent_tx })
-			.await
-		{
-			gum::debug!(target: LOG_TARGET, ?err, "Shutting down responder");
-			return
 		}
-		pending_out.push(pending_sent_rx);
 	}
-}
+}

polkadot/node/network/statement-distribution/src/v2/requests.rs

@@ -357,6 +357,7 @@ impl RequestManager {
 				id,
 				&props,
 				&peer_advertised,
+				&response_manager,
 			) {
 				None => continue,
 				Some(t) => t,
@@ -378,14 +379,17 @@ impl RequestManager {
 			);
 
 			let stored_id = id.clone();
-			response_manager.push(Box::pin(async move {
-				TaggedResponse {
-					identifier: stored_id,
-					requested_peer: target,
-					props,
-					response: response_fut.await,
-				}
-			}));
+			response_manager.push(
+				Box::pin(async move {
+					TaggedResponse {
+						identifier: stored_id,
+						requested_peer: target,
+						props,
+						response: response_fut.await,
+					}
+				}),
+				target,
+			);
 
 			entry.in_flight = true;
 
@@ -413,28 +417,40 @@ impl RequestManager {
 /// A manager for pending responses.
 pub struct ResponseManager {
 	pending_responses: FuturesUnordered<BoxFuture<'static, TaggedResponse>>,
+	active_peers: HashSet<PeerId>,
 }
 
 impl ResponseManager {
 	pub fn new() -> Self {
-		Self { pending_responses: FuturesUnordered::new() }
+		Self { 
+			pending_responses: FuturesUnordered::new(),
+			active_peers: HashSet::new(),
+		}
 	}
 
 	/// Await the next incoming response to a sent request, or immediately
 	/// return `None` if there are no pending responses.
 	pub async fn incoming(&mut self) -> Option<UnhandledResponse> {
-		self.pending_responses
-			.next()
-			.await
-			.map(|response| UnhandledResponse { response })
+		if let Some(response) = self.pending_responses.next().await {
+			// Upon receiving a response, remove the peer from the active set.
+			self.active_peers.remove(&response.requested_peer);
+			return Some(UnhandledResponse { response });
+		}
+		None
 	}
 
 	fn len(&self) -> usize {
 		self.pending_responses.len()
 	}
 
-	fn push(&mut self, response: BoxFuture<'static, TaggedResponse>) {
+	fn push(&mut self, response: BoxFuture<'static, TaggedResponse>, target: PeerId) {
 		self.pending_responses.push(response);
+		self.active_peers.insert(target);
+	}
+
+	/// Returns true if we are currently sending a request to the peer.
+	fn is_sending_to(&self, peer: &PeerId) -> bool {
+		self.active_peers.contains(peer)
 	}
 }
 
@@ -462,10 +478,17 @@ fn find_request_target_with_update(
 	candidate_identifier: &CandidateIdentifier,
 	props: &RequestProperties,
 	peer_advertised: impl Fn(&CandidateIdentifier, &PeerId) -> Option<StatementFilter>,
+	response_manager: &ResponseManager,
 ) -> Option<PeerId> {
 	let mut prune = Vec::new();
 	let mut target = None;
 	for (i, p) in known_by.iter().enumerate() {
+
+		// If we are already sending to that peer, skip for now
+		if response_manager.is_sending_to(p) {
+			continue
+		}
+
 		let mut filter = match peer_advertised(candidate_identifier, p) {
 			None => {
 				prune.push(i);

polkadot/node/network/statement-distribution/src/v2/tests/requests.rs

@@ -1891,7 +1891,7 @@ fn local_node_sanity_checks_incoming_requests() {
 			let mask = StatementFilter::blank(state.config.group_size + 1);
 			let response = state
 				.send_request(
-					peer_c,
+					peer_a,
 					request_v2::AttestedCandidateRequest { candidate_hash: candidate.hash(), mask },
 				)
 				.await