refactor grid topology to expose more info to subsystems (#6140)

* refactor grid topology to expose more info to subsystems

* fix grid_topology test

* fix overseer test

* Update node/network/protocol/src/grid_topology.rs

Co-authored-by: Vsevolod Stakhov <vsevolod.stakhov@parity.io>

* Update node/network/protocol/src/grid_topology.rs

Co-authored-by: Andronik <write@reusable.software>

* Update node/network/protocol/src/grid_topology.rs

Co-authored-by: Andronik <write@reusable.software>

* fix bug in populating topology

* fmt

Co-authored-by: Vsevolod Stakhov <vsevolod.stakhov@parity.io>
Co-authored-by: Andronik <write@reusable.software>

polkadot/node/network/approval-distribution/src/lib.rs

@@ -343,9 +343,13 @@ impl State {
 				})
 			},
 			NetworkBridgeEvent::NewGossipTopology(topology) => {
-				let session = topology.session;
-				self.handle_new_session_topology(ctx, session, SessionGridTopology::from(topology))
-					.await;
+				self.handle_new_session_topology(
+					ctx,
+					topology.session,
+					topology.topology,
+					topology.local_index,
+				)
+				.await;
 			},
 			NetworkBridgeEvent::PeerViewChange(peer_id, view) => {
 				self.handle_peer_view_change(ctx, metrics, peer_id, view, rng).await;
@@ -500,8 +504,14 @@ impl State {
 		ctx: &mut Context,
 		session: SessionIndex,
 		topology: SessionGridTopology,
+		local_index: Option<ValidatorIndex>,
 	) {
-		self.topologies.insert_topology(session, topology);
+		if local_index.is_none() {
+			// this subsystem only matters to validators.
+			return
+		}
+
+		self.topologies.insert_topology(session, topology, local_index);
 		let topology = self.topologies.get_topology(session).expect("just inserted above; qed");
 
 		adjust_required_routing_and_propagate(
@@ -511,7 +521,9 @@ impl State {
 			|block_entry| block_entry.session == session,
 			|required_routing, local, validator_index| {
 				if *required_routing == RequiredRouting::PendingTopology {
-					*required_routing = topology.required_routing_by_index(*validator_index, local);
+					*required_routing = topology
+						.local_grid_neighbors()
+						.required_routing_by_index(*validator_index, local);
 				}
 			},
 		)
@@ -861,7 +873,7 @@ impl State {
 		let local = source == MessageSource::Local;
 
 		let required_routing = topology.map_or(RequiredRouting::PendingTopology, |t| {
-			t.required_routing_by_index(validator_index, local)
+			t.local_grid_neighbors().required_routing_by_index(validator_index, local)
 		});
 
 		let message_state = match entry.candidates.get_mut(claimed_candidate_index as usize) {
@@ -902,7 +914,10 @@ impl State {
 				return false
 			}
 
-			if let Some(true) = topology.as_ref().map(|t| t.route_to_peer(required_routing, peer)) {
+			if let Some(true) = topology
+				.as_ref()
+				.map(|t| t.local_grid_neighbors().route_to_peer(required_routing, peer))
+			{
 				return true
 			}
 
@@ -1169,7 +1184,8 @@ impl State {
 			//      the assignment to all aware peers in the required routing _except_ the original
 			//      source of the assignment. Hence the `in_topology_check`.
 			//   3. Any randomly selected peers have been sent the assignment already.
-			let in_topology = topology.map_or(false, |t| t.route_to_peer(required_routing, peer));
+			let in_topology = topology
+				.map_or(false, |t| t.local_grid_neighbors().route_to_peer(required_routing, peer));
 			in_topology || knowledge.sent.contains(message_subject, MessageKind::Assignment)
 		};
 
@@ -1301,9 +1317,9 @@ impl State {
 						let required_routing = message_state.required_routing;
 						let rng = &mut *rng;
 						let mut peer_filter = move |peer_id| {
-							let in_topology = topology
-								.as_ref()
-								.map_or(false, |t| t.route_to_peer(required_routing, peer_id));
+							let in_topology = topology.as_ref().map_or(false, |t| {
+								t.local_grid_neighbors().route_to_peer(required_routing, peer_id)
+							});
 							in_topology || {
 								let route_random = random_routing.sample(total_peers, rng);
 								if route_random {
@@ -1564,7 +1580,10 @@ async fn adjust_required_routing_and_propagate<Context, BlockFilter, RoutingModi
 				});
 
 			for (peer, peer_knowledge) in &mut block_entry.known_by {
-				if !topology.route_to_peer(message_state.required_routing, peer) {
+				if !topology
+					.local_grid_neighbors()
+					.route_to_peer(message_state.required_routing, peer)
+				{
 					continue
 				}
 

polkadot/node/network/approval-distribution/src/tests.rs

@@ -17,7 +17,12 @@
 use super::*;
 use assert_matches::assert_matches;
 use futures::{executor, future, Future};
-use polkadot_node_network_protocol::{our_view, peer_set::ValidationVersion, view, ObservedRole};
+use polkadot_node_network_protocol::{
+	grid_topology::{SessionGridTopology, TopologyPeerInfo},
+	our_view,
+	peer_set::ValidationVersion,
+	view, ObservedRole,
+};
 use polkadot_node_primitives::approval::{
 	AssignmentCertKind, VRFOutput, VRFProof, RELAY_VRF_MODULO_CONTEXT,
 };
@@ -119,33 +124,79 @@ fn make_gossip_topology(
 	neighbors_x: &[usize],
 	neighbors_y: &[usize],
 ) -> network_bridge_event::NewGossipTopology {
-	let mut t = network_bridge_event::NewGossipTopology {
-		session,
-		our_neighbors_x: HashMap::new(),
-		our_neighbors_y: HashMap::new(),
+	// This builds a grid topology which is a square matrix.
+	// The local validator occupies the top left-hand corner.
+	// The X peers occupy the same row and the Y peers occupy
+	// the same column.
+
+	let local_index = 1;
+
+	assert_eq!(
+		neighbors_x.len(),
+		neighbors_y.len(),
+		"mocking grid topology only implemented for squares",
+	);
+
+	let d = neighbors_x.len() + 1;
+
+	let grid_size = d * d;
+	assert!(grid_size > 0);
+	assert!(all_peers.len() >= grid_size);
+
+	let peer_info = |i: usize| TopologyPeerInfo {
+		peer_ids: vec![all_peers[i].0.clone()],
+		validator_index: ValidatorIndex::from(i as u32),
+		discovery_id: all_peers[i].1.clone(),
 	};
 
-	for &i in neighbors_x {
-		t.our_neighbors_x.insert(
-			all_peers[i].1.clone(),
-			network_bridge_event::TopologyPeerInfo {
-				peer_ids: vec![all_peers[i].0.clone()],
-				validator_index: ValidatorIndex::from(i as u32),
-			},
-		);
+	let mut canonical_shuffling: Vec<_> = (0..)
+		.filter(|i| local_index != *i)
+		.filter(|i| !neighbors_x.contains(i))
+		.filter(|i| !neighbors_y.contains(i))
+		.take(grid_size)
+		.map(peer_info)
+		.collect();
+
+	// filled with junk except for own.
+	let mut shuffled_indices = vec![d + 1; grid_size];
+	shuffled_indices[local_index] = 0;
+	canonical_shuffling[0] = peer_info(local_index);
+
+	for (x_pos, v) in neighbors_x.iter().enumerate() {
+		let pos = 1 + x_pos;
+		canonical_shuffling[pos] = peer_info(*v);
 	}
 
-	for &i in neighbors_y {
-		t.our_neighbors_y.insert(
-			all_peers[i].1.clone(),
-			network_bridge_event::TopologyPeerInfo {
-				peer_ids: vec![all_peers[i].0.clone()],
-				validator_index: ValidatorIndex::from(i as u32),
-			},
-		);
+	for (y_pos, v) in neighbors_y.iter().enumerate() {
+		let pos = d * (1 + y_pos);
+		canonical_shuffling[pos] = peer_info(*v);
 	}
 
-	t
+	let topology = SessionGridTopology::new(shuffled_indices, canonical_shuffling);
+
+	// sanity check.
+	{
+		let g_n = topology
+			.compute_grid_neighbors_for(ValidatorIndex(local_index as _))
+			.expect("topology just constructed with this validator index");
+
+		assert_eq!(g_n.validator_indices_x.len(), neighbors_x.len());
+		assert_eq!(g_n.validator_indices_y.len(), neighbors_y.len());
+
+		for i in neighbors_x {
+			assert!(g_n.validator_indices_x.contains(&ValidatorIndex(*i as _)));
+		}
+
+		for i in neighbors_y {
+			assert!(g_n.validator_indices_y.contains(&ValidatorIndex(*i as _)));
+		}
+	}
+
+	network_bridge_event::NewGossipTopology {
+		session,
+		topology,
+		local_index: Some(ValidatorIndex(local_index as _)),
+	}
 }
 
 async fn setup_gossip_topology(