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

@@ -26,6 +26,9 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
+extern crate self as sp_weights;
+
+mod weight_meter;
 mod weight_v2;
 
 use codec::{CompactAs, Decode, Encode, MaxEncodedLen};
@@ -40,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,176 @@
+// 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 still be 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 consumed_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 the given weight after checking that it can be consumed. Otherwise do nothing.
+ 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));
+
+ assert!(meter.check_accrue(Weight::from_parts(5, 0)));
+ assert_eq!(meter.consumed, Weight::from_parts(5, 0));
+ assert_eq!(meter.remaining(), Weight::from_parts(5, 20));
+
+ assert!(meter.check_accrue(Weight::from_parts(2, 10)));
+ assert_eq!(meter.consumed, Weight::from_parts(7, 10));
+ assert_eq!(meter.remaining(), Weight::from_parts(3, 10));
+
+ assert!(meter.check_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_can_accrue_works() {
+ let meter = WeightMeter::from_limit(Weight::from_parts(1, 1));
+
+ assert!(meter.can_accrue(Weight::from_parts(0, 0)));
+ assert!(meter.can_accrue(Weight::from_parts(1, 1)));
+ assert!(!meter.can_accrue(Weight::from_parts(0, 2)));
+ assert!(!meter.can_accrue(Weight::from_parts(2, 0)));
+ assert!(!meter.can_accrue(Weight::from_parts(2, 2)));
+ }
+
+ #[test]
+ fn weight_meter_check_accrue_works() {
+ let mut meter = WeightMeter::from_limit(Weight::from_parts(2, 2));
+
+ assert!(meter.check_accrue(Weight::from_parts(0, 0)));
+ assert!(meter.check_accrue(Weight::from_parts(1, 1)));
+ assert!(!meter.check_accrue(Weight::from_parts(0, 2)));
+ assert!(!meter.check_accrue(Weight::from_parts(2, 0)));
+ assert!(!meter.check_accrue(Weight::from_parts(2, 2)));
+ assert!(meter.check_accrue(Weight::from_parts(0, 1)));
+ assert!(meter.check_accrue(Weight::from_parts(1, 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]
+ fn consumed_ratio_works() {
+ let mut meter = WeightMeter::from_limit(Weight::from_parts(10, 20));
+
+ assert!(meter.check_accrue(Weight::from_parts(5, 0)));
+ assert_eq!(meter.consumed_ratio(), Perbill::from_percent(50));
+ assert!(meter.check_accrue(Weight::from_parts(0, 12)));
+ assert_eq!(meter.consumed_ratio(), Perbill::from_percent(60));
+
+ assert!(meter.check_accrue(Weight::from_parts(2, 0)));
+ assert_eq!(meter.consumed_ratio(), Perbill::from_percent(70));
+ assert!(meter.check_accrue(Weight::from_parts(0, 4)));
+ assert_eq!(meter.consumed_ratio(), Perbill::from_percent(80));
+
+ assert!(meter.check_accrue(Weight::from_parts(3, 0)));
+ assert_eq!(meter.consumed_ratio(), Perbill::from_percent(100));
+ assert!(meter.check_accrue(Weight::from_parts(0, 4)));
+ assert_eq!(meter.consumed_ratio(), Perbill::from_percent(100));
+ }
+}