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>
2026-04-26 15:47:58 +00:00 · 2022-10-12 18:30:12 -05:00
parent bccffcad12
commit a7780e0797
17 changed files with 614 additions and 328 deletions
@@ -30,7 +30,7 @@
 //!

 use crate::PeerId;
-use polkadot_primitives::v2::{SessionIndex, ValidatorIndex};
+use polkadot_primitives::v2::{AuthorityDiscoveryId, SessionIndex, ValidatorIndex};
 use rand::{CryptoRng, Rng};
 use std::{
 	collections::{hash_map, HashMap, HashSet},
@@ -48,9 +48,106 @@ pub const DEFAULT_RANDOM_SAMPLE_RATE: usize = crate::MIN_GOSSIP_PEERS;
 /// The number of peers to randomly propagate messages to.
 pub const DEFAULT_RANDOM_CIRCULATION: usize = 4;

-/// Topology representation
-#[derive(Default, Clone, Debug)]
+/// Information about a peer in the gossip topology for a session.
+#[derive(Debug, Clone, PartialEq)]
+pub struct TopologyPeerInfo {
+	/// The validator's known peer IDs.
+	pub peer_ids: Vec<PeerId>,
+	/// The index of the validator in the discovery keys of the corresponding
+	/// `SessionInfo`. This can extend _beyond_ the set of active parachain validators.
+	pub validator_index: ValidatorIndex,
+	/// The authority discovery public key of the validator in the corresponding
+	/// `SessionInfo`.
+	pub discovery_id: AuthorityDiscoveryId,
+}
+
+/// Topology representation for a session.
+#[derive(Default, Clone, Debug, PartialEq)]
 pub struct SessionGridTopology {
+	/// An array mapping validator indices to their indices in the
+	/// shuffling itself. This has the same size as the number of validators
+	/// in the session.
+	shuffled_indices: Vec<usize>,
+	/// The canonical shuffling of validators for the session.
+	canonical_shuffling: Vec<TopologyPeerInfo>,
+}
+
+impl SessionGridTopology {
+	/// Create a new session grid topology.
+	pub fn new(shuffled_indices: Vec<usize>, canonical_shuffling: Vec<TopologyPeerInfo>) -> Self {
+		SessionGridTopology { shuffled_indices, canonical_shuffling }
+	}
+
+	/// Produces the outgoing routing logic for a particular peer.
+	///
+	/// Returns `None` if the validator index is out of bounds.
+	pub fn compute_grid_neighbors_for(&self, v: ValidatorIndex) -> Option<GridNeighbors> {
+		if self.shuffled_indices.len() != self.canonical_shuffling.len() {
+			return None
+		}
+		let shuffled_val_index = *self.shuffled_indices.get(v.0 as usize)?;
+
+		let neighbors = matrix_neighbors(shuffled_val_index, self.shuffled_indices.len())?;
+
+		let mut grid_subset = GridNeighbors::empty();
+		for r_n in neighbors.row_neighbors {
+			let n = &self.canonical_shuffling[r_n];
+			grid_subset.validator_indices_x.insert(n.validator_index);
+			for p in &n.peer_ids {
+				grid_subset.peers_x.insert(p.clone());
+			}
+		}
+
+		for c_n in neighbors.column_neighbors {
+			let n = &self.canonical_shuffling[c_n];
+			grid_subset.validator_indices_y.insert(n.validator_index);
+			for p in &n.peer_ids {
+				grid_subset.peers_y.insert(p.clone());
+			}
+		}
+
+		Some(grid_subset)
+	}
+}
+
+struct MatrixNeighbors<R, C> {
+	row_neighbors: R,
+	column_neighbors: C,
+}
+
+/// Compute the row and column neighbors of `val_index` in a matrix
+fn matrix_neighbors(
+	val_index: usize,
+	len: usize,
+) -> Option<MatrixNeighbors<impl Iterator<Item = usize>, impl Iterator<Item = usize>>> {
+	if val_index >= len {
+		return None
+	}
+
+	// e.g. for size 11 the matrix would be
+	//
+	// 0  1  2
+	// 3  4  5
+	// 6  7  8
+	// 9 10
+	//
+	// and for index 10, the neighbors would be 1, 4, 7, 9
+
+	let sqrt = (len as f64).sqrt() as usize;
+	let our_row = val_index / sqrt;
+	let our_column = val_index % sqrt;
+	let row_neighbors = our_row * sqrt..std::cmp::min(our_row * sqrt + sqrt, len);
+	let column_neighbors = (our_column..len).step_by(sqrt);
+
+	Some(MatrixNeighbors {
+		row_neighbors: row_neighbors.filter(move |i| *i != val_index),
+		column_neighbors: column_neighbors.filter(move |i| *i != val_index),
+	})
+}
+
+/// Information about the grid neighbors for a particular node in the topology.
+#[derive(Debug, Clone, PartialEq)]
+pub struct GridNeighbors {
 	/// Represent peers in the X axis
 	pub peers_x: HashSet<PeerId>,
 	/// Represent validators in the X axis
@@ -61,7 +158,18 @@ pub struct SessionGridTopology {
 	pub validator_indices_y: HashSet<ValidatorIndex>,
 }

-impl SessionGridTopology {
+impl GridNeighbors {
+	/// Utility function for creating an empty set of grid neighbors.
+	/// Useful for testing.
+	pub fn empty() -> Self {
+		GridNeighbors {
+			peers_x: HashSet::new(),
+			validator_indices_x: HashSet::new(),
+			peers_y: HashSet::new(),
+			validator_indices_y: HashSet::new(),
+		}
+	}
+
 	/// Given the originator of a message as a validator index, indicates the part of the topology
 	/// we're meant to send the message to.
 	pub fn required_routing_by_index(
@@ -123,7 +231,7 @@ impl SessionGridTopology {
 	}

 	/// Returns the difference between this and the `other` topology as a vector of peers
-	pub fn peers_diff(&self, other: &SessionGridTopology) -> Vec<PeerId> {
+	pub fn peers_diff(&self, other: &Self) -> Vec<PeerId> {
 		self.peers_x
 			.iter()
 			.chain(self.peers_y.iter())
@@ -138,15 +246,39 @@ impl SessionGridTopology {
 	}
 }

+/// An entry tracking a session grid topology and some cached local neighbors.
+#[derive(Debug)]
+pub struct SessionGridTopologyEntry {
+	topology: SessionGridTopology,
+	local_neighbors: GridNeighbors,
+}
+
+impl SessionGridTopologyEntry {
+	/// Access the local grid neighbors.
+	pub fn local_grid_neighbors(&self) -> &GridNeighbors {
+		&self.local_neighbors
+	}
+
+	/// Access the local grid neighbors mutably.
+	pub fn local_grid_neighbors_mut(&mut self) -> &mut GridNeighbors {
+		&mut self.local_neighbors
+	}
+
+	/// Access the underlying topology.
+	pub fn get(&self) -> &SessionGridTopology {
+		&self.topology
+	}
+}
+
 /// A set of topologies indexed by session
 #[derive(Default)]
 pub struct SessionGridTopologies {
-	inner: HashMap<SessionIndex, (Option<SessionGridTopology>, usize)>,
+	inner: HashMap<SessionIndex, (Option<SessionGridTopologyEntry>, usize)>,
 }

 impl SessionGridTopologies {
 	/// Returns a topology for the specific session index
-	pub fn get_topology(&self, session: SessionIndex) -> Option<&SessionGridTopology> {
+	pub fn get_topology(&self, session: SessionIndex) -> Option<&SessionGridTopologyEntry> {
 		self.inner.get(&session).and_then(|val| val.0.as_ref())
 	}

@@ -166,63 +298,112 @@ impl SessionGridTopologies {
 	}

 	/// Insert a new topology, no-op if already present.
-	pub fn insert_topology(&mut self, session: SessionIndex, topology: SessionGridTopology) {
+	pub fn insert_topology(
+		&mut self,
+		session: SessionIndex,
+		topology: SessionGridTopology,
+		local_index: Option<ValidatorIndex>,
+	) {
 		let entry = self.inner.entry(session).or_insert((None, 0));
 		if entry.0.is_none() {
-			entry.0 = Some(topology);
+			let local_neighbors = local_index
+				.and_then(|l| topology.compute_grid_neighbors_for(l))
+				.unwrap_or_else(GridNeighbors::empty);
+
+			entry.0 = Some(SessionGridTopologyEntry { topology, local_neighbors });
 		}
 	}
 }

 /// A simple storage for a topology and the corresponding session index
-#[derive(Default, Debug)]
-pub struct GridTopologySessionBound {
-	topology: SessionGridTopology,
+#[derive(Debug)]
+struct GridTopologySessionBound {
+	entry: SessionGridTopologyEntry,
 	session_index: SessionIndex,
 }

 /// A storage for the current and maybe previous topology
-#[derive(Default, Debug)]
+#[derive(Debug)]
 pub struct SessionBoundGridTopologyStorage {
 	current_topology: GridTopologySessionBound,
 	prev_topology: Option<GridTopologySessionBound>,
 }

+impl Default for SessionBoundGridTopologyStorage {
+	fn default() -> Self {
+		// having this struct be `Default` is objectively stupid
+		// but used in a few places
+		SessionBoundGridTopologyStorage {
+			current_topology: GridTopologySessionBound {
+				// session 0 is valid so we should use the upper bound
+				// as the default instead of the lower bound.
+				session_index: SessionIndex::max_value(),
+				entry: SessionGridTopologyEntry {
+					topology: SessionGridTopology {
+						shuffled_indices: Vec::new(),
+						canonical_shuffling: Vec::new(),
+					},
+					local_neighbors: GridNeighbors::empty(),
+				},
+			},
+			prev_topology: None,
+		}
+	}
+}
+
 impl SessionBoundGridTopologyStorage {
 	/// Return a grid topology based on the session index:
 	/// If we need a previous session and it is registered in the storage, then return that session.
 	/// Otherwise, return a current session to have some grid topology in any case
-	pub fn get_topology_or_fallback(&self, idx: SessionIndex) -> &SessionGridTopology {
-		self.get_topology(idx).unwrap_or(&self.current_topology.topology)
+	pub fn get_topology_or_fallback(&self, idx: SessionIndex) -> &SessionGridTopologyEntry {
+		self.get_topology(idx).unwrap_or(&self.current_topology.entry)
 	}

 	/// Return the grid topology for the specific session index, if no such a session is stored
 	/// returns `None`.
-	pub fn get_topology(&self, idx: SessionIndex) -> Option<&SessionGridTopology> {
+	pub fn get_topology(&self, idx: SessionIndex) -> Option<&SessionGridTopologyEntry> {
 		if let Some(prev_topology) = &self.prev_topology {
 			if idx == prev_topology.session_index {
-				return Some(&prev_topology.topology)
+				return Some(&prev_topology.entry)
 			}
 		}
 		if self.current_topology.session_index == idx {
-			return Some(&self.current_topology.topology)
+			return Some(&self.current_topology.entry)
 		}

 		None
 	}

 	/// Update the current topology preserving the previous one
-	pub fn update_topology(&mut self, session_index: SessionIndex, topology: SessionGridTopology) {
+	pub fn update_topology(
+		&mut self,
+		session_index: SessionIndex,
+		topology: SessionGridTopology,
+		local_index: Option<ValidatorIndex>,
+	) {
+		let local_neighbors = local_index
+			.and_then(|l| topology.compute_grid_neighbors_for(l))
+			.unwrap_or_else(GridNeighbors::empty);
+
 		let old_current = std::mem::replace(
 			&mut self.current_topology,
-			GridTopologySessionBound { topology, session_index },
+			GridTopologySessionBound {
+				entry: SessionGridTopologyEntry { topology, local_neighbors },
+				session_index,
+			},
 		);
 		self.prev_topology.replace(old_current);
 	}

 	/// Returns a current grid topology
-	pub fn get_current_topology(&self) -> &SessionGridTopology {
-		&self.current_topology.topology
+	pub fn get_current_topology(&self) -> &SessionGridTopologyEntry {
+		&self.current_topology.entry
+	}
+
+	/// Access the current grid topology mutably. Dangerous and intended
+	/// to be used in tests.
+	pub fn get_current_topology_mut(&mut self) -> &mut SessionGridTopologyEntry {
+		&mut self.current_topology.entry
 	}
 }

@@ -365,4 +546,27 @@ mod tests {
 		let mut random_routing = RandomRouting { target: 10, sent: 0, sample_rate: 10 };
 		assert_eq!(run_random_routing(&mut random_routing, &mut rng, 10, 100), 10);
 	}
+
+	#[test]
+	fn test_matrix_neighbors() {
+		for (our_index, len, expected_row, expected_column) in vec![
+			(0usize, 1usize, vec![], vec![]),
+			(1, 2, vec![], vec![0usize]),
+			(0, 9, vec![1, 2], vec![3, 6]),
+			(9, 10, vec![], vec![0, 3, 6]),
+			(10, 11, vec![9], vec![1, 4, 7]),
+			(7, 11, vec![6, 8], vec![1, 4, 10]),
+		]
+		.into_iter()
+		{
+			let matrix = matrix_neighbors(our_index, len).unwrap();
+			let mut row_result: Vec<_> = matrix.row_neighbors.collect();
+			let mut column_result: Vec<_> = matrix.column_neighbors.collect();
+			row_result.sort();
+			column_result.sort();
+
+			assert_eq!(row_result, expected_row);
+			assert_eq!(column_result, expected_column);
+		}
+	}
 }