node_bench/
tempdb.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19use kvdb::{DBKeyValue, DBTransaction, KeyValueDB};
20use kvdb_rocksdb::{Database, DatabaseConfig};
21use std::{io, path::PathBuf, sync::Arc};
22
23#[derive(Clone, Copy, Debug)]
24pub enum DatabaseType {
25	RocksDb,
26	ParityDb,
27}
28
29pub struct TempDatabase(tempfile::TempDir);
30
31struct ParityDbWrapper(parity_db::Db);
32
33impl KeyValueDB for ParityDbWrapper {
34	/// Get a value by key.
35	fn get(&self, col: u32, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
36		Ok(self.0.get(col as u8, &key[key.len() - 32..]).expect("db error"))
37	}
38
39	/// Get a value by partial key. Only works for flushed data.
40	fn get_by_prefix(&self, _col: u32, _prefix: &[u8]) -> io::Result<Option<Vec<u8>>> {
41		unimplemented!()
42	}
43
44	/// Write a transaction of changes to the buffer.
45	fn write(&self, transaction: DBTransaction) -> io::Result<()> {
46		self.0
47			.commit(transaction.ops.iter().map(|op| match op {
48				kvdb::DBOp::Insert { col, key, value } =>
49					(*col as u8, &key[key.len() - 32..], Some(value.to_vec())),
50				kvdb::DBOp::Delete { col, key } => (*col as u8, &key[key.len() - 32..], None),
51				kvdb::DBOp::DeletePrefix { col: _, prefix: _ } => unimplemented!(),
52			}))
53			.expect("db error");
54		Ok(())
55	}
56
57	/// Iterate over flushed data for a given column.
58	fn iter<'a>(&'a self, _col: u32) -> Box<dyn Iterator<Item = io::Result<DBKeyValue>> + 'a> {
59		unimplemented!()
60	}
61
62	/// Iterate over flushed data for a given column, starting from a given prefix.
63	fn iter_with_prefix<'a>(
64		&'a self,
65		_col: u32,
66		_prefix: &'a [u8],
67	) -> Box<dyn Iterator<Item = io::Result<DBKeyValue>> + 'a> {
68		unimplemented!()
69	}
70}
71
72impl TempDatabase {
73	pub fn new() -> Self {
74		let dir = tempfile::tempdir().expect("temp dir creation failed");
75		log::trace!(
76			target: "bench-logistics",
77			"Created temp db at {}",
78			dir.path().to_string_lossy(),
79		);
80
81		TempDatabase(dir)
82	}
83
84	pub fn open(&mut self, db_type: DatabaseType) -> Arc<dyn KeyValueDB> {
85		match db_type {
86			DatabaseType::RocksDb => {
87				let db_cfg = DatabaseConfig::with_columns(1);
88				let db = Database::open(&db_cfg, &self.0.path()).expect("Database backend error");
89				Arc::new(db)
90			},
91			DatabaseType::ParityDb => Arc::new(ParityDbWrapper({
92				let mut options = parity_db::Options::with_columns(self.0.path(), 1);
93				let column_options = &mut options.columns[0];
94				column_options.ref_counted = true;
95				column_options.preimage = true;
96				column_options.uniform = true;
97				parity_db::Db::open_or_create(&options).expect("db open error")
98			})),
99		}
100	}
101}
102
103impl Clone for TempDatabase {
104	fn clone(&self) -> Self {
105		let new_dir = tempfile::tempdir().expect("temp dir creation failed");
106		let self_dir = self.0.path();
107
108		log::trace!(
109			target: "bench-logistics",
110			"Cloning db ({}) to {}",
111			self_dir.to_string_lossy(),
112			new_dir.path().to_string_lossy(),
113		);
114		let self_db_files = std::fs::read_dir(self_dir)
115			.expect("failed to list file in seed dir")
116			.map(|f_result| f_result.expect("failed to read file in seed db").path())
117			.collect::<Vec<PathBuf>>();
118		fs_extra::copy_items(&self_db_files, new_dir.path(), &fs_extra::dir::CopyOptions::new())
119			.expect("Copy of seed database is ok");
120
121		TempDatabase(new_dir)
122	}
123}
node_bench/tempdb.rs

node_bench/
tempdb.rs
