From cb57906f204d3f46925a7c69b2e298d6b5e26d44 Mon Sep 17 00:00:00 2001
From: Alexandru Gheorghe <alexandru.gheorghe@parity.io>
Date: Thu, 20 Jun 2024 14:38:53 +0300
Subject: [PATCH] Optimize logic for gossiping assignments

Signed-off-by: Alexandru Gheorghe <alexandru.gheorghe@parity.io>
---
 .../network/approval-distribution/src/lib.rs  | 30 ++++++++++++-------
 .../network/protocol/src/grid_topology.rs     | 24 +++++++++++++++
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/polkadot/node/network/approval-distribution/src/lib.rs b/polkadot/node/network/approval-distribution/src/lib.rs
index 56b993a60f17..63c1af6cad6d 100644
--- a/polkadot/node/network/approval-distribution/src/lib.rs
+++ b/polkadot/node/network/approval-distribution/src/lib.rs
@@ -1510,6 +1510,21 @@ impl State {
 		let required_routing = topology.map_or(RequiredRouting::PendingTopology, |t| {
 			t.local_grid_neighbors().required_routing_by_index(validator_index, local)
 		});
+		// Peers that we will send the assignment to.
+		let mut peers = Vec::new();
+
+		let peers_to_route_to = topology
+			.as_ref()
+			.map(|t| t.peers_to_route(required_routing))
+			.unwrap_or_default();
+
+		for peer in peers_to_route_to {
+			if entry.known_by.get(&peer).is_none() {
+				continue
+			}
+
+			peers.push(peer);
+		}
 
 		// All the peers that know the relay chain block.
 		let peers_to_filter = entry.known_by();
@@ -1535,23 +1550,12 @@ impl State {
 		let n_peers_total = self.peer_views.len();
 		let source_peer = source.peer_id();
 
-		// Peers that we will send the assignment to.
-		let mut peers = Vec::new();
-
 		// Filter destination peers
 		for peer in peers_to_filter.into_iter() {
 			if Some(peer) == source_peer {
 				continue
 			}
 
-			if let Some(true) = topology
-				.as_ref()
-				.map(|t| t.local_grid_neighbors().route_to_peer(required_routing, &peer))
-			{
-				peers.push(peer);
-				continue
-			}
-
 			if !topology.map(|topology| topology.is_validator(&peer)).unwrap_or(false) {
 				continue
 			}
@@ -1566,6 +1570,10 @@ impl State {
 				approval_entry.routing_info_mut().mark_randomly_sent(peer);
 				peers.push(peer);
 			}
+
+			if approval_entry.routing_info().random_routing.is_complete() {
+				break
+			}
 		}
 
 		// Add the metadata of the assignment to the knowledge of each peer.
diff --git a/polkadot/node/network/protocol/src/grid_topology.rs b/polkadot/node/network/protocol/src/grid_topology.rs
index a14d24610722..bed10280ffa9 100644
--- a/polkadot/node/network/protocol/src/grid_topology.rs
+++ b/polkadot/node/network/protocol/src/grid_topology.rs
@@ -278,6 +278,18 @@ impl GridNeighbors {
 			.collect::<Vec<_>>()
 	}
 
+	/// Returns the the Peer Ids for the required routing.
+	pub fn peers_to_route(&self, required_routing: RequiredRouting) -> Vec<PeerId> {
+		match required_routing {
+			RequiredRouting::All =>
+				self.peers_x.iter().chain(self.peers_y.iter()).copied().collect(),
+			RequiredRouting::GridX => self.peers_x.iter().copied().collect(),
+			RequiredRouting::GridY => self.peers_y.iter().copied().collect(),
+			RequiredRouting::GridXY =>
+				self.peers_x.iter().chain(self.peers_y.iter()).copied().collect(),
+			RequiredRouting::None | RequiredRouting::PendingTopology => Vec::new(),
+		}
+	}
 	/// A convenience method that returns total number of peers in the topology
 	pub fn len(&self) -> usize {
 		self.peers_x.len().saturating_add(self.peers_y.len())
@@ -313,6 +325,14 @@ impl SessionGridTopologyEntry {
 		self.topology.is_validator(peer)
 	}
 
+	/// Returns peers to route based on the required routing
+	pub fn peers_to_route(&self, required_routing: RequiredRouting) -> Vec<PeerId> {
+		if required_routing == RequiredRouting::All {
+			return self.topology.peer_ids.iter().copied().collect()
+		}
+		self.local_neighbors.peers_to_route(required_routing)
+	}
+
 	/// Updates the known peer ids for the passed authorities ids.
 	pub fn update_authority_ids(
 		&mut self,
@@ -524,6 +544,10 @@ impl RandomRouting {
 	pub fn inc_sent(&mut self) {
 		self.sent += 1
 	}
+
+	pub fn is_complete(&self) -> bool {
+		self.sent >= self.target
+	}
 }
 
 /// Routing mode