fix bitfield selection (#1913)

polkadot/node/core/provisioner/src/lib.rs

@@ -39,9 +39,10 @@ use polkadot_node_subsystem_util::{
 };
 use polkadot_primitives::v1::{
 	BackedCandidate, BlockNumber, CoreState, Hash, OccupiedCoreAssumption,
-	SignedAvailabilityBitfield,
+	SignedAvailabilityBitfield, ValidatorIndex,
 };
-use std::{collections::HashSet, convert::TryFrom, pin::Pin};
+use std::{convert::TryFrom, pin::Pin};
+use std::collections::BTreeMap;
 use thiserror::Error;
 
 struct ProvisioningJob {
@@ -313,7 +314,7 @@ async fn send_inherent_data(
 /// In general, we want to pick all the bitfields. However, we have the following constraints:
 ///
 /// - not more than one per validator
-/// - each must correspond to an occupied core
+/// - each 1 bit must correspond to an occupied core
 ///
 /// If we have too many, an arbitrary selection policy is fine. For purposes of maximizing availability,
 /// we pick the one with the greatest number of 1 bits.
@@ -324,32 +325,30 @@ fn select_availability_bitfields(
 	cores: &[CoreState],
 	bitfields: &[SignedAvailabilityBitfield],
 ) -> Vec<SignedAvailabilityBitfield> {
-	let mut bitfield_per_core: Vec<Option<SignedAvailabilityBitfield>> = vec![None; cores.len()];
+	let mut selected: BTreeMap<ValidatorIndex, SignedAvailabilityBitfield> = BTreeMap::new();
-	let mut seen_validators = HashSet::new();
 
-	for mut bitfield in bitfields.iter().cloned() {
+	'a:
-		// If we have seen the validator already, ignore it.
+	for bitfield in bitfields.iter().cloned() {
-		if !seen_validators.insert(bitfield.validator_index()) {
+		if bitfield.payload().0.len() != cores.len() {
-			continue;
+			continue
 		}
 
-		for (idx, _) in cores.iter().enumerate().filter(|v| v.1.is_occupied()) {
+		let is_better = selected.get(&bitfield.validator_index())
+			.map_or(true, |b| b.payload().0.count_ones() < bitfield.payload().0.count_ones());
+
+		if !is_better { continue }
+
+		for (idx, _) in cores.iter().enumerate().filter(|v| !v.1.is_occupied()) {
+			// Bit is set for an unoccupied core - invalid
 			if *bitfield.payload().0.get(idx).unwrap_or(&false) {
-				if let Some(ref mut occupied) = bitfield_per_core[idx] {
+				continue 'a
-					if occupied.payload().0.count_ones() < bitfield.payload().0.count_ones() {
-						// We found a better bitfield, lets swap them and search a new spot for the old
-						// best one
-						std::mem::swap(occupied, &mut bitfield);
-					}
-				} else {
-					bitfield_per_core[idx] = Some(bitfield);
-					break;
-				}
 			}
 		}
+
+		let _ = selected.insert(bitfield.validator_index(), bitfield);
 	}
 
-	bitfield_per_core.into_iter().filter_map(|v| v).collect()
+	selected.into_iter().map(|(_, b)| b).collect()
 }
 
 /// Determine which cores are free, and then to the degree possible, pick a candidate appropriate to each free core.

polkadot/node/core/provisioner/src/tests.rs

@@ -10,7 +10,7 @@ pub fn occupied_core(para_id: u32) -> CoreState {
 		occupied_since: 100_u32,
 		time_out_at: 200_u32,
 		next_up_on_time_out: None,
-		availability: default_bitvec(),
+		availability: bitvec![bitvec::order::Lsb0, u8; 0; 32],
 	})
 }
 
@@ -28,8 +28,8 @@ where
 	CoreState::Occupied(core)
 }
 
-pub fn default_bitvec() -> CoreAvailability {
+pub fn default_bitvec(n_cores: usize) -> CoreAvailability {
-	bitvec![bitvec::order::Lsb0, u8; 0; 32]
+	bitvec![bitvec::order::Lsb0, u8; 0; n_cores]
 }
 
 pub fn scheduled_core(id: u32) -> ScheduledCore {
@@ -68,7 +68,7 @@ mod select_availability_bitfields {
 	#[test]
 	fn not_more_than_one_per_validator() {
 		let keystore: SyncCryptoStorePtr = Arc::new(KeyStore::new());
-		let mut bitvec = default_bitvec();
+		let mut bitvec = default_bitvec(2);
 		bitvec.set(0, true);
 		bitvec.set(1, true);
 
@@ -94,102 +94,98 @@ mod select_availability_bitfields {
 	#[test]
 	fn each_corresponds_to_an_occupied_core() {
 		let keystore: SyncCryptoStorePtr = Arc::new(KeyStore::new());
-		let bitvec = default_bitvec();
+		let bitvec = default_bitvec(3);
 
-		let cores = vec![CoreState::Free, CoreState::Scheduled(Default::default())];
+		// invalid: bit on free core
+		let mut bitvec0 = bitvec.clone();
+		bitvec0.set(0, true);
 
-		let bitfields = vec![
+		// invalid: bit on scheduled core
-			block_on(signed_bitfield(&keystore, bitvec.clone(), 0)),
+		let mut bitvec1 = bitvec.clone();
-			block_on(signed_bitfield(&keystore, bitvec.clone(), 1)),
+		bitvec1.set(1, true);
-			block_on(signed_bitfield(&keystore, bitvec, 1)),
+
+		// valid: bit on occupied core.
+		let mut bitvec2 = bitvec.clone();
+		bitvec2.set(2, true);
+
+		let cores = vec![
+			CoreState::Free,
+			CoreState::Scheduled(Default::default()),
+			occupied_core(2),
 		];
 
-		let mut selected_bitfields = select_availability_bitfields(&cores, &bitfields);
+		let bitfields = vec![
-		selected_bitfields.sort_by_key(|bitfield| bitfield.validator_index());
+			block_on(signed_bitfield(&keystore, bitvec0, 0)),
+			block_on(signed_bitfield(&keystore, bitvec1, 1)),
+			block_on(signed_bitfield(&keystore, bitvec2.clone(), 2)),
+		];
 
-		// bitfields not corresponding to occupied cores are not selected
+		let selected_bitfields = select_availability_bitfields(&cores, &bitfields);
-		assert!(selected_bitfields.is_empty());
+
+		// selects only the valid bitfield
+		assert_eq!(selected_bitfields.len(), 1);
+		assert_eq!(selected_bitfields[0].payload().0, bitvec2);
 	}
 
 	#[test]
 	fn more_set_bits_win_conflicts() {
 		let keystore: SyncCryptoStorePtr = Arc::new(KeyStore::new());
-		let mut bitvec = default_bitvec();
+		let mut bitvec = default_bitvec(2);
 		bitvec.set(0, true);
 
 		let mut bitvec1 = bitvec.clone();
 		bitvec1.set(1, true);
 
-		let cores = vec![occupied_core(0)];
+		let cores = vec![occupied_core(0), occupied_core(1)];
 
 		let bitfields = vec![
-			block_on(signed_bitfield(&keystore, bitvec, 0)),
+			block_on(signed_bitfield(&keystore, bitvec, 1)),
 			block_on(signed_bitfield(&keystore, bitvec1.clone(), 1)),
 		];
 
-		// this test is probablistic: chances are excellent that it does what it claims to.
-		// it cannot fail unless things are broken.
-		// however, there is a (very small) chance that it passes when things are broken.
-		for _ in 0..64 {
-			let selected_bitfields = select_availability_bitfields(&cores, &bitfields);
-			assert_eq!(selected_bitfields.len(), 1);
-			assert_eq!(selected_bitfields[0].payload().0, bitvec1.clone());
-		}
-	}
-
-	#[test]
-	fn more_validators_than_parachains() {
-		let keystore: SyncCryptoStorePtr = Arc::new(KeyStore::new());
-		let mut bitvec = default_bitvec();
-		bitvec.set(0, true);
-
-		let cores = vec![occupied_core(0)];
-
-		let bitfields = vec![
-			block_on(signed_bitfield(&keystore, bitvec.clone(), 0)),
-			block_on(signed_bitfield(&keystore, bitvec.clone(), 1)),
-			block_on(signed_bitfield(&keystore, bitvec.clone(), 2)),
-			block_on(signed_bitfield(&keystore, bitvec.clone(), 3)),
-		];
-
 		let selected_bitfields = select_availability_bitfields(&cores, &bitfields);
 		assert_eq!(selected_bitfields.len(), 1);
-		assert_eq!(selected_bitfields[0].payload().0, bitvec);
+		assert_eq!(selected_bitfields[0].payload().0, bitvec1.clone());
 	}
 
 	#[test]
 	fn more_complex_bitfields() {
 		let keystore: SyncCryptoStorePtr = Arc::new(KeyStore::new());
-		let mut bitvec0 = default_bitvec();
+
+		let cores = vec![occupied_core(0), occupied_core(1), occupied_core(2), occupied_core(3)];
+
+		let mut bitvec0 = default_bitvec(4);
 		bitvec0.set(0, true);
 		bitvec0.set(2, true);
 
-		let mut bitvec1 = default_bitvec();
+		let mut bitvec1 = default_bitvec(4);
 		bitvec1.set(1, true);
 
-		let mut bitvec2 = default_bitvec();
+		let mut bitvec2 = default_bitvec(4);
 		bitvec2.set(2, true);
 
-		let mut bitvec3 = default_bitvec();
+		let mut bitvec3 = default_bitvec(4);
 		bitvec3.set(0, true);
 		bitvec3.set(1, true);
 		bitvec3.set(2, true);
 		bitvec3.set(3, true);
 
-		let cores = vec![occupied_core(0), occupied_core(1), occupied_core(2), occupied_core(3)];
+		// these are out of order but will be selected in order. The better
-
+		// bitfield for 3 will be selected.
 		let bitfields = vec![
-			block_on(signed_bitfield(&keystore, bitvec0.clone(), 0)),
+			block_on(signed_bitfield(&keystore, bitvec2.clone(), 3)),
-			block_on(signed_bitfield(&keystore, bitvec1.clone(), 1)),
-			block_on(signed_bitfield(&keystore, bitvec2.clone(), 2)),
 			block_on(signed_bitfield(&keystore, bitvec3.clone(), 3)),
+			block_on(signed_bitfield(&keystore, bitvec0.clone(), 0)),
+			block_on(signed_bitfield(&keystore, bitvec2.clone(), 2)),
+			block_on(signed_bitfield(&keystore, bitvec1.clone(), 1)),
 		];
 
 		let selected_bitfields = select_availability_bitfields(&cores, &bitfields);
-		assert_eq!(selected_bitfields.len(), 3);
+		assert_eq!(selected_bitfields.len(), 4);
-		assert_eq!(selected_bitfields[0].payload().0, bitvec3);
+		assert_eq!(selected_bitfields[0].payload().0, bitvec0);
 		assert_eq!(selected_bitfields[1].payload().0, bitvec1);
-		assert_eq!(selected_bitfields[2].payload().0, bitvec0);
+		assert_eq!(selected_bitfields[2].payload().0, bitvec2);
+		assert_eq!(selected_bitfields[3].payload().0, bitvec3);
 	}
 }
 
@@ -358,6 +354,7 @@ mod select_candidates {
 	#[test]
 	fn selects_correct_candidates() {
 		let mock_cores = mock_availability_cores();
+		let n_cores = mock_cores.len();
 
 		let empty_hash = PersistedValidationData::<BlockNumber>::default().hash();
 
@@ -370,7 +367,7 @@ mod select_candidates {
 				..Default::default()
 			},
 			validity_votes: Vec::new(),
-			validator_indices: default_bitvec(),
+			validator_indices: default_bitvec(n_cores),
 		};
 
 		let candidates: Vec<_> = std::iter::repeat(candidate_template)