1use codec::{Decode, Encode};
22use log::info;
23
24use crate::{migration::EpochV0, Epoch, LOG_TARGET};
25use sc_client_api::backend::AuxStore;
26use sc_consensus_epochs::{
27 migration::{EpochChangesV0For, EpochChangesV1For},
28 EpochChangesFor, SharedEpochChanges,
29};
30use sp_blockchain::{Error as ClientError, Result as ClientResult};
31use sp_consensus_babe::{BabeBlockWeight, BabeConfiguration};
32use sp_runtime::traits::Block as BlockT;
33
34const BABE_EPOCH_CHANGES_VERSION: &[u8] = b"babe_epoch_changes_version";
35const BABE_EPOCH_CHANGES_KEY: &[u8] = b"babe_epoch_changes";
36const BABE_EPOCH_CHANGES_CURRENT_VERSION: u32 = 3;
37
38pub fn block_weight_key<H: Encode>(block_hash: H) -> Vec<u8> {
40 (b"block_weight", block_hash).encode()
41}
42
43fn load_decode<B, T>(backend: &B, key: &[u8]) -> ClientResult<Option<T>>
44where
45 B: AuxStore,
46 T: Decode,
47{
48 let corrupt = |e: codec::Error| {
49 ClientError::Backend(format!("BABE DB is corrupted. Decode error: {}", e))
50 };
51 match backend.get_aux(key)? {
52 None => Ok(None),
53 Some(t) => T::decode(&mut &t[..]).map(Some).map_err(corrupt),
54 }
55}
56
57pub fn load_epoch_changes<Block: BlockT, B: AuxStore>(
59 backend: &B,
60 config: &BabeConfiguration,
61) -> ClientResult<SharedEpochChanges<Block, Epoch>> {
62 let version = load_decode::<_, u32>(backend, BABE_EPOCH_CHANGES_VERSION)?;
63
64 let maybe_epoch_changes = match version {
65 None =>
66 load_decode::<_, EpochChangesV0For<Block, EpochV0>>(backend, BABE_EPOCH_CHANGES_KEY)?
67 .map(|v0| v0.migrate().map(|_, _, epoch| epoch.migrate(config))),
68 Some(1) =>
69 load_decode::<_, EpochChangesV1For<Block, EpochV0>>(backend, BABE_EPOCH_CHANGES_KEY)?
70 .map(|v1| v1.migrate().map(|_, _, epoch| epoch.migrate(config))),
71 Some(2) => {
72 load_decode::<_, EpochChangesV1For<Block, Epoch>>(backend, BABE_EPOCH_CHANGES_KEY)?
74 .map(|v2| v2.migrate())
75 },
76 Some(BABE_EPOCH_CHANGES_CURRENT_VERSION) =>
77 load_decode::<_, EpochChangesFor<Block, Epoch>>(backend, BABE_EPOCH_CHANGES_KEY)?,
78 Some(other) =>
79 return Err(ClientError::Backend(format!("Unsupported BABE DB version: {:?}", other))),
80 };
81
82 let epoch_changes =
83 SharedEpochChanges::<Block, Epoch>::new(maybe_epoch_changes.unwrap_or_else(|| {
84 info!(
85 target: LOG_TARGET,
86 "👶 Creating empty BABE epoch changes on what appears to be first startup.",
87 );
88 EpochChangesFor::<Block, Epoch>::default()
89 }));
90
91 epoch_changes.shared_data().rebalance();
96
97 Ok(epoch_changes)
98}
99
100pub(crate) fn write_epoch_changes<Block: BlockT, F, R>(
102 epoch_changes: &EpochChangesFor<Block, Epoch>,
103 write_aux: F,
104) -> R
105where
106 F: FnOnce(&[(&'static [u8], &[u8])]) -> R,
107{
108 BABE_EPOCH_CHANGES_CURRENT_VERSION.using_encoded(|version| {
109 let encoded_epoch_changes = epoch_changes.encode();
110 write_aux(&[
111 (BABE_EPOCH_CHANGES_KEY, encoded_epoch_changes.as_slice()),
112 (BABE_EPOCH_CHANGES_VERSION, version),
113 ])
114 })
115}
116
117pub(crate) fn write_block_weight<H: Encode, F, R>(
119 block_hash: H,
120 block_weight: BabeBlockWeight,
121 write_aux: F,
122) -> R
123where
124 F: FnOnce(&[(Vec<u8>, &[u8])]) -> R,
125{
126 let key = block_weight_key(block_hash);
127 block_weight.using_encoded(|s| write_aux(&[(key, s)]))
128}
129
130pub fn load_block_weight<H: Encode, B: AuxStore>(
132 backend: &B,
133 block_hash: H,
134) -> ClientResult<Option<BabeBlockWeight>> {
135 load_decode(backend, block_weight_key(block_hash).as_slice())
136}
137
138#[cfg(test)]
139mod test {
140 use super::*;
141 use crate::migration::EpochV0;
142 use fork_tree::ForkTree;
143 use sc_consensus_epochs::{EpochHeader, PersistedEpoch, PersistedEpochHeader};
144 use sc_network_test::Block as TestBlock;
145 use sp_consensus::Error as ConsensusError;
146 use sp_consensus_babe::AllowedSlots;
147 use sp_core::H256;
148 use sp_runtime::traits::NumberFor;
149 use substrate_test_runtime_client;
150
151 #[test]
152 fn load_decode_from_v0_epoch_changes() {
153 let epoch = EpochV0 {
154 start_slot: 0.into(),
155 authorities: vec![],
156 randomness: [0; 32],
157 epoch_index: 1,
158 duration: 100,
159 };
160 let client = substrate_test_runtime_client::new();
161 let mut v0_tree = ForkTree::<H256, NumberFor<TestBlock>, _>::new();
162 v0_tree
163 .import::<_, ConsensusError>(
164 Default::default(),
165 Default::default(),
166 PersistedEpoch::Regular(epoch),
167 &|_, _| Ok(false), )
169 .unwrap();
170
171 client
172 .insert_aux(
173 &[(
174 BABE_EPOCH_CHANGES_KEY,
175 &EpochChangesV0For::<TestBlock, EpochV0>::from_raw(v0_tree).encode()[..],
176 )],
177 &[],
178 )
179 .unwrap();
180
181 assert_eq!(load_decode::<_, u32>(&client, BABE_EPOCH_CHANGES_VERSION).unwrap(), None);
182
183 let epoch_changes = load_epoch_changes::<TestBlock, _>(
184 &client,
185 &BabeConfiguration {
186 slot_duration: 10,
187 epoch_length: 4,
188 c: (3, 10),
189 authorities: Vec::new(),
190 randomness: Default::default(),
191 allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
192 },
193 )
194 .unwrap();
195
196 assert!(
197 epoch_changes
198 .shared_data()
199 .tree()
200 .iter()
201 .map(|(_, _, epoch)| epoch.clone())
202 .collect::<Vec<_>>() ==
203 vec![PersistedEpochHeader::Regular(EpochHeader {
204 start_slot: 0.into(),
205 end_slot: 100.into(),
206 })],
207 ); write_epoch_changes::<TestBlock, _, _>(&epoch_changes.shared_data(), |values| {
210 client.insert_aux(values, &[]).unwrap();
211 });
212
213 assert_eq!(load_decode::<_, u32>(&client, BABE_EPOCH_CHANGES_VERSION).unwrap(), Some(3));
214 }
215}