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//! Mechanism to reclaim PoV proof size weight after an extrinsic has been applied.
17
18#![cfg_attr(not(feature = "std"), no_std)]
19
20use codec::{Decode, DecodeWithMemTracking, Encode};
21use core::marker::PhantomData;
22use cumulus_primitives_core::Weight;
23use cumulus_primitives_proof_size_hostfunction::{
24	storage_proof_size::storage_proof_size, PROOF_RECORDING_DISABLED,
25};
26use frame_support::{
27	dispatch::{DispatchInfo, PostDispatchInfo},
28	weights::WeightMeter,
29};
30use frame_system::Config;
31use scale_info::TypeInfo;
32use sp_runtime::{
33	impl_tx_ext_default,
34	traits::{DispatchInfoOf, Dispatchable, PostDispatchInfoOf, TransactionExtension},
35	transaction_validity::TransactionValidityError,
36	DispatchResult,
37};
38
39#[cfg(test)]
40mod tests;
41
42const LOG_TARGET: &'static str = "runtime::storage_reclaim";
43
44/// `StorageWeightReclaimer` is a mechanism for manually reclaiming storage weight.
45///
46/// It internally keeps track of the proof size and storage weight at initialization time. At
47/// reclaim  it computes the real consumed storage weight and refunds excess weight.
48///
49/// # Example
50#[doc = docify::embed!("src/tests.rs", simple_reclaimer_example)]
51pub struct StorageWeightReclaimer {
52	previous_remaining_proof_size: u64,
53	previous_reported_proof_size: Option<u64>,
54}
55
56impl StorageWeightReclaimer {
57	/// Creates a new `StorageWeightReclaimer` instance and initializes it with the storage
58	/// size provided by `weight_meter` and reported proof size from the node.
59	#[must_use = "Must call `reclaim_with_meter` to reclaim the weight"]
60	pub fn new(weight_meter: &WeightMeter) -> StorageWeightReclaimer {
61		let previous_remaining_proof_size = weight_meter.remaining().proof_size();
62		let previous_reported_proof_size = get_proof_size();
63		Self { previous_remaining_proof_size, previous_reported_proof_size }
64	}
65
66	/// Check the consumed storage weight and calculate the consumed excess weight.
67	fn reclaim(&mut self, remaining_weight: Weight) -> Option<Weight> {
68		let current_remaining_weight = remaining_weight.proof_size();
69		let current_storage_proof_size = get_proof_size()?;
70		let previous_storage_proof_size = self.previous_reported_proof_size?;
71		let used_weight =
72			self.previous_remaining_proof_size.saturating_sub(current_remaining_weight);
73		let reported_used_size =
74			current_storage_proof_size.saturating_sub(previous_storage_proof_size);
75		let reclaimable = used_weight.saturating_sub(reported_used_size);
76		log::trace!(
77			target: LOG_TARGET,
78			"Found reclaimable storage weight. benchmarked: {used_weight}, consumed: {reported_used_size}"
79		);
80
81		self.previous_remaining_proof_size = current_remaining_weight.saturating_add(reclaimable);
82		self.previous_reported_proof_size = Some(current_storage_proof_size);
83		Some(Weight::from_parts(0, reclaimable))
84	}
85
86	/// Check the consumed storage weight and add the reclaimed
87	/// weight budget back to `weight_meter`.
88	pub fn reclaim_with_meter(&mut self, weight_meter: &mut WeightMeter) -> Option<Weight> {
89		let reclaimed = self.reclaim(weight_meter.remaining())?;
90		weight_meter.reclaim_proof_size(reclaimed.proof_size());
91		Some(reclaimed)
92	}
93}
94
95/// Returns the current storage proof size from the host side.
96///
97/// Returns `None` if proof recording is disabled on the host.
98pub fn get_proof_size() -> Option<u64> {
99	let proof_size = storage_proof_size();
100	(proof_size != PROOF_RECORDING_DISABLED).then_some(proof_size)
101}
102
103// Encapsulate into a mod so that macro generated code doesn't trigger a warning about deprecated
104// usage.
105#[allow(deprecated)]
106mod allow_deprecated {
107	use super::*;
108
109	/// Storage weight reclaim mechanism.
110	///
111	/// This extension checks the size of the node-side storage proof
112	/// before and after executing a given extrinsic. The difference between
113	/// benchmarked and spent weight can be reclaimed.
114	#[deprecated(note = "This extension doesn't provide accurate reclaim for storage intensive \
115		transaction extension pipeline; it ignores the validation and preparation of extensions prior \
116		to itself and ignores the post dispatch logic for extensions subsequent to itself, it also
117		doesn't provide weight information. \
118		Use `StorageWeightReclaim` in the `cumulus-pallet-weight-reclaim` crate")]
119	#[derive(Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, Default, TypeInfo)]
120	#[scale_info(skip_type_params(T))]
121	pub struct StorageWeightReclaim<T: Config + Send + Sync>(pub(super) PhantomData<T>);
122}
123#[allow(deprecated)]
124pub use allow_deprecated::StorageWeightReclaim;
125
126#[allow(deprecated)]
127impl<T: Config + Send + Sync> StorageWeightReclaim<T> {
128	/// Create a new `StorageWeightReclaim` instance.
129	pub fn new() -> Self {
130		Self(Default::default())
131	}
132}
133
134#[allow(deprecated)]
135impl<T: Config + Send + Sync> core::fmt::Debug for StorageWeightReclaim<T> {
136	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
137		let _ = write!(f, "StorageWeightReclaim");
138		Ok(())
139	}
140}
141
142#[allow(deprecated)]
143impl<T: Config + Send + Sync> TransactionExtension<T::RuntimeCall> for StorageWeightReclaim<T>
144where
145	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
146{
147	const IDENTIFIER: &'static str = "StorageWeightReclaim";
148	type Implicit = ();
149	type Val = ();
150	type Pre = Option<u64>;
151
152	fn prepare(
153		self,
154		_val: Self::Val,
155		_origin: &T::RuntimeOrigin,
156		_call: &T::RuntimeCall,
157		_info: &DispatchInfoOf<T::RuntimeCall>,
158		_len: usize,
159	) -> Result<Self::Pre, TransactionValidityError> {
160		Ok(get_proof_size())
161	}
162
163	fn post_dispatch_details(
164		pre: Self::Pre,
165		info: &DispatchInfoOf<T::RuntimeCall>,
166		post_info: &PostDispatchInfoOf<T::RuntimeCall>,
167		_len: usize,
168		_result: &DispatchResult,
169	) -> Result<Weight, TransactionValidityError> {
170		let Some(pre_dispatch_proof_size) = pre else {
171			return Ok(Weight::zero());
172		};
173
174		let Some(post_dispatch_proof_size) = get_proof_size() else {
175			log::debug!(
176				target: LOG_TARGET,
177				"Proof recording enabled during pre-dispatch, now disabled. This should not happen."
178			);
179			return Ok(Weight::zero())
180		};
181		// Unspent weight according to the `actual_weight` from `PostDispatchInfo`
182		// This unspent weight will be refunded by the `CheckWeight` extension, so we need to
183		// account for that.
184		let unspent = post_info.calc_unspent(info).proof_size();
185		let benchmarked_weight = info.total_weight().proof_size().saturating_sub(unspent);
186		let consumed_weight = post_dispatch_proof_size.saturating_sub(pre_dispatch_proof_size);
187
188		let storage_size_diff = benchmarked_weight.abs_diff(consumed_weight as u64);
189
190		let extrinsic_len = frame_system::AllExtrinsicsLen::<T>::get().unwrap_or(0);
191		let node_side_pov_size = post_dispatch_proof_size.saturating_add(extrinsic_len.into());
192
193		// This value will be reclaimed by [`frame_system::CheckWeight`], so we need to calculate
194		// that in.
195		frame_system::BlockWeight::<T>::mutate(|current| {
196			if consumed_weight > benchmarked_weight {
197				log::error!(
198					target: LOG_TARGET,
199					"Benchmarked storage weight smaller than consumed storage weight. extrinsic: {} benchmarked: {benchmarked_weight} consumed: {consumed_weight} unspent: {unspent}",
200					frame_system::Pallet::<T>::extrinsic_index().unwrap_or(0)
201				);
202				current.accrue(Weight::from_parts(0, storage_size_diff), info.class)
203			} else {
204				log::trace!(
205					target: LOG_TARGET,
206					"Reclaiming storage weight. extrinsic: {} benchmarked: {benchmarked_weight} consumed: {consumed_weight} unspent: {unspent}",
207					frame_system::Pallet::<T>::extrinsic_index().unwrap_or(0)
208				);
209				current.reduce(Weight::from_parts(0, storage_size_diff), info.class)
210			}
211
212			// If we encounter a situation where the node-side proof size is already higher than
213			// what we have in the runtime bookkeeping, we add the difference to the `BlockWeight`.
214			// This prevents that the proof size grows faster than the runtime proof size.
215			let block_weight_proof_size = current.total().proof_size();
216			let missing_from_node = node_side_pov_size.saturating_sub(block_weight_proof_size);
217			if missing_from_node > 0 {
218				log::debug!(
219					target: LOG_TARGET,
220					"Node-side PoV size higher than runtime proof size weight. node-side: {node_side_pov_size} extrinsic_len: {extrinsic_len} runtime: {block_weight_proof_size}, missing: {missing_from_node}. Setting to node-side proof size."
221				);
222				current.accrue(Weight::from_parts(0, missing_from_node), info.class);
223			}
224		});
225		Ok(Weight::zero())
226	}
227
228	impl_tx_ext_default!(T::RuntimeCall; weight validate);
229}
cumulus_primitives_storage_weight_reclaim/lib.rs

cumulus_primitives_storage_weight_reclaim/
lib.rs
