Fix offchain election to respect the weight (#7215)

* Mockup

* Fix offchain election to respect the weight

* Fix builds a bit

* Update frame/staking/src/offchain_election.rs

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

* Update frame/staking/src/offchain_election.rs

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

* Make it build, binary search

* Fix a number of grumbles

* one more fix.

* remove unwrap.

* better alg.

* Better alg again.

* Final fixes

* Fix

* Rollback to normal

* Final touches.

* Better tests.

* Update frame/staking/src/lib.rs

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

* Proper maxExtWeight

* Final fix

* Final fix for the find_voter

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>
Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

substrate/frame/staking/src/benchmarking.rs

@@ -522,7 +522,7 @@ benchmarks! {
 			compact,
 			score,
 			size
-		) = offchain_election::prepare_submission::<T>(assignments, winners, false).unwrap();
+		) = offchain_election::prepare_submission::<T>(assignments, winners, false, T::MaximumBlockWeight::get()).unwrap();
 
 		assert_eq!(
 			winners.len(), compact.unique_targets().len(),
@@ -590,7 +590,7 @@ benchmarks! {
 			compact,
 			score,
 			size
-		) = offchain_election::prepare_submission::<T>(assignments, winners, false).unwrap();
+		) = offchain_election::prepare_submission::<T>(assignments, winners, false, T::MaximumBlockWeight::get()).unwrap();
 
 		assert_eq!(
 			winners.len(), compact.unique_targets().len(),

substrate/frame/staking/src/lib.rs

@@ -695,7 +695,7 @@ pub enum ElectionStatus<BlockNumber> {
 /// Note that these values must reflect the __total__ number, not only those that are present in the
 /// solution. In short, these should be the same size as the size of the values dumped in
 /// `SnapshotValidators` and `SnapshotNominators`.
-#[derive(PartialEq, Eq, Clone, Encode, Decode, RuntimeDebug, Default)]
+#[derive(PartialEq, Eq, Clone, Copy, Encode, Decode, RuntimeDebug, Default)]
 pub struct ElectionSize {
 	/// Number of validators in the snapshot of the current election round.
 	#[codec(compact)]
@@ -883,6 +883,13 @@ pub trait Trait: frame_system::Trait + SendTransactionTypes<Call<Self>> {
 	/// multiple pallets send unsigned transactions.
 	type UnsignedPriority: Get<TransactionPriority>;
 
+	/// Maximum weight that the unsigned transaction can have.
+	///
+	/// Chose this value with care. On one hand, it should be as high as possible, so the solution
+	/// can contain as many nominators/validators as possible. On the other hand, it should be small
+	/// enough to fit in the block.
+	type OffchainSolutionWeightLimit: Get<Weight>;
+
 	/// Weight information for extrinsics in this pallet.
 	type WeightInfo: WeightInfo;
 }
@@ -1294,6 +1301,7 @@ decl_module! {
 				consumed_weight += T::DbWeight::get().reads_writes(reads, writes);
 				consumed_weight += weight;
 			};
+
 			if
 				// if we don't have any ongoing offchain compute.
 				Self::era_election_status().is_closed() &&
@@ -1339,12 +1347,12 @@ decl_module! {
 			if Self::era_election_status().is_open_at(now) {
 				let offchain_status = set_check_offchain_execution_status::<T>(now);
 				if let Err(why) = offchain_status {
-					log!(debug, "skipping offchain worker in open election window due to [{}]", why);
+					log!(warn, "💸 skipping offchain worker in open election window due to [{}]", why);
 				} else {
 					if let Err(e) = compute_offchain_election::<T>() {
 						log!(error, "💸 Error in election offchain worker: {:?}", e);
 					} else {
-						log!(debug, "Executed offchain worker thread without errors.");
+						log!(debug, "💸 Executed offchain worker thread without errors.");
 					}
 				}
 			}
@@ -2141,7 +2149,7 @@ decl_module! {
 		/// transaction in the block.
 		///
 		/// # <weight>
-		/// See `crate::weight` module.
+		/// See [`submit_election_solution`].
 		/// # </weight>
 		#[weight = T::WeightInfo::submit_solution_better(
 			size.validators.into(),
@@ -2171,6 +2179,7 @@ decl_module! {
 				effectively depriving the validators from their authoring reward. Hence, this panic
 				is expected."
 			);
+
 			Ok(adjustments)
 		}
 	}
@@ -3082,7 +3091,6 @@ impl<T: Trait> Module<T> {
 	pub fn set_slash_reward_fraction(fraction: Perbill) {
 		SlashRewardFraction::put(fraction);
 	}
-
 }
 
 /// In this implementation `new_session(session)` must be called before `end_session(session-1)`
@@ -3367,13 +3375,13 @@ impl<T: Trait> frame_support::unsigned::ValidateUnsigned for Module<T> {
 				let invalid = to_invalid(error_with_post_info);
 				log!(
 					debug,
-					"validate unsigned pre dispatch checks failed due to error #{:?}.",
+					"💸 validate unsigned pre dispatch checks failed due to error #{:?}.",
 					invalid,
 				);
-				return invalid .into();
+				return invalid.into();
 			}
 
-			log!(debug, "validateUnsigned succeeded for a solution at era {}.", era);
+			log!(debug, "💸 validateUnsigned succeeded for a solution at era {}.", era);
 
 			ValidTransaction::with_tag_prefix("StakingOffchain")
 				// The higher the score[0], the better a solution is.

substrate/frame/staking/src/mock.rs

@@ -17,24 +17,29 @@
 
 //! Test utilities
 
-use std::{collections::HashSet, cell::RefCell};
-use sp_runtime::Perbill;
-use sp_runtime::curve::PiecewiseLinear;
-use sp_runtime::traits::{IdentityLookup, Convert, SaturatedConversion, Zero};
-use sp_runtime::testing::{Header, UintAuthorityId, TestXt};
-use sp_staking::{SessionIndex, offence::{OffenceDetails, OnOffenceHandler}};
-use sp_core::H256;
+use crate::*;
 use frame_support::{
-	assert_ok, impl_outer_origin, parameter_types, impl_outer_dispatch, impl_outer_event,
-	StorageValue, StorageMap, StorageDoubleMap, IterableStorageMap,
-	traits::{Currency, Get, FindAuthor, OnFinalize, OnInitialize},
-	weights::{Weight, constants::RocksDbWeight},
+	assert_ok, impl_outer_dispatch, impl_outer_event, impl_outer_origin, parameter_types,
+	traits::{Currency, FindAuthor, Get, OnFinalize, OnInitialize},
+	weights::{constants::RocksDbWeight, Weight},
+	IterableStorageMap, StorageDoubleMap, StorageMap, StorageValue,
 };
+use sp_core::H256;
 use sp_io;
 use sp_npos_elections::{
-	build_support_map, evaluate_support, reduce, ExtendedBalance, StakedAssignment, ElectionScore,
+	build_support_map, evaluate_support, reduce, ElectionScore, ExtendedBalance, StakedAssignment,
 };
-use crate::*;
+use sp_runtime::{
+	curve::PiecewiseLinear,
+	testing::{Header, TestXt, UintAuthorityId},
+	traits::{Convert, IdentityLookup, SaturatedConversion, Zero},
+	Perbill,
+};
+use sp_staking::{
+	offence::{OffenceDetails, OnOffenceHandler},
+	SessionIndex,
+};
+use std::{cell::RefCell, collections::HashSet};
 
 pub const INIT_TIMESTAMP: u64 = 30_000;
 
@@ -194,7 +199,7 @@ pub struct Test;
 
 parameter_types! {
 	pub const BlockHashCount: u64 = 250;
-	pub const MaximumBlockWeight: Weight = 1024;
+	pub const MaximumBlockWeight: Weight = frame_support::weights::constants::WEIGHT_PER_SECOND * 2;
 	pub const MaximumBlockLength: u32 = 2 * 1024;
 	pub const AvailableBlockRatio: Perbill = Perbill::one();
 	pub const MaxLocks: u32 = 1024;
@@ -293,6 +298,7 @@ parameter_types! {
 	pub const MaxNominatorRewardedPerValidator: u32 = 64;
 	pub const UnsignedPriority: u64 = 1 << 20;
 	pub const MinSolutionScoreBump: Perbill = Perbill::zero();
+	pub const OffchainSolutionWeightLimit: Weight = MaximumBlockWeight::get();
 }
 
 thread_local! {
@@ -331,10 +337,12 @@ impl Trait for Test {
 	type MinSolutionScoreBump = MinSolutionScoreBump;
 	type MaxNominatorRewardedPerValidator = MaxNominatorRewardedPerValidator;
 	type UnsignedPriority = UnsignedPriority;
+	type OffchainSolutionWeightLimit = OffchainSolutionWeightLimit;
 	type WeightInfo = ();
 }
 
-impl<LocalCall> frame_system::offchain::SendTransactionTypes<LocalCall> for Test where
+impl<LocalCall> frame_system::offchain::SendTransactionTypes<LocalCall> for Test
+where
 	Call: From<LocalCall>,
 {
 	type OverarchingCall = Call;

substrate/frame/staking/src/offchain_election.rs

@@ -17,19 +17,20 @@
 
 //! Helpers for offchain worker election.
 
-use codec::Decode;
 use crate::{
-	Call, CompactAssignments, Module, NominatorIndex, OffchainAccuracy, Trait, ValidatorIndex,
-	ElectionSize,
+	Call, CompactAssignments, ElectionSize, Module, NominatorIndex, Nominators, OffchainAccuracy,
+	Trait, ValidatorIndex, WeightInfo,
 };
+use codec::Decode;
+use frame_support::{traits::Get, weights::Weight, IterableStorageMap};
 use frame_system::offchain::SubmitTransaction;
 use sp_npos_elections::{
-	build_support_map, evaluate_support, reduce, Assignment, ExtendedBalance, ElectionResult,
-	ElectionScore,
+	build_support_map, evaluate_support, reduce, Assignment, ElectionResult, ElectionScore,
+	ExtendedBalance,
+};
+use sp_runtime::{
+	offchain::storage::StorageValueRef, traits::TrailingZeroInput, PerThing, RuntimeDebug,
 };
-use sp_runtime::offchain::storage::StorageValueRef;
-use sp_runtime::{PerThing, RuntimeDebug, traits::TrailingZeroInput};
-use frame_support::traits::Get;
 use sp_std::{convert::TryInto, prelude::*};
 
 /// Error types related to the offchain election machinery.
@@ -46,6 +47,8 @@ pub enum OffchainElectionError {
 	InternalElectionError(sp_npos_elections::Error),
 	/// One of the computed winners is invalid.
 	InvalidWinner,
+	/// A nominator is not available in the snapshot.
+	NominatorSnapshotCorrupt,
 }
 
 impl From<sp_npos_elections::Error> for OffchainElectionError {
@@ -115,11 +118,16 @@ pub(crate) fn compute_offchain_election<T: Trait>() -> Result<(), OffchainElecti
 		.ok_or(OffchainElectionError::ElectionFailed)?;
 
 	// process and prepare it for submission.
-	let (winners, compact, score, size) = prepare_submission::<T>(assignments, winners, true)?;
+	let (winners, compact, score, size) = prepare_submission::<T>(
+		assignments,
+		winners,
+		true,
+		T::OffchainSolutionWeightLimit::get(),
+	)?;
 
 	crate::log!(
 		info,
-		"prepared a seq-phragmen solution with {} balancing iterations and score {:?}",
+		"💸 prepared a seq-phragmen solution with {} balancing iterations and score {:?}",
 		iters,
 		score,
 	);
@@ -155,6 +163,162 @@ pub fn get_balancing_iters<T: Trait>() -> usize {
 	}
 }
 
+/// Find the maximum `len` that a compact can have in order to fit into the block weight.
+///
+/// This only returns a value between zero and `size.nominators`.
+pub fn maximum_compact_len<W: crate::WeightInfo>(
+	winners_len: u32,
+	size: ElectionSize,
+	max_weight: Weight,
+) -> u32 {
+	use sp_std::cmp::Ordering;
+
+	if size.nominators < 1 {
+		return size.nominators;
+	}
+
+	let max_voters = size.nominators.max(1);
+	let mut voters = max_voters;
+
+	// helper closures.
+	let weight_with = |voters: u32| -> Weight {
+		W::submit_solution_better(
+			size.validators.into(),
+			size.nominators.into(),
+			voters,
+			winners_len,
+		)
+	};
+
+	let next_voters = |current_weight: Weight, voters: u32, step: u32| -> Result<u32, ()> {
+		match current_weight.cmp(&max_weight) {
+			Ordering::Less => {
+				let next_voters = voters.checked_add(step);
+				match next_voters {
+					Some(voters) if voters < max_voters => Ok(voters),
+					_ => Err(()),
+				}
+			},
+			Ordering::Greater => voters.checked_sub(step).ok_or(()),
+			Ordering::Equal => Ok(voters),
+		}
+	};
+
+	// First binary-search the right amount of voters
+	let mut step = voters / 2;
+	let mut current_weight = weight_with(voters);
+	while step > 0 {
+		match next_voters(current_weight, voters, step) {
+			// proceed with the binary search
+			Ok(next) if next != voters => {
+				voters = next;
+			},
+			// we are out of bounds, break out of the loop.
+			Err(()) => {
+				break;
+			},
+			// we found the right value - early exit the function.
+			Ok(next) => return next
+		}
+		step = step / 2;
+		current_weight = weight_with(voters);
+	}
+
+
+	// Time to finish.
+	// We might have reduced less than expected due to rounding error. Increase one last time if we
+	// have any room left, the reduce until we are sure we are below limit.
+	while voters + 1 <= max_voters && weight_with(voters + 1) < max_weight {
+		voters += 1;
+	}
+	while voters.checked_sub(1).is_some() && weight_with(voters) > max_weight {
+		voters -= 1;
+	}
+
+	debug_assert!(
+		weight_with(voters.min(size.nominators)) <= max_weight,
+		"weight_with({}) <= {}", voters.min(size.nominators), max_weight,
+	);
+	voters.min(size.nominators)
+}
+
+/// Greedily reduce the size of the a solution to fit into the block, w.r.t. weight.
+///
+/// The weight of the solution is foremost a function of the number of voters (i.e.
+/// `compact.len()`). Aside from this, the other components of the weight are invariant. The number
+/// of winners shall not be changed (otherwise the solution is invalid) and the `ElectionSize` is
+/// merely a representation of the total number of stakers.
+///
+/// Thus, we reside to stripping away some voters. This means only changing the `compact` struct.
+///
+/// Note that the solution is already computed, and the winners are elected based on the merit of
+/// teh entire stake in the system. Nonetheless, some of the voters will be removed further down the
+/// line.
+///
+/// Indeed, the score must be computed **after** this step. If this step reduces the score too much,
+/// then the solution will be discarded.
+pub fn trim_to_weight<T: Trait, FN>(
+	maximum_allowed_voters: u32,
+	mut compact: CompactAssignments,
+	nominator_index: FN,
+) -> Result<CompactAssignments, OffchainElectionError>
+where
+	for<'r> FN: Fn(&'r T::AccountId) -> Option<NominatorIndex>,
+{
+	match compact.len().checked_sub(maximum_allowed_voters as usize) {
+		Some(to_remove) if to_remove > 0 => {
+			// grab all voters and sort them by least stake.
+			let mut voters_sorted = <Nominators<T>>::iter()
+				.map(|(who, _)| {
+					(
+						who.clone(),
+						<Module<T>>::slashable_balance_of_vote_weight(&who),
+					)
+				})
+				.collect::<Vec<_>>();
+			voters_sorted.sort_by_key(|(_, y)| *y);
+
+			// start removing from the least stake. Iterate until we know enough have been removed.
+			let mut removed = 0;
+			for (maybe_index, _stake) in voters_sorted
+				.iter()
+				.map(|(who, stake)| (nominator_index(&who), stake))
+			{
+				let index = maybe_index.ok_or(OffchainElectionError::NominatorSnapshotCorrupt)?;
+				if compact.remove_voter(index) {
+					crate::log!(
+						trace,
+						"💸 removed a voter at index {} with stake {:?} from compact to reduce the size",
+						index,
+						_stake,
+					);
+					removed += 1
+				}
+
+				if removed >= to_remove {
+					break;
+				}
+			}
+
+			crate::log!(
+					warn,
+					"💸 {} nominators out of {} had to be removed from compact solution due to size limits.",
+					removed,
+					compact.len() + removed,
+				);
+			Ok(compact)
+		}
+		_ => {
+			// nada, return as-is
+			crate::log!(
+				info,
+				"💸 Compact solution did not get trimmed due to block weight limits.",
+			);
+			Ok(compact)
+		}
+	}
+}
+
 /// Takes an election result and spits out some data that can be submitted to the chain.
 ///
 /// This does a lot of stuff; read the inline comments.
@@ -162,12 +326,17 @@ pub fn prepare_submission<T: Trait>(
 	assignments: Vec<Assignment<T::AccountId, OffchainAccuracy>>,
 	winners: Vec<(T::AccountId, ExtendedBalance)>,
 	do_reduce: bool,
-) -> Result<(
-	Vec<ValidatorIndex>,
-	CompactAssignments,
-	ElectionScore,
-	ElectionSize,
-), OffchainElectionError> where
+	maximum_weight: Weight,
+) -> Result<
+	(
+		Vec<ValidatorIndex>,
+		CompactAssignments,
+		ElectionScore,
+		ElectionSize,
+	),
+	OffchainElectionError,
+>
+where
 	ExtendedBalance: From<<OffchainAccuracy as PerThing>::Inner>,
 {
 	// make sure that the snapshot is available.
@@ -176,7 +345,7 @@ pub fn prepare_submission<T: Trait>(
 	let snapshot_nominators =
 		<Module<T>>::snapshot_nominators().ok_or(OffchainElectionError::SnapshotUnavailable)?;
 
-	// all helper closures
+	// all helper closures that we'd ever need.
 	let nominator_index = |a: &T::AccountId| -> Option<NominatorIndex> {
 		snapshot_nominators
 			.iter()
@@ -189,6 +358,19 @@ pub fn prepare_submission<T: Trait>(
 			.position(|x| x == a)
 			.and_then(|i| <usize as TryInto<ValidatorIndex>>::try_into(i).ok())
 	};
+	let nominator_at = |i: NominatorIndex| -> Option<T::AccountId> {
+		snapshot_nominators.get(i as usize).cloned()
+	};
+
+	let validator_at = |i: ValidatorIndex| -> Option<T::AccountId> {
+		snapshot_validators.get(i as usize).cloned()
+	};
+
+	// both conversions are safe; snapshots are not created if they exceed.
+	let size = ElectionSize {
+		validators: snapshot_validators.len() as ValidatorIndex,
+		nominators: snapshot_nominators.len() as NominatorIndex,
+	};
 
 	// Clean winners.
 	let winners = sp_npos_elections::to_without_backing(winners);
@@ -208,17 +390,40 @@ pub fn prepare_submission<T: Trait>(
 	let low_accuracy_assignment = sp_npos_elections::assignment_staked_to_ratio_normalized(staked)
 		.map_err(|e| OffchainElectionError::from(e))?;
 
-	// convert back to staked to compute the score in the receiver's accuracy. This can be done
-	// nicer, for now we do it as such since this code is not time-critical. This ensure that the
-	// score _predicted_ here is the same as the one computed on chain and you will not get a
-	// `PhragmenBogusScore` error. This is totally NOT needed if we don't do reduce. This whole
-	// _accuracy glitch_ happens because reduce breaks that assumption of rounding and **scale**.
-	// The initial phragmen results are computed in `OffchainAccuracy` and the initial `staked`
-	// assignment set is also all multiples of this value. After reduce, this no longer holds. Hence
-	// converting to ratio thereafter is not trivially reversible.
+	// compact encode the assignment.
+	let compact = CompactAssignments::from_assignment(
+		low_accuracy_assignment,
+		nominator_index,
+		validator_index,
+	)
+	.map_err(|e| OffchainElectionError::from(e))?;
+
+	// potentially reduce the size of the compact to fit weight.
+	let maximum_allowed_voters =
+		maximum_compact_len::<T::WeightInfo>(winners.len() as u32, size, maximum_weight);
+
+	crate::log!(debug, "💸 Maximum weight = {:?} // current weight = {:?} // maximum voters = {:?} // current votes = {:?}",
+		maximum_weight,
+		T::WeightInfo::submit_solution_better(
+				size.validators.into(),
+				size.nominators.into(),
+				compact.len() as u32,
+				winners.len() as u32,
+		),
+		maximum_allowed_voters,
+		compact.len(),
+	);
+
+	let compact = trim_to_weight::<T, _>(maximum_allowed_voters, compact, &nominator_index)?;
+
+	// re-compute the score. We re-create what the chain will do. This is a bit verbose and wastes
+	// CPU time, but it is necessary to ensure that the score that we claim is the same as the one
+	// calculated by the chain.
 	let score = {
+		let compact = compact.clone();
+		let assignments = compact.into_assignment(nominator_at, validator_at).unwrap();
 		let staked = sp_npos_elections::assignment_ratio_to_staked(
-			low_accuracy_assignment.clone(),
+			assignments,
 			<Module<T>>::slashable_balance_of_vote_weight,
 		);
 
@@ -227,13 +432,6 @@ pub fn prepare_submission<T: Trait>(
 		evaluate_support::<T::AccountId>(&support_map)
 	};
 
-	// compact encode the assignment.
-	let compact = CompactAssignments::from_assignment(
-		low_accuracy_assignment,
-		nominator_index,
-		validator_index,
-	).map_err(|e| OffchainElectionError::from(e))?;
-
 	// winners to index. Use a simple for loop for a more expressive early exit in case of error.
 	let mut winners_indexed: Vec<ValidatorIndex> = Vec::with_capacity(winners.len());
 	for w in winners {
@@ -247,11 +445,152 @@ pub fn prepare_submission<T: Trait>(
 		}
 	}
 
-	// both conversions are safe; snapshots are not created if they exceed.
-	let size = ElectionSize {
-		validators: snapshot_validators.len() as ValidatorIndex,
-		nominators: snapshot_nominators.len() as NominatorIndex,
-	};
-
 	Ok((winners_indexed, compact, score, size))
 }
+
+#[cfg(test)]
+mod test {
+	#![allow(unused_variables)]
+	use super::*;
+	use crate::ElectionSize;
+
+	struct Staking;
+
+	impl crate::WeightInfo for Staking {
+		fn bond() -> Weight {
+			unimplemented!()
+		}
+		fn bond_extra() -> Weight {
+			unimplemented!()
+		}
+		fn unbond() -> Weight {
+			unimplemented!()
+		}
+		fn withdraw_unbonded_update(s: u32) -> Weight {
+			unimplemented!()
+		}
+		fn withdraw_unbonded_kill(s: u32) -> Weight {
+			unimplemented!()
+		}
+		fn validate() -> Weight {
+			unimplemented!()
+		}
+		fn nominate(n: u32) -> Weight {
+			unimplemented!()
+		}
+		fn chill() -> Weight {
+			unimplemented!()
+		}
+		fn set_payee() -> Weight {
+			unimplemented!()
+		}
+		fn set_controller() -> Weight {
+			unimplemented!()
+		}
+		fn set_validator_count() -> Weight {
+			unimplemented!()
+		}
+		fn force_no_eras() -> Weight {
+			unimplemented!()
+		}
+		fn force_new_era() -> Weight {
+			unimplemented!()
+		}
+		fn force_new_era_always() -> Weight {
+			unimplemented!()
+		}
+		fn set_invulnerables(v: u32) -> Weight {
+			unimplemented!()
+		}
+		fn force_unstake(s: u32) -> Weight {
+			unimplemented!()
+		}
+		fn cancel_deferred_slash(s: u32) -> Weight {
+			unimplemented!()
+		}
+		fn payout_stakers_dead_controller(n: u32) -> Weight {
+			unimplemented!()
+		}
+		fn payout_stakers_alive_staked(n: u32) -> Weight {
+			unimplemented!()
+		}
+		fn rebond(l: u32) -> Weight {
+			unimplemented!()
+		}
+		fn set_history_depth(e: u32) -> Weight {
+			unimplemented!()
+		}
+		fn reap_stash(s: u32) -> Weight {
+			unimplemented!()
+		}
+		fn new_era(v: u32, n: u32) -> Weight {
+			unimplemented!()
+		}
+		fn submit_solution_better(v: u32, n: u32, a: u32, w: u32) -> Weight {
+			(0 * v + 0 * n + 1000 * a + 0 * w) as Weight
+		}
+	}
+
+	#[test]
+	fn find_max_voter_binary_search_works() {
+		let size = ElectionSize {
+			validators: 0,
+			nominators: 10,
+		};
+
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 0), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 999), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1000), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1001), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1990), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1999), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2000), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2001), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2010), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2990), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2999), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 3000), 3);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 3333), 3);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 5500), 5);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 7777), 7);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 9999), 9);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 10_000), 10);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 10_999), 10);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 11_000), 10);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 22_000), 10);
+
+		let size = ElectionSize {
+			validators: 0,
+			nominators: 1,
+		};
+
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 0), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 999), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1000), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1001), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1990), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1999), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2000), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2001), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2010), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 3333), 1);
+
+		let size = ElectionSize {
+			validators: 0,
+			nominators: 2,
+		};
+
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 0), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 999), 0);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1000), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1001), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 1999), 1);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2000), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2001), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 2010), 2);
+		assert_eq!(maximum_compact_len::<Staking>(0, size, 3333), 2);
+	}
+}

substrate/frame/staking/src/testing_utils.rs

@@ -277,7 +277,12 @@ pub fn get_weak_solution<T: Trait>(
 /// worker code.
 pub fn get_seq_phragmen_solution<T: Trait>(
 	do_reduce: bool,
-) -> (Vec<ValidatorIndex>, CompactAssignments, ElectionScore, ElectionSize) {
+) -> (
+	Vec<ValidatorIndex>,
+	CompactAssignments,
+	ElectionScore,
+	ElectionSize,
+) {
 	let iters = offchain_election::get_balancing_iters::<T>();
 
 	let sp_npos_elections::ElectionResult {
@@ -285,22 +290,42 @@ pub fn get_seq_phragmen_solution<T: Trait>(
 		assignments,
 	} = <Module<T>>::do_phragmen::<OffchainAccuracy>(iters).unwrap();
 
-	offchain_election::prepare_submission::<T>(assignments, winners, do_reduce).unwrap()
+	offchain_election::prepare_submission::<T>(
+		assignments,
+		winners,
+		do_reduce,
+		T::MaximumBlockWeight::get(),
+	)
+	.unwrap()
 }
 
 /// Returns a solution in which only one winner is elected with just a self vote.
 pub fn get_single_winner_solution<T: Trait>(
-	winner: T::AccountId
-) -> Result<(Vec<ValidatorIndex>, CompactAssignments, ElectionScore, ElectionSize), &'static str> {
+	winner: T::AccountId,
+) -> Result<
+	(
+		Vec<ValidatorIndex>,
+		CompactAssignments,
+		ElectionScore,
+		ElectionSize,
+	),
+	&'static str,
+> {
 	let snapshot_validators = <Module<T>>::snapshot_validators().unwrap();
 	let snapshot_nominators = <Module<T>>::snapshot_nominators().unwrap();
 
-	let val_index = snapshot_validators.iter().position(|x| *x == winner).ok_or("not a validator")?;
-	let nom_index = snapshot_nominators.iter().position(|x| *x == winner).ok_or("not a nominator")?;
+	let val_index = snapshot_validators
+		.iter()
+		.position(|x| *x == winner)
+		.ok_or("not a validator")?;
+	let nom_index = snapshot_nominators
+		.iter()
+		.position(|x| *x == winner)
+		.ok_or("not a nominator")?;
 
 	let stake = <Staking<T>>::slashable_balance_of(&winner);
-	let stake = <T::CurrencyToVote as Convert<BalanceOf<T>, VoteWeight>>::convert(stake)
-		as ExtendedBalance;
+	let stake =
+		<T::CurrencyToVote as Convert<BalanceOf<T>, VoteWeight>>::convert(stake) as ExtendedBalance;
 
 	let val_index = val_index as ValidatorIndex;
 	let nom_index = nom_index as NominatorIndex;

substrate/frame/staking/src/tests.rs

@@ -3100,7 +3100,7 @@ mod offchain_election {
 	}
 
 	#[test]
-	#[ignore] // This takes a few mins
+	#[ignore]
 	fn offchain_wont_work_if_snapshot_fails() {
 		ExtBuilder::default()
 			.offchain_election_ext()
@@ -3382,8 +3382,8 @@ mod offchain_election {
 	}
 
 	#[test]
-	fn offchain_worker_runs_with_equalise() {
-		// Offchain worker equalises based on the number provided by randomness. See the difference
+	fn offchain_worker_runs_with_balancing() {
+		// Offchain worker balances based on the number provided by randomness. See the difference
 		// in the priority, which comes from the computed score.
 		let mut ext = ExtBuilder::default()
 			.offchain_election_ext()