cumulus_primitives_storage_weight_reclaim/
lib.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// SPDX-License-Identifier: Apache-2.0
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// 	http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Utilities to reclaim PoV proof size weight.
17//!
18//! Provides [`StorageWeightReclaimer`] to manually reclaim storage weight and [`get_proof_size`]
19//! to read the current storage proof size from the host.
20//!
21//! To reclaim storage weight after an extrinsic has been applied, use the `StorageWeightReclaim`
22//! transaction extension from the `cumulus-pallet-weight-reclaim` crate.
23
24#![cfg_attr(not(feature = "std"), no_std)]
25
26use cumulus_primitives_core::Weight;
27use cumulus_primitives_proof_size_hostfunction::{
28	storage_proof_size::storage_proof_size, PROOF_RECORDING_DISABLED,
29};
30use frame_support::weights::WeightMeter;
31
32#[cfg(test)]
33mod tests;
34
35const LOG_TARGET: &'static str = "runtime::storage_reclaim";
36
37/// `StorageWeightReclaimer` is a mechanism for manually reclaiming storage weight.
38///
39/// It internally keeps track of the proof size and storage weight at initialization time. At
40/// reclaim  it computes the real consumed storage weight and refunds excess weight.
41///
42/// # Example
43#[doc = docify::embed!("src/tests.rs", simple_reclaimer_example)]
44pub struct StorageWeightReclaimer {
45	previous_remaining_proof_size: u64,
46	previous_reported_proof_size: Option<u64>,
47}
48
49impl StorageWeightReclaimer {
50	/// Creates a new `StorageWeightReclaimer` instance and initializes it with the storage
51	/// size provided by `weight_meter` and reported proof size from the node.
52	#[must_use = "Must call `reclaim_with_meter` to reclaim the weight"]
53	pub fn new(weight_meter: &WeightMeter) -> StorageWeightReclaimer {
54		let previous_remaining_proof_size = weight_meter.remaining().proof_size();
55		let previous_reported_proof_size = get_proof_size();
56		Self { previous_remaining_proof_size, previous_reported_proof_size }
57	}
58
59	/// Check the consumed storage weight and calculate the consumed excess weight.
60	fn reclaim(&mut self, remaining_weight: Weight) -> Option<Weight> {
61		let current_remaining_weight = remaining_weight.proof_size();
62		let current_storage_proof_size = get_proof_size()?;
63		let previous_storage_proof_size = self.previous_reported_proof_size?;
64		let used_weight =
65			self.previous_remaining_proof_size.saturating_sub(current_remaining_weight);
66		let reported_used_size =
67			current_storage_proof_size.saturating_sub(previous_storage_proof_size);
68		let reclaimable = used_weight.saturating_sub(reported_used_size);
69		log::trace!(
70			target: LOG_TARGET,
71			"Found reclaimable storage weight. benchmarked: {used_weight}, consumed: {reported_used_size}"
72		);
73
74		self.previous_remaining_proof_size = current_remaining_weight.saturating_add(reclaimable);
75		self.previous_reported_proof_size = Some(current_storage_proof_size);
76		Some(Weight::from_parts(0, reclaimable))
77	}
78
79	/// Check the consumed storage weight and add the reclaimed
80	/// weight budget back to `weight_meter`.
81	pub fn reclaim_with_meter(&mut self, weight_meter: &mut WeightMeter) -> Option<Weight> {
82		let reclaimed = self.reclaim(weight_meter.remaining())?;
83		weight_meter.reclaim_proof_size(reclaimed.proof_size());
84		Some(reclaimed)
85	}
86}
87
88/// Returns the current storage proof size from the host side.
89///
90/// Returns `None` if proof recording is disabled on the host.
91pub fn get_proof_size() -> Option<u64> {
92	let proof_size = storage_proof_size();
93	(proof_size != PROOF_RECORDING_DISABLED).then_some(proof_size)
94}
cumulus_primitives_storage_weight_reclaim/lib.rs

cumulus_primitives_storage_weight_reclaim/
lib.rs
