pallet_staking/
election_size_tracker.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! ## A static size tracker for the election snapshot data.
19//!
20//! ### Overview
21//!
22//! The goal of the size tracker is to provide a static, no-allocation byte tracker to be
23//! used by the election data provider when preparing the results of
24//! [`ElectionDataProvider::electing_voters`]. The [`StaticTracker`] implementation uses
25//! [`codec::Encode::size_hint`] to estimate the SCALE encoded size of the snapshot voters struct
26//! as it is being constructed without requiring extra stack allocations.
27//!
28//! The [`StaticTracker::try_register_voter`] is called to update the static tracker internal
29//! state, if It will return an error if the resulting SCALE encoded size (in bytes) is larger than
30//! the provided `DataProviderBounds`.
31//!
32//! ### Example
33//!
34//! ```ignore
35//! use pallet_staking::election_size_tracker::*;
36//!
37//! // instantiates a new tracker.
38//! let mut size_tracker = StaticTracker::<Staking>::default();
39//!
40//! let voter_bounds = ElectionBoundsBuilder::default().voter_size(1_00.into()).build().voters;
41//!
42//! let mut sorted_voters = T::VoterList.iter();
43//! let mut selected_voters = vec![];
44//!
45//! // fit as many voters in the vec as the bounds permit.
46//! for v in sorted_voters {
47//!     let voter = (v, weight_of(&v), targets_of(&v));
48//!     if size_tracker.try_register_voter(&voter, &voter_bounds).is_err() {
49//!         // voter bounds size exhausted
50//!         break;
51//!     }
52//!     selected_voters.push(voter);
53//! }
54//!
55//! // The SCALE encoded size in bytes of `selected_voters` is guaranteed to be below
56//! // `voter_bounds`.
57//! debug_assert!(
58//!     selected_voters.encoded_size() <=
59//!     SizeTracker::<Staking>::final_byte_size_of(size_tracker.num_voters, size_tracker.size)
60//! );
61//! ```
62//!
63//! ### Implementation Details
64//!
65//! The current implementation of the static tracker is tightly coupled with the staking pallet
66//! implementation, namely the representation of a voter ([`VoterOf`]). The SCALE encoded byte size
67//! is calculated using [`Encode::size_hint`] of each type in the voter tuple. Each voter's byte
68//! size is the sum of:
69//! - 1 * [`Encode::size_hint`] of the `AccountId` type;
70//! - 1 * [`Encode::size_hint`] of the `VoteWeight` type;
71//! - `num_votes` * [`Encode::size_hint`] of the `AccountId` type.
72
73use codec::Encode;
74use frame_election_provider_support::{
75	bounds::{DataProviderBounds, SizeBound},
76	ElectionDataProvider, VoterOf,
77};
78
79/// Keeps track of the SCALE encoded byte length of the snapshot's voters or targets.
80///
81/// The tracker calculates the bytes used based on static rules, without requiring any actual
82/// encoding or extra allocations.
83#[derive(Clone, Copy, Debug)]
84pub struct StaticTracker<DataProvider> {
85	pub size: usize,
86	pub counter: usize,
87	_marker: core::marker::PhantomData<DataProvider>,
88}
89
90impl<DataProvider> Default for StaticTracker<DataProvider> {
91	fn default() -> Self {
92		Self { size: 0, counter: 0, _marker: Default::default() }
93	}
94}
95
96impl<DataProvider> StaticTracker<DataProvider>
97where
98	DataProvider: ElectionDataProvider,
99{
100	/// Tries to register a new voter.
101	///
102	/// If the new voter exhausts the provided bounds, return an error. Otherwise, the internal
103	/// state of the tracker is updated with the new registered voter.
104	pub fn try_register_voter(
105		&mut self,
106		voter: &VoterOf<DataProvider>,
107		bounds: &DataProviderBounds,
108	) -> Result<(), ()> {
109		let tracker_size_after = {
110			let voter_hint = Self::voter_size_hint(voter);
111			Self::final_byte_size_of(self.counter + 1, self.size.saturating_add(voter_hint))
112		};
113
114		match bounds.size_exhausted(SizeBound(tracker_size_after as u32)) {
115			true => Err(()),
116			false => {
117				self.size = tracker_size_after;
118				self.counter += 1;
119				Ok(())
120			},
121		}
122	}
123
124	/// Calculates the size of the voter to register based on [`Encode::size_hint`].
125	fn voter_size_hint(voter: &VoterOf<DataProvider>) -> usize {
126		let (voter_account, vote_weight, targets) = voter;
127
128		voter_account
129			.size_hint()
130			.saturating_add(vote_weight.size_hint())
131			.saturating_add(voter_account.size_hint().saturating_mul(targets.len()))
132	}
133
134	/// Tries to register a new target.
135	///
136	/// If the new target exhausts the provided bounds, return an error. Otherwise, the internal
137	/// state of the tracker is updated with the new registered target.
138	pub fn try_register_target(
139		&mut self,
140		target: DataProvider::AccountId,
141		bounds: &DataProviderBounds,
142	) -> Result<(), ()> {
143		let tracker_size_after = Self::final_byte_size_of(
144			self.counter + 1,
145			self.size.saturating_add(target.size_hint()),
146		);
147
148		match bounds.size_exhausted(SizeBound(tracker_size_after as u32)) {
149			true => Err(()),
150			false => {
151				self.size = tracker_size_after;
152				self.counter += 1;
153				Ok(())
154			},
155		}
156	}
157
158	/// Size of the SCALE encoded prefix with a given length.
159	#[inline]
160	fn length_prefix(len: usize) -> usize {
161		use codec::{Compact, CompactLen};
162		Compact::<u32>::compact_len(&(len as u32))
163	}
164
165	/// Calculates the final size in bytes of the SCALE encoded snapshot voter struct.
166	fn final_byte_size_of(num_voters: usize, size: usize) -> usize {
167		Self::length_prefix(num_voters).saturating_add(size)
168	}
169}
170
171#[cfg(test)]
172mod tests {
173	use super::*;
174	use crate::{
175		mock::{AccountId, Staking, Test},
176		BoundedVec, MaxNominationsOf,
177	};
178	use frame_election_provider_support::bounds::ElectionBoundsBuilder;
179	use sp_core::bounded_vec;
180
181	type Voters = BoundedVec<AccountId, MaxNominationsOf<Test>>;
182
183	#[test]
184	pub fn election_size_tracker_works() {
185		let mut voters: Vec<(u64, u64, Voters)> = vec![];
186		let mut size_tracker = StaticTracker::<Staking>::default();
187		let voter_bounds = ElectionBoundsBuilder::default().voters_size(1_50.into()).build().voters;
188
189		// register 1 voter with 1 vote.
190		let voter = (1, 10, bounded_vec![2]);
191		assert!(size_tracker.try_register_voter(&voter, &voter_bounds).is_ok());
192		voters.push(voter);
193
194		assert_eq!(
195			StaticTracker::<Staking>::final_byte_size_of(size_tracker.counter, size_tracker.size),
196			voters.encoded_size()
197		);
198
199		// register another voter, now with 3 votes.
200		let voter = (2, 20, bounded_vec![3, 4, 5]);
201		assert!(size_tracker.try_register_voter(&voter, &voter_bounds).is_ok());
202		voters.push(voter);
203
204		assert_eq!(
205			StaticTracker::<Staking>::final_byte_size_of(size_tracker.counter, size_tracker.size),
206			voters.encoded_size()
207		);
208
209		// register noop vote (unlikely to happen).
210		let voter = (3, 30, bounded_vec![]);
211		assert!(size_tracker.try_register_voter(&voter, &voter_bounds).is_ok());
212		voters.push(voter);
213
214		assert_eq!(
215			StaticTracker::<Staking>::final_byte_size_of(size_tracker.counter, size_tracker.size),
216			voters.encoded_size()
217		);
218	}
219
220	#[test]
221	pub fn election_size_tracker_bounds_works() {
222		let mut voters: Vec<(u64, u64, Voters)> = vec![];
223		let mut size_tracker = StaticTracker::<Staking>::default();
224		let voter_bounds = ElectionBoundsBuilder::default().voters_size(1_00.into()).build().voters;
225
226		let voter = (1, 10, bounded_vec![2]);
227		assert!(size_tracker.try_register_voter(&voter, &voter_bounds).is_ok());
228		voters.push(voter);
229
230		assert_eq!(
231			StaticTracker::<Staking>::final_byte_size_of(size_tracker.counter, size_tracker.size),
232			voters.encoded_size()
233		);
234
235		assert!(size_tracker.size > 0 && size_tracker.size < 1_00);
236		let size_before_overflow = size_tracker.size;
237
238		// try many voters that will overflow the tracker's buffer.
239		let voter = (2, 10, bounded_vec![2, 3, 4, 5, 6, 7, 8, 9]);
240		voters.push(voter.clone());
241
242		assert!(size_tracker.try_register_voter(&voter, &voter_bounds).is_err());
243		assert!(size_tracker.size > 0 && size_tracker.size < 1_00);
244
245		// size of the tracker did not update when trying to register votes failed.
246		assert_eq!(size_tracker.size, size_before_overflow);
247	}
248
249	#[test]
250	fn len_prefix_works() {
251		let length_samples =
252			vec![0usize, 1, 62, 63, 64, 16383, 16384, 16385, 1073741822, 1073741823, 1073741824];
253
254		for s in length_samples {
255			// the encoded size of a vector of n bytes should be n + the length prefix
256			assert_eq!(vec![1u8; s].encoded_size(), StaticTracker::<Staking>::length_prefix(s) + s);
257		}
258	}
259}
pallet_staking/election_size_tracker.rs

pallet_staking/
election_size_tracker.rs
