1use crate::{DbState, DbStateBuilder};
22use hash_db::{Hasher as DbHasher, Prefix};
23use kvdb::{DBTransaction, KeyValueDB};
24use linked_hash_map::LinkedHashMap;
25use parking_lot::Mutex;
26use sp_core::{
27 hexdisplay::HexDisplay,
28 storage::{ChildInfo, TrackedStorageKey},
29};
30use sp_runtime::{traits::Hash, StateVersion, Storage};
31use sp_state_machine::{
32 backend::Backend as StateBackend, BackendTransaction, ChildStorageCollection, DBValue,
33 IterArgs, StorageCollection, StorageIterator, StorageKey, StorageValue,
34};
35use sp_trie::{
36 cache::{CacheSize, SharedTrieCache},
37 prefixed_key, MemoryDB, MerkleValue,
38};
39use std::{
40 cell::{Cell, RefCell},
41 collections::HashMap,
42 sync::Arc,
43};
44
45type State<H> = DbState<H>;
46
47struct StorageDb<Hasher> {
48 db: Arc<dyn KeyValueDB>,
49 _phantom: std::marker::PhantomData<Hasher>,
50}
51
52impl<Hasher: Hash> sp_state_machine::Storage<Hasher> for StorageDb<Hasher> {
53 fn get(&self, key: &Hasher::Output, prefix: Prefix) -> Result<Option<DBValue>, String> {
54 let prefixed_key = prefixed_key::<Hasher>(key, prefix);
55 self.db
56 .get(0, &prefixed_key)
57 .map_err(|e| format!("Database backend error: {:?}", e))
58 }
59}
60
61struct KeyTracker {
62 enable_tracking: bool,
63 main_keys: LinkedHashMap<Vec<u8>, TrackedStorageKey>,
67 child_keys: LinkedHashMap<Vec<u8>, LinkedHashMap<Vec<u8>, TrackedStorageKey>>,
72}
73
74pub struct BenchmarkingState<Hasher: Hash> {
76 root: Cell<Hasher::Output>,
77 genesis_root: Hasher::Output,
78 state: RefCell<Option<State<Hasher>>>,
79 db: Cell<Option<Arc<dyn KeyValueDB>>>,
80 genesis: HashMap<Vec<u8>, (Vec<u8>, i32)>,
81 record: Cell<Vec<Vec<u8>>>,
82 key_tracker: Arc<Mutex<KeyTracker>>,
83 whitelist: RefCell<Vec<TrackedStorageKey>>,
84 proof_recorder: Option<sp_trie::recorder::Recorder<Hasher>>,
85 proof_recorder_root: Cell<Hasher::Output>,
86 shared_trie_cache: SharedTrieCache<Hasher>,
87}
88
89pub struct RawIter<Hasher: Hash> {
91 inner: <DbState<Hasher> as StateBackend<Hasher>>::RawIter,
92 child_trie: Option<Vec<u8>>,
93 key_tracker: Arc<Mutex<KeyTracker>>,
94}
95
96impl<Hasher: Hash> StorageIterator<Hasher> for RawIter<Hasher> {
97 type Backend = BenchmarkingState<Hasher>;
98 type Error = String;
99
100 fn next_key(&mut self, backend: &Self::Backend) -> Option<Result<StorageKey, Self::Error>> {
101 match self.inner.next_key(backend.state.borrow().as_ref()?) {
102 Some(Ok(key)) => {
103 self.key_tracker.lock().add_read_key(self.child_trie.as_deref(), &key);
104 Some(Ok(key))
105 },
106 result => result,
107 }
108 }
109
110 fn next_pair(
111 &mut self,
112 backend: &Self::Backend,
113 ) -> Option<Result<(StorageKey, StorageValue), Self::Error>> {
114 match self.inner.next_pair(backend.state.borrow().as_ref()?) {
115 Some(Ok((key, value))) => {
116 self.key_tracker.lock().add_read_key(self.child_trie.as_deref(), &key);
117 Some(Ok((key, value)))
118 },
119 result => result,
120 }
121 }
122
123 fn was_complete(&self) -> bool {
124 self.inner.was_complete()
125 }
126}
127
128impl<Hasher: Hash> BenchmarkingState<Hasher> {
129 pub fn new(
131 genesis: Storage,
132 _cache_size_mb: Option<usize>,
133 record_proof: bool,
134 enable_tracking: bool,
135 ) -> Result<Self, String> {
136 let state_version = sp_runtime::StateVersion::default();
137 let mut root = Default::default();
138 let mut mdb = MemoryDB::<Hasher>::default();
139 sp_trie::trie_types::TrieDBMutBuilderV1::<Hasher>::new(&mut mdb, &mut root).build();
140
141 let mut state = BenchmarkingState {
142 state: RefCell::new(None),
143 db: Cell::new(None),
144 root: Cell::new(root),
145 genesis: Default::default(),
146 genesis_root: Default::default(),
147 record: Default::default(),
148 key_tracker: Arc::new(Mutex::new(KeyTracker {
149 main_keys: Default::default(),
150 child_keys: Default::default(),
151 enable_tracking,
152 })),
153 whitelist: Default::default(),
154 proof_recorder: record_proof.then(Default::default),
155 proof_recorder_root: Cell::new(root),
156 shared_trie_cache: SharedTrieCache::new(CacheSize::new(0), None),
158 };
159
160 state.add_whitelist_to_tracker();
161
162 state.reopen()?;
163 let child_delta = genesis.children_default.values().map(|child_content| {
164 (
165 &child_content.child_info,
166 child_content.data.iter().map(|(k, v)| (k.as_ref(), Some(v.as_ref()))),
167 )
168 });
169 let (root, transaction): (Hasher::Output, _) =
170 state.state.borrow().as_ref().unwrap().full_storage_root(
171 genesis.top.iter().map(|(k, v)| (k.as_ref(), Some(v.as_ref()))),
172 child_delta,
173 state_version,
174 );
175 state.genesis = transaction.clone().drain();
176 state.genesis_root = root;
177 state.commit(root, transaction, Vec::new(), Vec::new())?;
178 state.record.take();
179 Ok(state)
180 }
181
182 pub fn recorder(&self) -> Option<sp_trie::recorder::Recorder<Hasher>> {
184 self.proof_recorder.clone()
185 }
186
187 fn reopen(&self) -> Result<(), String> {
188 *self.state.borrow_mut() = None;
189 let db = match self.db.take() {
190 Some(db) => db,
191 None => Arc::new(kvdb_memorydb::create(1)),
192 };
193 self.db.set(Some(db.clone()));
194 if let Some(recorder) = &self.proof_recorder {
195 recorder.reset();
196 self.proof_recorder_root.set(self.root.get());
197 }
198 let storage_db = Arc::new(StorageDb::<Hasher> { db, _phantom: Default::default() });
199 *self.state.borrow_mut() = Some(
200 DbStateBuilder::<Hasher>::new(storage_db, self.root.get())
201 .with_optional_recorder(self.proof_recorder.clone())
202 .with_cache(self.shared_trie_cache.local_cache_trusted())
203 .build(),
204 );
205 Ok(())
206 }
207
208 fn add_whitelist_to_tracker(&self) {
209 self.key_tracker.lock().add_whitelist(&self.whitelist.borrow());
210 }
211
212 fn wipe_tracker(&self) {
213 let mut key_tracker = self.key_tracker.lock();
214 key_tracker.main_keys = LinkedHashMap::new();
215 key_tracker.child_keys = LinkedHashMap::new();
216 key_tracker.add_whitelist(&self.whitelist.borrow());
217 }
218
219 fn add_read_key(&self, childtrie: Option<&[u8]>, key: &[u8]) {
220 self.key_tracker.lock().add_read_key(childtrie, key);
221 }
222
223 fn add_write_key(&self, childtrie: Option<&[u8]>, key: &[u8]) {
224 self.key_tracker.lock().add_write_key(childtrie, key);
225 }
226
227 fn all_trackers(&self) -> Vec<TrackedStorageKey> {
228 self.key_tracker.lock().all_trackers()
229 }
230}
231
232impl KeyTracker {
233 fn add_whitelist(&mut self, whitelist: &[TrackedStorageKey]) {
234 whitelist.iter().for_each(|key| {
235 let mut whitelisted = TrackedStorageKey::new(key.key.clone());
236 whitelisted.whitelist();
237 self.main_keys.insert(key.key.clone(), whitelisted);
238 });
239 }
240
241 fn add_read_key(&mut self, childtrie: Option<&[u8]>, key: &[u8]) {
243 if !self.enable_tracking {
244 return
245 }
246
247 let child_key_tracker = &mut self.child_keys;
248 let main_key_tracker = &mut self.main_keys;
249
250 let key_tracker = if let Some(childtrie) = childtrie {
251 child_key_tracker.entry(childtrie.to_vec()).or_insert_with(LinkedHashMap::new)
252 } else {
253 main_key_tracker
254 };
255
256 let should_log = match key_tracker.get_mut(key) {
257 None => {
258 let mut has_been_read = TrackedStorageKey::new(key.to_vec());
259 has_been_read.add_read();
260 key_tracker.insert(key.to_vec(), has_been_read);
261 true
262 },
263 Some(tracker) => {
264 let should_log = !tracker.has_been_read();
265 tracker.add_read();
266 should_log
267 },
268 };
269
270 if should_log {
271 if let Some(childtrie) = childtrie {
272 log::trace!(
273 target: "benchmark",
274 "Childtrie Read: {} {}", HexDisplay::from(&childtrie), HexDisplay::from(&key)
275 );
276 } else {
277 log::trace!(target: "benchmark", "Read: {}", HexDisplay::from(&key));
278 }
279 }
280 }
281
282 fn add_write_key(&mut self, childtrie: Option<&[u8]>, key: &[u8]) {
284 if !self.enable_tracking {
285 return
286 }
287
288 let child_key_tracker = &mut self.child_keys;
289 let main_key_tracker = &mut self.main_keys;
290
291 let key_tracker = if let Some(childtrie) = childtrie {
292 child_key_tracker.entry(childtrie.to_vec()).or_insert_with(LinkedHashMap::new)
293 } else {
294 main_key_tracker
295 };
296
297 let should_log = match key_tracker.get_mut(key) {
299 None => {
300 let mut has_been_written = TrackedStorageKey::new(key.to_vec());
301 has_been_written.add_write();
302 key_tracker.insert(key.to_vec(), has_been_written);
303 true
304 },
305 Some(tracker) => {
306 let should_log = !tracker.has_been_written();
307 tracker.add_write();
308 should_log
309 },
310 };
311
312 if should_log {
313 if let Some(childtrie) = childtrie {
314 log::trace!(
315 target: "benchmark",
316 "Childtrie Write: {} {}", HexDisplay::from(&childtrie), HexDisplay::from(&key)
317 );
318 } else {
319 log::trace!(target: "benchmark", "Write: {}", HexDisplay::from(&key));
320 }
321 }
322 }
323
324 fn all_trackers(&self) -> Vec<TrackedStorageKey> {
326 let mut all_trackers = Vec::new();
327
328 self.main_keys.iter().for_each(|(_, tracker)| {
329 all_trackers.push(tracker.clone());
330 });
331
332 self.child_keys.iter().for_each(|(_, child_tracker)| {
333 child_tracker.iter().for_each(|(_, tracker)| {
334 all_trackers.push(tracker.clone());
335 });
336 });
337
338 all_trackers
339 }
340}
341
342fn state_err() -> String {
343 "State is not open".into()
344}
345
346impl<Hasher: Hash> StateBackend<Hasher> for BenchmarkingState<Hasher> {
347 type Error = <DbState<Hasher> as StateBackend<Hasher>>::Error;
348 type TrieBackendStorage = <DbState<Hasher> as StateBackend<Hasher>>::TrieBackendStorage;
349 type RawIter = RawIter<Hasher>;
350
351 fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
352 self.add_read_key(None, key);
353 self.state.borrow().as_ref().ok_or_else(state_err)?.storage(key)
354 }
355
356 fn storage_hash(&self, key: &[u8]) -> Result<Option<Hasher::Output>, Self::Error> {
357 self.add_read_key(None, key);
358 self.state.borrow().as_ref().ok_or_else(state_err)?.storage_hash(key)
359 }
360
361 fn child_storage(
362 &self,
363 child_info: &ChildInfo,
364 key: &[u8],
365 ) -> Result<Option<Vec<u8>>, Self::Error> {
366 self.add_read_key(Some(child_info.storage_key()), key);
367 self.state
368 .borrow()
369 .as_ref()
370 .ok_or_else(state_err)?
371 .child_storage(child_info, key)
372 }
373
374 fn child_storage_hash(
375 &self,
376 child_info: &ChildInfo,
377 key: &[u8],
378 ) -> Result<Option<Hasher::Output>, Self::Error> {
379 self.add_read_key(Some(child_info.storage_key()), key);
380 self.state
381 .borrow()
382 .as_ref()
383 .ok_or_else(state_err)?
384 .child_storage_hash(child_info, key)
385 }
386
387 fn closest_merkle_value(
388 &self,
389 key: &[u8],
390 ) -> Result<Option<MerkleValue<Hasher::Output>>, Self::Error> {
391 self.add_read_key(None, key);
392 self.state.borrow().as_ref().ok_or_else(state_err)?.closest_merkle_value(key)
393 }
394
395 fn child_closest_merkle_value(
396 &self,
397 child_info: &ChildInfo,
398 key: &[u8],
399 ) -> Result<Option<MerkleValue<Hasher::Output>>, Self::Error> {
400 self.add_read_key(None, key);
401 self.state
402 .borrow()
403 .as_ref()
404 .ok_or_else(state_err)?
405 .child_closest_merkle_value(child_info, key)
406 }
407
408 fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
409 self.add_read_key(None, key);
410 self.state.borrow().as_ref().ok_or_else(state_err)?.exists_storage(key)
411 }
412
413 fn exists_child_storage(
414 &self,
415 child_info: &ChildInfo,
416 key: &[u8],
417 ) -> Result<bool, Self::Error> {
418 self.add_read_key(Some(child_info.storage_key()), key);
419 self.state
420 .borrow()
421 .as_ref()
422 .ok_or_else(state_err)?
423 .exists_child_storage(child_info, key)
424 }
425
426 fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
427 self.add_read_key(None, key);
428 self.state.borrow().as_ref().ok_or_else(state_err)?.next_storage_key(key)
429 }
430
431 fn next_child_storage_key(
432 &self,
433 child_info: &ChildInfo,
434 key: &[u8],
435 ) -> Result<Option<Vec<u8>>, Self::Error> {
436 self.add_read_key(Some(child_info.storage_key()), key);
437 self.state
438 .borrow()
439 .as_ref()
440 .ok_or_else(state_err)?
441 .next_child_storage_key(child_info, key)
442 }
443
444 fn storage_root<'a>(
445 &self,
446 delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
447 state_version: StateVersion,
448 ) -> (Hasher::Output, BackendTransaction<Hasher>) {
449 self.state
450 .borrow()
451 .as_ref()
452 .map_or(Default::default(), |s| s.storage_root(delta, state_version))
453 }
454
455 fn child_storage_root<'a>(
456 &self,
457 child_info: &ChildInfo,
458 delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
459 state_version: StateVersion,
460 ) -> (Hasher::Output, bool, BackendTransaction<Hasher>) {
461 self.state
462 .borrow()
463 .as_ref()
464 .map_or(Default::default(), |s| s.child_storage_root(child_info, delta, state_version))
465 }
466
467 fn raw_iter(&self, args: IterArgs) -> Result<Self::RawIter, Self::Error> {
468 let child_trie =
469 args.child_info.as_ref().map(|child_info| child_info.storage_key().to_vec());
470 self.state
471 .borrow()
472 .as_ref()
473 .map(|s| s.raw_iter(args))
474 .unwrap_or(Ok(Default::default()))
475 .map(|raw_iter| RawIter {
476 inner: raw_iter,
477 key_tracker: self.key_tracker.clone(),
478 child_trie,
479 })
480 }
481
482 fn commit(
483 &self,
484 storage_root: <Hasher as DbHasher>::Out,
485 mut transaction: BackendTransaction<Hasher>,
486 main_storage_changes: StorageCollection,
487 child_storage_changes: ChildStorageCollection,
488 ) -> Result<(), Self::Error> {
489 if let Some(db) = self.db.take() {
490 let mut db_transaction = DBTransaction::new();
491 let changes = transaction.drain();
492 let mut keys = Vec::with_capacity(changes.len());
493 for (key, (val, rc)) in changes {
494 if rc > 0 {
495 db_transaction.put(0, &key, &val);
496 } else if rc < 0 {
497 db_transaction.delete(0, &key);
498 }
499 keys.push(key);
500 }
501 let mut record = self.record.take();
502 record.extend(keys);
503 self.record.set(record);
504 db.write(db_transaction)
505 .map_err(|_| String::from("Error committing transaction"))?;
506 self.root.set(storage_root);
507 self.db.set(Some(db));
508
509 main_storage_changes.iter().for_each(|(key, _)| {
511 self.add_write_key(None, key);
512 });
513 child_storage_changes.iter().for_each(|(child_storage_key, storage_changes)| {
514 storage_changes.iter().for_each(|(key, _)| {
515 self.add_write_key(Some(child_storage_key), key);
516 })
517 });
518 } else {
519 return Err("Trying to commit to a closed db".into())
520 }
521 self.reopen()
522 }
523
524 fn wipe(&self) -> Result<(), Self::Error> {
525 let record = self.record.take();
527 if let Some(db) = self.db.take() {
528 let mut db_transaction = DBTransaction::new();
529 for key in record {
530 match self.genesis.get(&key) {
531 Some((v, _)) => db_transaction.put(0, &key, v),
532 None => db_transaction.delete(0, &key),
533 }
534 }
535 db.write(db_transaction)
536 .map_err(|_| String::from("Error committing transaction"))?;
537 self.db.set(Some(db));
538 }
539
540 self.root.set(self.genesis_root);
541 self.reopen()?;
542 self.wipe_tracker();
543 Ok(())
544 }
545
546 fn read_write_count(&self) -> (u32, u32, u32, u32) {
552 let mut reads = 0;
553 let mut repeat_reads = 0;
554 let mut writes = 0;
555 let mut repeat_writes = 0;
556
557 self.all_trackers().iter().for_each(|tracker| {
558 if !tracker.whitelisted {
559 if tracker.reads > 0 {
560 reads += 1;
561 repeat_reads += tracker.reads - 1;
562 }
563
564 if tracker.writes > 0 {
565 writes += 1;
566 repeat_writes += tracker.writes - 1;
567 }
568 }
569 });
570 (reads, repeat_reads, writes, repeat_writes)
571 }
572
573 fn reset_read_write_count(&self) {
575 self.wipe_tracker()
576 }
577
578 fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
579 self.whitelist.borrow().to_vec()
580 }
581
582 fn set_whitelist(&self, new: Vec<TrackedStorageKey>) {
583 *self.whitelist.borrow_mut() = new;
584 }
585
586 fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
587 let mut prefix_key_tracker = LinkedHashMap::<Vec<u8>, (u32, u32, bool)>::new();
591 self.all_trackers().iter().for_each(|tracker| {
592 if !tracker.whitelisted {
593 let prefix_length = tracker.key.len().min(32);
594 let prefix = tracker.key[0..prefix_length].to_vec();
595 let reads = tracker.reads.min(1);
598 let writes = tracker.writes.min(1);
599 if let Some(prefix_tracker) = prefix_key_tracker.get_mut(&prefix) {
600 prefix_tracker.0 += reads;
601 prefix_tracker.1 += writes;
602 } else {
603 prefix_key_tracker.insert(prefix, (reads, writes, tracker.whitelisted));
604 }
605 }
606 });
607
608 prefix_key_tracker
609 .iter()
610 .map(|(key, tracker)| -> (Vec<u8>, u32, u32, bool) {
611 (key.to_vec(), tracker.0, tracker.1, tracker.2)
612 })
613 .collect::<Vec<_>>()
614 }
615
616 fn register_overlay_stats(&self, stats: &sp_state_machine::StateMachineStats) {
617 self.state.borrow().as_ref().map(|s| s.register_overlay_stats(stats));
618 }
619
620 fn usage_info(&self) -> sp_state_machine::UsageInfo {
621 self.state
622 .borrow()
623 .as_ref()
624 .map_or(sp_state_machine::UsageInfo::empty(), |s| s.usage_info())
625 }
626
627 fn proof_size(&self) -> Option<u32> {
628 self.proof_recorder.as_ref().map(|recorder| {
629 let proof_size = recorder.estimate_encoded_size() as u32;
630
631 let proof = recorder.to_storage_proof();
632
633 let proof_recorder_root = self.proof_recorder_root.get();
634 if proof_recorder_root == Default::default() || proof_size == 1 {
635 log::debug!(target: "benchmark", "Some proof size: {}", &proof_size);
637 proof_size
638 } else {
639 if let Some(size) = proof.encoded_compact_size::<Hasher>(proof_recorder_root) {
640 size as u32
641 } else if proof_recorder_root == self.root.get() {
642 log::debug!(target: "benchmark", "No changes - no proof");
643 0
644 } else {
645 panic!(
646 "proof rec root {:?}, root {:?}, genesis {:?}, rec_len {:?}",
647 self.proof_recorder_root.get(),
648 self.root.get(),
649 self.genesis_root,
650 proof_size,
651 );
652 }
653 }
654 })
655 }
656}
657
658impl<Hasher: Hash> std::fmt::Debug for BenchmarkingState<Hasher> {
659 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
660 write!(f, "Bench DB")
661 }
662}
663
664#[cfg(test)]
665mod test {
666 use crate::bench::BenchmarkingState;
667 use sp_runtime::traits::HashingFor;
668 use sp_state_machine::backend::Backend as _;
669
670 fn hex(hex: &str) -> Vec<u8> {
671 array_bytes::hex2bytes(hex).unwrap()
672 }
673
674 #[test]
675 fn iteration_is_also_counted_in_rw_counts() {
676 let storage = sp_runtime::Storage {
677 top: vec![(
678 hex("ce6e1397e668c7fcf47744350dc59688455a2c2dbd2e2a649df4e55d93cd7158"),
679 hex("0102030405060708"),
680 )]
681 .into_iter()
682 .collect(),
683 ..sp_runtime::Storage::default()
684 };
685 let bench_state =
686 BenchmarkingState::<HashingFor<crate::tests::Block>>::new(storage, None, false, true)
687 .unwrap();
688
689 assert_eq!(bench_state.read_write_count(), (0, 0, 0, 0));
690 assert_eq!(bench_state.keys(Default::default()).unwrap().count(), 1);
691 assert_eq!(bench_state.read_write_count(), (1, 0, 0, 0));
692 }
693
694 #[test]
695 fn read_to_main_and_child_tries() {
696 let bench_state = BenchmarkingState::<HashingFor<crate::tests::Block>>::new(
697 Default::default(),
698 None,
699 false,
700 true,
701 )
702 .unwrap();
703
704 for _ in 0..2 {
705 let child1 = sp_core::storage::ChildInfo::new_default(b"child1");
706 let child2 = sp_core::storage::ChildInfo::new_default(b"child2");
707
708 bench_state.storage(b"foo").unwrap();
709 bench_state.child_storage(&child1, b"foo").unwrap();
710 bench_state.child_storage(&child2, b"foo").unwrap();
711
712 bench_state.storage(b"bar").unwrap();
713 bench_state.child_storage(&child1, b"bar").unwrap();
714 bench_state.child_storage(&child2, b"bar").unwrap();
715
716 bench_state
717 .commit(
718 Default::default(),
719 Default::default(),
720 vec![("foo".as_bytes().to_vec(), None)],
721 vec![("child1".as_bytes().to_vec(), vec![("foo".as_bytes().to_vec(), None)])],
722 )
723 .unwrap();
724
725 let rw_tracker = bench_state.read_write_count();
726 assert_eq!(rw_tracker.0, 6);
727 assert_eq!(rw_tracker.1, 0);
728 assert_eq!(rw_tracker.2, 2);
729 assert_eq!(rw_tracker.3, 0);
730 bench_state.wipe().unwrap();
731 }
732 }
733}