Move WeightCounter to sp-weights

frame/scheduler/src/benchmarking.rs

@@ -122,17 +122,13 @@ fn make_origin<T: Config>(signed: bool) -> <T as Config>::PalletsOrigin {
 	}
 }
 
-fn dummy_counter() -> WeightCounter {
-	WeightCounter { used: Weight::zero(), limit: Weight::MAX }
-}
-
 benchmarks! {
 	// `service_agendas` when no work is done.
 	service_agendas_base {
 		let now = T::BlockNumber::from(BLOCK_NUMBER);
 		IncompleteSince::<T>::put(now - One::one());
 	}: {
-		Scheduler::<T>::service_agendas(&mut dummy_counter(), now, 0);
+		Scheduler::<T>::service_agendas(&mut WeightMeter::max_limit(), now, 0);
 	} verify {
 		assert_eq!(IncompleteSince::<T>::get(), Some(now - One::one()));
 	}
@@ -144,7 +140,7 @@ benchmarks! {
 		fill_schedule::<T>(now, s)?;
 		let mut executed = 0;
 	}: {
-		Scheduler::<T>::service_agenda(&mut dummy_counter(), &mut executed, now, now, 0);
+		Scheduler::<T>::service_agenda(&mut WeightMeter::max_limit(), &mut executed, now, now, 0);
 	} verify {
 		assert_eq!(executed, 0);
 	}
@@ -155,7 +151,7 @@ benchmarks! {
 		let now = BLOCK_NUMBER.into();
 		let task = make_task::<T>(false, false, false, None, 0);
 		// prevent any tasks from actually being executed as we only want the surrounding weight.
-		let mut counter = WeightCounter { used: Weight::zero(), limit: Weight::zero() };
+		let mut counter = WeightMeter::from_limit(Weight::zero());
 	}: {
 		let result = Scheduler::<T>::service_task(&mut counter, now, now, 0, true, task);
 	} verify {
@@ -169,7 +165,7 @@ benchmarks! {
 		let now = BLOCK_NUMBER.into();
 		let task = make_task::<T>(false, false, false, Some(s), 0);
 		// prevent any tasks from actually being executed as we only want the surrounding weight.
-		let mut counter = WeightCounter { used: Weight::zero(), limit: Weight::zero() };
+		let mut counter = WeightMeter::from_limit(Weight::zero());
 	}: {
 		let result = Scheduler::<T>::service_task(&mut counter, now, now, 0, true, task);
 	} verify {
@@ -181,7 +177,7 @@ benchmarks! {
 		let now = BLOCK_NUMBER.into();
 		let task = make_task::<T>(false, true, false, None, 0);
 		// prevent any tasks from actually being executed as we only want the surrounding weight.
-		let mut counter = WeightCounter { used: Weight::zero(), limit: Weight::zero() };
+		let mut counter = WeightMeter::from_limit(Weight::zero());
 	}: {
 		let result = Scheduler::<T>::service_task(&mut counter, now, now, 0, true, task);
 	} verify {
@@ -193,15 +189,15 @@ benchmarks! {
 		let now = BLOCK_NUMBER.into();
 		let task = make_task::<T>(true, false, false, None, 0);
 		// prevent any tasks from actually being executed as we only want the surrounding weight.
-		let mut counter = WeightCounter { used: Weight::zero(), limit: Weight::zero() };
+		let mut counter = WeightMeter::from_limit(Weight::zero());
 	}: {
 		let result = Scheduler::<T>::service_task(&mut counter, now, now, 0, true, task);
 	} verify {
 	}
 
 	// `execute_dispatch` when the origin is `Signed`, not counting the dispatable's weight.
 	execute_dispatch_signed {
-		let mut counter = WeightCounter { used: Weight::zero(), limit: Weight::MAX };
+		let mut counter = WeightMeter::max_limit();
 		let origin = make_origin::<T>(true);
 		let call = T::Preimages::realize(&make_call::<T>(None)).unwrap().0;
 	}: {
@@ -212,7 +208,7 @@ benchmarks! {
 
 	// `execute_dispatch` when the origin is not `Signed`, not counting the dispatable's weight.
 	execute_dispatch_unsigned {
-		let mut counter = WeightCounter { used: Weight::zero(), limit: Weight::MAX };
+		let mut counter = WeightMeter::max_limit();
 		let origin = make_origin::<T>(false);
 		let call = T::Preimages::realize(&make_call::<T>(None)).unwrap().0;
 	}: {

frame/scheduler/src/lib.rs

@@ -70,7 +70,7 @@ use frame_support::{
 		Bounded, CallerTrait, EnsureOrigin, Get, Hash as PreimageHash, IsType, OriginTrait,
 		PalletInfoAccess, PrivilegeCmp, QueryPreimage, StorageVersion, StorePreimage,
 	},
-	weights::Weight,
+	weights::{Weight, WeightMeter},
 };
 use frame_system::{self as system};
 pub use pallet::*;
@@ -143,25 +143,6 @@ pub type ScheduledOf<T> = Scheduled<
 	<T as frame_system::Config>::AccountId,
 >;
 
-struct WeightCounter {
-	used: Weight,
-	limit: Weight,
-}
-impl WeightCounter {
-	fn check_accrue(&mut self, w: Weight) -> bool {
-		let test = self.used.saturating_add(w);
-		if test.any_gt(self.limit) {
-			false
-		} else {
-			self.used = test;
-			true
-		}
-	}
-	fn can_accrue(&mut self, w: Weight) -> bool {
-		self.used.saturating_add(w).all_lte(self.limit)
-	}
-}
-
 pub(crate) trait MarginalWeightInfo: WeightInfo {
 	fn service_task(maybe_lookup_len: Option<usize>, named: bool, periodic: bool) -> Weight {
 		let base = Self::service_task_base();
@@ -306,10 +287,9 @@ pub mod pallet {
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		/// Execute the scheduled calls
 		fn on_initialize(now: T::BlockNumber) -> Weight {
-			let mut weight_counter =
-				WeightCounter { used: Weight::zero(), limit: T::MaximumWeight::get() };
+			let mut weight_counter = WeightMeter::from_limit(T::MaximumWeight::get());
 			Self::service_agendas(&mut weight_counter, now, u32::max_value());
-			weight_counter.used
+			weight_counter.consumed
 		}
 	}
 
@@ -933,7 +913,7 @@ use ServiceTaskError::*;
 
 impl<T: Config> Pallet<T> {
 	/// Service up to `max` agendas queue starting from earliest incompletely executed agenda.
-	fn service_agendas(weight: &mut WeightCounter, now: T::BlockNumber, max: u32) {
+	fn service_agendas(weight: &mut WeightMeter, now: T::BlockNumber, max: u32) {
 		if !weight.check_accrue(T::WeightInfo::service_agendas_base()) {
 			return
 		}
@@ -961,7 +941,7 @@ impl<T: Config> Pallet<T> {
 	/// Returns `true` if the agenda was fully completed, `false` if it should be revisited at a
 	/// later block.
 	fn service_agenda(
-		weight: &mut WeightCounter,
+		weight: &mut WeightMeter,
 		executed: &mut u32,
 		now: T::BlockNumber,
 		when: T::BlockNumber,
@@ -1030,7 +1010,7 @@ impl<T: Config> Pallet<T> {
 	/// - realizing the task's call which can include a preimage lookup.
 	/// - Rescheduling the task for execution in a later agenda if periodic.
 	fn service_task(
-		weight: &mut WeightCounter,
+		weight: &mut WeightMeter,
 		now: T::BlockNumber,
 		when: T::BlockNumber,
 		agenda_index: u32,
@@ -1110,7 +1090,7 @@ impl<T: Config> Pallet<T> {
 	/// NOTE: Only the weight for this function will be counted (origin lookup, dispatch and the
 	/// call itself).
 	fn execute_dispatch(
-		weight: &mut WeightCounter,
+		weight: &mut WeightMeter,
 		origin: T::PalletsOrigin,
 		call: <T as Config>::RuntimeCall,
 	) -> Result<DispatchResult, ServiceTaskError> {

primitives/weights/src/lib.rs

@@ -28,6 +28,7 @@
 
 extern crate self as sp_weights;
 
+mod weight_meter;
 mod weight_v2;
 
 use codec::{CompactAs, Decode, Encode, MaxEncodedLen};
@@ -42,6 +43,7 @@ use sp_arithmetic::{
 use sp_core::Get;
 use sp_debug_derive::RuntimeDebug;
 
+pub use weight_meter::*;
 pub use weight_v2::*;
 
 pub mod constants {

primitives/weights/src/weight_meter.rs

@@ -0,0 +1,167 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Contains the `WeightMeter` primitive to meter weight usage.
+
+use super::Weight;
+
+use sp_arithmetic::Perbill;
+
+/// Meters consumed weight and a hard limit for the maximal consumable weight.
+///
+/// Can be used to check if enough weight for an operation is available before committing to it.
+///
+/// # Example
+///
+/// ```rust
+/// use sp_weights::{Weight, WeightMeter};
+///
+/// // The weight is limited to (10, 0).
+/// let mut meter = WeightMeter::from_limit(Weight::from_parts(10, 0));
+/// // There is enough weight remaining for an operation with (5, 0) weight.
+/// assert!(meter.check_accrue(Weight::from_parts(5, 0)));
+/// // There is not enough weight remaining for an operation with (6, 0) weight.
+/// assert!(!meter.check_accrue(Weight::from_parts(6, 0)));
+/// ```
+#[derive(Debug, Clone)]
+pub struct WeightMeter {
+	/// The already consumed weight.
+	pub consumed: Weight,
+
+	/// The maximal consumable weight.
+	pub limit: Weight,
+}
+
+impl WeightMeter {
+	/// Creates [`Self`] from a limit for the maximal consumable weight.
+	pub fn from_limit(limit: Weight) -> Self {
+		Self { consumed: Weight::zero(), limit }
+	}
+
+	/// Creates [`Self`] with the maximal possible limit for the consumable weight.
+	pub fn max_limit() -> Self {
+		Self::from_limit(Weight::MAX)
+	}
+
+	/// The remaining weight that can be still consumed.
+	pub fn remaining(&self) -> Weight {
+		self.limit.saturating_sub(self.consumed)
+	}
+
+	/// The ratio of consumed weight to the limit.
+	///
+	/// Calculates one ratio per component and returns the largest.
+	pub fn remaining_ratio(&self) -> Perbill {
+		let time = Perbill::from_rational(self.consumed.ref_time(), self.limit.ref_time());
+		let pov = Perbill::from_rational(self.consumed.proof_size(), self.limit.proof_size());
+		time.max(pov)
+	}
+
+	/// Consume some weight and defensively fail if it is over the limit or would saturate the
+	/// arithmetic type.
+	pub fn defensive_accrue(&mut self, w: Weight) {
+		let res = self.check_accrue(w);
+		debug_assert!(res, "Weight meter exhausted; {} + {} > {}", &self.consumed, &w, &self.limit);
+	}
+
+	/// Check if the given weight can be consumed. Do nothing if not.
+	pub fn check_accrue(&mut self, w: Weight) -> bool {
+		self.consumed.checked_add(&w).map_or(false, |test| {
+			if test.any_gt(self.limit) {
+				false
+			} else {
+				self.consumed = test;
+				true
+			}
+		})
+	}
+
+	/// Check if the given weight can be consumed.
+	pub fn can_accrue(&self, w: Weight) -> bool {
+		self.consumed.checked_add(&w).map_or(false, |t| t.all_lte(self.limit))
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use crate::*;
+
+	#[test]
+	fn weight_meter_remaining_works() {
+		let mut meter = WeightMeter::from_limit(Weight::from_parts(10, 20));
+
+		meter.defensive_accrue(Weight::from_parts(5, 0));
+		assert_eq!(meter.consumed, Weight::from_parts(5, 0));
+		assert_eq!(meter.remaining(), Weight::from_parts(5, 20));
+
+		meter.defensive_accrue(Weight::from_parts(2, 10));
+		assert_eq!(meter.consumed, Weight::from_parts(7, 10));
+		assert_eq!(meter.remaining(), Weight::from_parts(3, 10));
+
+		meter.defensive_accrue(Weight::from_parts(3, 10));
+		assert_eq!(meter.consumed, Weight::from_parts(10, 20));
+		assert_eq!(meter.remaining(), Weight::from_parts(0, 0));
+	}
+
+	#[test]
+	fn weight_meter_check_and_can_accrue_works() {
+		let mut meter = WeightMeter::max_limit();
+
+		assert!(meter.can_accrue(Weight::from_parts(u64::MAX, 0)));
+		assert!(meter.check_accrue(Weight::from_parts(u64::MAX, 0)));
+
+		assert!(meter.can_accrue(Weight::from_parts(0, u64::MAX)));
+		assert!(meter.check_accrue(Weight::from_parts(0, u64::MAX)));
+
+		assert!(!meter.can_accrue(Weight::from_parts(0, 1)));
+		assert!(!meter.check_accrue(Weight::from_parts(0, 1)));
+
+		assert!(!meter.can_accrue(Weight::from_parts(1, 0)));
+		assert!(!meter.check_accrue(Weight::from_parts(1, 0)));
+
+		assert!(meter.can_accrue(Weight::zero()));
+		assert!(meter.check_accrue(Weight::zero()));
+	}
+
+	#[test]
+	#[should_panic(expected = "Weight meter exhausted")]
+	fn weight_meter_defensive_accrue_works() {
+		let mut meter = WeightMeter::from_limit(Weight::from_parts(10, 0));
+
+		meter.defensive_accrue(Weight::from_parts(11, 0));
+	}
+
+	#[test]
+	fn remaining_ratio_works() {
+		let mut meter = WeightMeter::from_limit(Weight::from_parts(10, 20));
+
+		meter.defensive_accrue(Weight::from_parts(5, 0));
+		assert_eq!(meter.remaining_ratio(), Perbill::from_percent(50));
+		meter.defensive_accrue(Weight::from_parts(0, 12));
+		assert_eq!(meter.remaining_ratio(), Perbill::from_percent(60));
+
+		meter.defensive_accrue(Weight::from_parts(2, 0));
+		assert_eq!(meter.remaining_ratio(), Perbill::from_percent(70));
+		meter.defensive_accrue(Weight::from_parts(0, 4));
+		assert_eq!(meter.remaining_ratio(), Perbill::from_percent(80));
+
+		meter.defensive_accrue(Weight::from_parts(3, 0));
+		assert_eq!(meter.remaining_ratio(), Perbill::from_percent(100));
+		meter.defensive_accrue(Weight::from_parts(0, 4));
+		assert_eq!(meter.remaining_ratio(), Perbill::from_percent(100));
+	}
+}