Optimize offchain worker memory usage a bit.

client/allocator/src/freeing_bump.rs

@@ -412,9 +412,18 @@ impl FreeingBumpHeapAllocator {
 
 		// Write the order in the occupied header.
 		Header::Occupied(order).write_into(mem, header_ptr)?;
-
 		self.total_size += order.size() + HEADER_SIZE;
 
+		if order.size() + HEADER_SIZE > 8 * 1024 * 1024 {
+			log::warn!(
+				target: LOG_TARGET,
+				"large allocation of {} detected, after allocation, total_size = {}, bumper = {}.",
+				order.size() + HEADER_SIZE,
+				self.total_size,
+				self.bumper,
+			);
+		}
+
 		log::trace!(
 			target: LOG_TARGET,
 			"after allocation, total_size = {}, bumper = {}.",

frame/election-provider-multi-phase/src/lib.rs

@@ -948,6 +948,11 @@ pub mod pallet {
 				compute: ElectionCompute::Emergency,
 			};
 
+			Self::deposit_event(Event::SolutionStored {
+				election_compute: ElectionCompute::Emergency,
+				prev_ejected: Self::queued_solution().is_some(),
+			});
+
 			<QueuedSolution<T>>::put(solution);
 			Ok(())
 		}
@@ -1057,6 +1062,11 @@ pub mod pallet {
 				compute: ElectionCompute::Fallback,
 			};
 
+			Self::deposit_event(Event::SolutionStored {
+				election_compute: ElectionCompute::Fallback,
+				prev_ejected: Self::queued_solution().is_some(),
+			});
+
 			<QueuedSolution<T>>::put(solution);
 			Ok(())
 		}
@@ -1106,6 +1116,10 @@ pub mod pallet {
 		OcwCallWrongEra,
 		/// Snapshot metadata should exist but didn't.
 		MissingSnapshotMetadata,
+		/// Snapshot should exist but didn't.
+		MissingSnapshot,
+		/// Codec error.
+		Codec,
 		/// `Self::insert_submission` returned an invalid index.
 		InvalidSubmissionIndex,
 		/// The call is not allowed at this point.
@@ -1444,7 +1458,8 @@ impl<T: Config> Pallet<T> {
 
 		// Read the entire snapshot.
 		let RoundSnapshot { voters: snapshot_voters, targets: snapshot_targets } =
-			Self::snapshot().ok_or(FeasibilityError::SnapshotUnavailable)?;
+			Self::read_snapshot_with_preallocate()
+				.map_err(|_| FeasibilityError::SnapshotUnavailable)?;
 
 		// ----- Start building. First, we need some closures.
 		let cache = helpers::generate_voter_cache::<T::MinerConfig>(&snapshot_voters);
@@ -1495,6 +1510,27 @@ impl<T: Config> Pallet<T> {
 		Ok(ReadySolution { supports, compute, score })
 	}
 
+	/// Should be used instead of the getter of `Snapshot` in memory-bounded code paths.
+	fn read_snapshot_with_preallocate() -> Result<RoundSnapshot<T>, Error<T>> {
+		use codec::MaxEncodedLen;
+		let snap = Self::snapshot_metadata().ok_or(Error::<T>::MissingSnapshot)?;
+		let voters_size = snap.voters as usize * <VoterOf<T>>::max_encoded_len();
+		let targets_size = snap.targets as usize * T::AccountId::max_encoded_len();
+
+		// we want to decode two vecs, which need two u32s at most for their size.
+		let initial_capacity = voters_size + targets_size + 4 + 4;
+		let mut buffer = Vec::<u8>::with_capacity(initial_capacity);
+
+		// fill this whole buffer, and decode into it.
+		buffer.resize(buffer.capacity(), 0);
+		let _leftover = sp_io::storage::read(&<Snapshot<T>>::hashed_key(), &mut buffer, 0)
+			.ok_or(Error::<T>::MissingSnapshot)?;
+
+		// buffer should have not re-allocated
+		debug_assert!(buffer.capacity() == initial_capacity);
+		<RoundSnapshot<T> as codec::Decode>::decode(&mut &*buffer).map_err(|_| Error::<T>::Codec)
+	}
+
 	/// Perform the tasks to be done after a new `elect` has been triggered:
 	///
 	/// 1. Increment round.

frame/election-provider-multi-phase/src/unsigned.rs

@@ -160,7 +160,7 @@ impl<T: Config> Pallet<T> {
 	pub fn mine_solution(
 	) -> Result<(RawSolution<SolutionOf<T::MinerConfig>>, SolutionOrSnapshotSize), MinerError> {
 		let RoundSnapshot { voters, targets } =
-			Self::snapshot().ok_or(MinerError::SnapshotUnAvailable)?;
+			Self::read_snapshot_with_preallocate().map_err(|_| MinerError::SnapshotUnAvailable)?;
 		let desired_targets = Self::desired_targets().ok_or(MinerError::SnapshotUnAvailable)?;
 		let (solution, score, size) = Miner::<T::MinerConfig>::mine_solution_with_snapshot::<
 			T::Solver,
@@ -169,27 +169,6 @@ impl<T: Config> Pallet<T> {
 		Ok((RawSolution { solution, score, round }, size))
 	}
 
-	/// Convert a raw solution from [`sp_npos_elections::ElectionResult`] to [`RawSolution`], which
-	/// is ready to be submitted to the chain.
-	///
-	/// Will always reduce the solution as well.
-	pub fn prepare_election_result<Accuracy: PerThing128>(
-		election_result: ElectionResult<T::AccountId, Accuracy>,
-	) -> Result<(RawSolution<SolutionOf<T::MinerConfig>>, SolutionOrSnapshotSize), MinerError> {
-		let RoundSnapshot { voters, targets } =
-			Self::snapshot().ok_or(MinerError::SnapshotUnAvailable)?;
-		let desired_targets = Self::desired_targets().ok_or(MinerError::SnapshotUnAvailable)?;
-		let (solution, score, size) =
-			Miner::<T::MinerConfig>::prepare_election_result_with_snapshot(
-				election_result,
-				voters,
-				targets,
-				desired_targets,
-			)?;
-		let round = Self::round();
-		Ok((RawSolution { solution, score, round }, size))
-	}
-
 	/// Attempt to restore a solution from cache. Otherwise, compute it fresh. Either way, submit
 	/// if our call's score is greater than that of the cached solution.
 	pub fn restore_or_compute_then_maybe_submit() -> Result<(), MinerError> {
@@ -441,7 +420,10 @@ impl<T: MinerConfig> Miner<T> {
 			})
 	}
 
-	/// Same as [`Pallet::prepare_election_result`], but the input snapshot mut be given as inputs.
+	/// Convert a raw solution from [`sp_npos_elections::ElectionResult`] to [`RawSolution`], which
+	/// is ready to be submitted to the chain.
+	///
+	/// Will always reduce the solution as well.
 	pub fn prepare_election_result_with_snapshot<Accuracy: PerThing128>(
 		election_result: ElectionResult<T::AccountId, Accuracy>,
 		voters: Vec<(T::AccountId, VoteWeight, BoundedVec<T::AccountId, T::MaxVotesPerVoter>)>,
@@ -463,7 +445,7 @@ impl<T: MinerConfig> Miner<T> {
 			SolutionOf::<T>::try_from(assignments).map(|s| s.encoded_size())
 		};
 
-		let ElectionResult { assignments, winners: _ } = election_result;
+		let ElectionResult { assignments, winners: _w } = election_result;
 
 		// Reduce (requires round-trip to staked form)
 		let sorted_assignments = {
@@ -1118,7 +1100,19 @@ mod tests {
 						distribution: vec![(10, PerU16::one())],
 					}],
 				};
-				let (solution, witness) = MultiPhase::prepare_election_result(result).unwrap();
+
+				let RoundSnapshot { voters, targets } = MultiPhase::snapshot().unwrap();
+				let desired_targets = MultiPhase::desired_targets().unwrap();
+
+				let (raw, score, witness) =
+					Miner::<Runtime>::prepare_election_result_with_snapshot(
+						result,
+						voters.clone(),
+						targets.clone(),
+						desired_targets,
+					)
+					.unwrap();
+				let solution = RawSolution { solution: raw, score, round: MultiPhase::round() };
 				assert_ok!(MultiPhase::unsigned_pre_dispatch_checks(&solution));
 				assert_ok!(MultiPhase::submit_unsigned(
 					Origin::none(),
@@ -1139,7 +1133,14 @@ mod tests {
 						},
 					],
 				};
-				let (solution, _) = MultiPhase::prepare_election_result(result).unwrap();
+				let (raw, score, _) = Miner::<Runtime>::prepare_election_result_with_snapshot(
+					result,
+					voters.clone(),
+					targets.clone(),
+					desired_targets,
+				)
+				.unwrap();
+				let solution = RawSolution { solution: raw, score, round: MultiPhase::round() };
 				// 12 is not 50% more than 10
 				assert_eq!(solution.score.minimal_stake, 12);
 				assert_noop!(
@@ -1161,7 +1162,15 @@ mod tests {
 						},
 					],
 				};
-				let (solution, witness) = MultiPhase::prepare_election_result(result).unwrap();
+				let (raw, score, witness) =
+					Miner::<Runtime>::prepare_election_result_with_snapshot(
+						result,
+						voters.clone(),
+						targets.clone(),
+						desired_targets,
+					)
+					.unwrap();
+				let solution = RawSolution { solution: raw, score, round: MultiPhase::round() };
 				assert_eq!(solution.score.minimal_stake, 17);
 
 				// and it is fine

frame/support/src/storage/types/nmap.rs

@@ -544,7 +544,7 @@ mod test {
 	use crate::{
 		hash::{StorageHasher as _, *},
 		metadata::{StorageEntryModifier, StorageHasher},
-		storage::types::{Key, Key as NMapKey, ValueQuery},
+		storage::types::{Key, ValueQuery},
 	};
 	use sp_io::{hashing::twox_128, TestExternalities};
 

primitives/npos-elections/src/assignments.rs

@@ -120,18 +120,17 @@ impl<AccountId> StakedAssignment<AccountId> {
 		AccountId: IdentifierT,
 	{
 		let stake = self.total();
-		let distribution = self
-			.distribution
-			.into_iter()
-			.filter_map(|(target, w)| {
-				let per_thing = P::from_rational(w, stake);
-				if per_thing == Bounded::min_value() {
-					None
-				} else {
-					Some((target, per_thing))
-				}
-			})
-			.collect::<Vec<(AccountId, P)>>();
+		// most likely, the size of the staked assignment and normal assignments will be the same,
+		// so we pre-allocate it to prevent a sudden 2x allocation. `filter_map` starts with a size
+		// of 0 by default.
+		// https://www.reddit.com/r/rust/comments/3spfh1/does_collect_allocate_more_than_once_while/
+		let mut distribution = Vec::<(AccountId, P)>::with_capacity(self.distribution.len());
+		self.distribution.into_iter().for_each(|(target, w)| {
+			let per_thing = P::from_rational(w, stake);
+			if per_thing != Bounded::min_value() {
+				distribution.push((target, per_thing));
+			}
+		});
 
 		Assignment { who: self.who, distribution }
 	}

utils/frame/try-runtime/cli/src/commands/offchain_worker.rs

@@ -131,6 +131,11 @@ where
 		let builder = command.state.builder::<Block>()?;
 
 		let builder = if command.overwrite_wasm_code {
+			log::info!(
+				target: LOG_TARGET,
+				"replacing the in-storage :code: with the local code from {}'s chain_spec (your local repo)",
+				config.chain_spec.name(),
+			);
 			let (code_key, code) = extract_code(&config.chain_spec)?;
 			builder.inject_hashed_key_value(&[(code_key, code)])
 		} else {

utils/frame/try-runtime/cli/src/lib.rs

@@ -710,7 +710,7 @@ pub(crate) fn state_machine_call<Block: BlockT, D: NativeExecutionDispatch + 'st
 		sp_core::testing::TaskExecutor::new(),
 	)
 	.execute(execution.into())
-	.map_err(|e| format!("failed to execute 'TryRuntime_on_runtime_upgrade': {}", e))
+	.map_err(|e| format!("failed to execute '{}': {}", method, e))
 	.map_err::<sc_cli::Error, _>(Into::into)?;
 
 	Ok((changes, encoded_results))