cumulus_client_proof_size_recording/
lib.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// This file is part of Cumulus.
3// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
4
5// Cumulus is free software: you can redistribute it and/or modify
6// it under the terms of the GNU General Public License as published by
7// the Free Software Foundation, either version 3 of the License, or
8// (at your option) any later version.
9
10// Cumulus is distributed in the hope that it will be useful,
11// but WITHOUT ANY WARRANTY; without even the implied warranty of
12// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13// GNU General Public License for more details.
14
15// You should have received a copy of the GNU General Public License
16// along with Cumulus. If not, see <https://www.gnu.org/licenses/>.
17
18//! Proof size recording utilities.
19
20use codec::{Decode, Encode};
21use sc_client_api::{
22	backend::AuxStore,
23	client::{AuxDataOperations, FinalityNotification, PreCommitActions},
24};
25use sp_blockchain::{Error as ClientError, Result as ClientResult};
26use sp_runtime::traits::Block as BlockT;
27use sp_trie::proof_size_extension::RecordedProofSizeEstimations;
28use std::sync::Arc;
29
30const PROOF_SIZE_RECORDING_VERSION: &[u8] = b"cumulus_proof_size_recording_version";
31const PROOF_SIZE_RECORDING_CURRENT_VERSION: u32 = 1;
32
33/// The aux storage key used to store the proof size recordings for the given block hash.
34fn proof_size_recording_key<H: Encode>(block_hash: H) -> Vec<u8> {
35	(b"cumulus_proof_size_recording", block_hash).encode()
36}
37
38fn load_decode<B, T>(backend: &B, key: &[u8]) -> ClientResult<Option<T>>
39where
40	B: AuxStore,
41	T: Decode,
42{
43	let corrupt = |e: codec::Error| {
44		ClientError::Backend(format!("Proof size recording DB is corrupted. Decode error: {}", e))
45	};
46	match backend.get_aux(key)? {
47		None => Ok(None),
48		Some(t) => T::decode(&mut &t[..]).map(Some).map_err(corrupt),
49	}
50}
51
52/// Prepare aux storage key-value pairs for persisting proof size recordings.
53///
54/// Returns the key-value pairs that need to be written to the aux storage.
55pub fn prepare_proof_size_recording_aux_data<H: Encode>(
56	block_hash: H,
57	recordings: Vec<u32>,
58) -> impl Iterator<Item = (Vec<u8>, Vec<u8>)> {
59	let current_version = PROOF_SIZE_RECORDING_CURRENT_VERSION.encode();
60	let key = proof_size_recording_key(block_hash);
61	let recordings = recordings.encode();
62
63	[(key, recordings), (PROOF_SIZE_RECORDING_VERSION.to_vec(), current_version)].into_iter()
64}
65
66/// Load the proof size recordings associated with a block.
67pub fn load_proof_size_recording<H: Encode, B: AuxStore>(
68	backend: &B,
69	block_hash: H,
70) -> ClientResult<Option<RecordedProofSizeEstimations>> {
71	let version = load_decode::<_, u32>(backend, PROOF_SIZE_RECORDING_VERSION)?;
72
73	match version {
74		None => Ok(None),
75		Some(PROOF_SIZE_RECORDING_CURRENT_VERSION) => {
76			load_decode::<_, Vec<u32>>(backend, proof_size_recording_key(block_hash).as_slice())
77				.map(|recordings| recordings.map(Into::into))
78		},
79		Some(other) => Err(ClientError::Backend(format!(
80			"Unsupported proof size recording DB version: {:?}",
81			other
82		))),
83	}
84}
85
86/// Cleanup auxiliary storage for finalized blocks.
87///
88/// This function removes proof size recordings for blocks that are no longer needed
89/// after finalization. It processes the finalized blocks and their stale heads to
90/// determine which recordings can be safely removed.
91fn aux_storage_cleanup<Block>(notification: &FinalityNotification<Block>) -> AuxDataOperations
92where
93	Block: BlockT,
94{
95	// Convert the hashes to deletion operations
96	notification
97		.stale_blocks
98		.iter()
99		.map(|b| (proof_size_recording_key(b.hash), None))
100		.collect()
101}
102
103/// Register a finality action for cleaning up proof size recordings.
104///
105/// This should be called during consensus initialization to automatically clean up
106/// proof size recordings when blocks are finalized.
107pub fn register_proof_size_recording_cleanup<C, Block>(client: Arc<C>)
108where
109	C: PreCommitActions<Block> + 'static,
110	Block: BlockT,
111{
112	let on_finality = move |notification: &FinalityNotification<Block>| -> AuxDataOperations {
113		aux_storage_cleanup(notification)
114	};
115
116	client.register_finality_action(Box::new(on_finality));
117}
cumulus_client_proof_size_recording/lib.rs

cumulus_client_proof_size_recording/
lib.rs
