polkadot_node_subsystem_util/
database.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// This file is part of Polkadot.
3
4// Polkadot is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Polkadot is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
16
17//! Database trait for polkadot db.
18
19pub use kvdb::{DBKeyValue, DBTransaction, DBValue, KeyValueDB};
20
21/// Database trait with ordered key capacity.
22pub trait Database: KeyValueDB {
23	/// Check if column allows content iteration
24	/// and removal by prefix.
25	fn is_indexed_column(&self, col: u32) -> bool;
26}
27
28/// Implementation for database supporting `KeyValueDB` already.
29pub mod kvdb_impl {
30	use super::{DBKeyValue, DBTransaction, DBValue, Database, KeyValueDB};
31	use kvdb::{DBOp, IoStats, IoStatsKind};
32	use std::{collections::BTreeSet, io::Result};
33
34	/// Adapter implementing subsystem database
35	/// for `KeyValueDB`.
36	#[derive(Clone)]
37	pub struct DbAdapter<D> {
38		db: D,
39		indexed_columns: BTreeSet<u32>,
40	}
41
42	impl<D: KeyValueDB> DbAdapter<D> {
43		/// Instantiate new subsystem database, with
44		/// the columns that allow ordered iteration.
45		pub fn new(db: D, indexed_columns: &[u32]) -> Self {
46			DbAdapter { db, indexed_columns: indexed_columns.iter().cloned().collect() }
47		}
48
49		fn ensure_is_indexed(&self, col: u32) {
50			debug_assert!(
51				self.is_indexed_column(col),
52				"Invalid configuration of database, column {} is not ordered.",
53				col
54			);
55		}
56
57		fn ensure_ops_indexing(&self, transaction: &DBTransaction) {
58			debug_assert!({
59				let mut pass = true;
60				for op in &transaction.ops {
61					if let DBOp::DeletePrefix { col, .. } = op {
62						if !self.is_indexed_column(*col) {
63							pass = false;
64							break;
65						}
66					}
67				}
68				pass
69			})
70		}
71	}
72
73	impl<D: KeyValueDB> Database for DbAdapter<D> {
74		fn is_indexed_column(&self, col: u32) -> bool {
75			self.indexed_columns.contains(&col)
76		}
77	}
78
79	impl<D: KeyValueDB> KeyValueDB for DbAdapter<D> {
80		fn transaction(&self) -> DBTransaction {
81			self.db.transaction()
82		}
83
84		fn get(&self, col: u32, key: &[u8]) -> Result<Option<DBValue>> {
85			self.db.get(col, key)
86		}
87
88		fn get_by_prefix(&self, col: u32, prefix: &[u8]) -> Result<Option<DBValue>> {
89			self.ensure_is_indexed(col);
90			self.db.get_by_prefix(col, prefix)
91		}
92
93		fn write(&self, transaction: DBTransaction) -> Result<()> {
94			self.ensure_ops_indexing(&transaction);
95			self.db.write(transaction)
96		}
97
98		fn iter<'a>(&'a self, col: u32) -> Box<dyn Iterator<Item = Result<DBKeyValue>> + 'a> {
99			self.ensure_is_indexed(col);
100			self.db.iter(col)
101		}
102
103		fn iter_with_prefix<'a>(
104			&'a self,
105			col: u32,
106			prefix: &'a [u8],
107		) -> Box<dyn Iterator<Item = Result<DBKeyValue>> + 'a> {
108			self.ensure_is_indexed(col);
109			self.db.iter_with_prefix(col, prefix)
110		}
111
112		fn io_stats(&self, kind: IoStatsKind) -> IoStats {
113			self.db.io_stats(kind)
114		}
115
116		fn has_key(&self, col: u32, key: &[u8]) -> Result<bool> {
117			self.db.has_key(col, key)
118		}
119
120		fn has_prefix(&self, col: u32, prefix: &[u8]) -> Result<bool> {
121			self.ensure_is_indexed(col);
122			self.db.has_prefix(col, prefix)
123		}
124	}
125}
126
127/// Utilities for using parity-db database.
128pub mod paritydb_impl {
129	use super::{DBKeyValue, DBTransaction, DBValue, Database, KeyValueDB};
130	use kvdb::DBOp;
131	use parity_db::Db;
132	use parking_lot::Mutex;
133	use std::{collections::BTreeSet, io::Result, sync::Arc};
134
135	fn handle_err<T>(result: parity_db::Result<T>) -> T {
136		match result {
137			Ok(r) => r,
138			Err(e) => {
139				panic!("Critical database error: {:?}", e);
140			},
141		}
142	}
143
144	fn map_err<T>(result: parity_db::Result<T>) -> Result<T> {
145		result.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, format!("{:?}", e)))
146	}
147
148	/// Implementation of of `Database` for parity-db adapter.
149	pub struct DbAdapter {
150		db: Db,
151		indexed_columns: BTreeSet<u32>,
152		write_lock: Arc<Mutex<()>>,
153	}
154
155	impl KeyValueDB for DbAdapter {
156		fn transaction(&self) -> DBTransaction {
157			DBTransaction::new()
158		}
159
160		fn get(&self, col: u32, key: &[u8]) -> Result<Option<DBValue>> {
161			map_err(self.db.get(col as u8, key))
162		}
163
164		fn get_by_prefix(&self, col: u32, prefix: &[u8]) -> Result<Option<DBValue>> {
165			self.iter_with_prefix(col, prefix)
166				.next()
167				.transpose()
168				.map(|mb| mb.map(|(_, v)| v))
169		}
170
171		fn iter<'a>(&'a self, col: u32) -> Box<dyn Iterator<Item = Result<DBKeyValue>> + 'a> {
172			let mut iter = match self.db.iter(col as u8) {
173				Ok(iter) => iter,
174				Err(e) => return Box::new(std::iter::once(map_err(Err(e)))),
175			};
176			Box::new(std::iter::from_fn(move || {
177				iter.next().transpose().map(|r| map_err(r.map(|(k, v)| (k.into(), v))))
178			}))
179		}
180
181		fn iter_with_prefix<'a>(
182			&'a self,
183			col: u32,
184			prefix: &'a [u8],
185		) -> Box<dyn Iterator<Item = Result<DBKeyValue>> + 'a> {
186			if prefix.len() == 0 {
187				return self.iter(col);
188			}
189			let mut iter = match self.db.iter(col as u8) {
190				Ok(iter) => iter,
191				Err(e) => return Box::new(std::iter::once(map_err(Err(e)))),
192			};
193			if let Err(e) = iter.seek(prefix) {
194				return Box::new(std::iter::once(map_err(Err(e))));
195			}
196			Box::new(std::iter::from_fn(move || {
197				iter.next().transpose().and_then(|r| {
198					map_err(r.map(|(k, v)| k.starts_with(prefix).then(|| (k.into(), v))))
199						.transpose()
200				})
201			}))
202		}
203
204		fn write(&self, transaction: DBTransaction) -> Result<()> {
205			let mut ops = transaction.ops.into_iter();
206			// TODO using a key iterator or native delete here would be faster.
207			let mut current_prefix_iter: Option<(parity_db::BTreeIterator, u8, Vec<u8>)> = None;
208			let current_prefix_iter = &mut current_prefix_iter;
209			let transaction = std::iter::from_fn(move || loop {
210				if let Some((prefix_iter, col, prefix)) = current_prefix_iter {
211					if let Some((key, _value)) = handle_err(prefix_iter.next()) {
212						if key.starts_with(prefix) {
213							return Some((*col, key.to_vec(), None));
214						}
215					}
216					*current_prefix_iter = None;
217				}
218				return match ops.next() {
219					None => None,
220					Some(DBOp::Insert { col, key, value }) => {
221						Some((col as u8, key.to_vec(), Some(value)))
222					},
223					Some(DBOp::Delete { col, key }) => Some((col as u8, key.to_vec(), None)),
224					Some(DBOp::DeletePrefix { col, prefix }) => {
225						let col = col as u8;
226						let mut iter = handle_err(self.db.iter(col));
227						handle_err(iter.seek(&prefix[..]));
228						*current_prefix_iter = Some((iter, col, prefix.to_vec()));
229						continue;
230					},
231				};
232			});
233
234			// Locking is required due to possible racy change of the content of a deleted prefix.
235			let _lock = self.write_lock.lock();
236			map_err(self.db.commit(transaction))
237		}
238	}
239
240	impl Database for DbAdapter {
241		fn is_indexed_column(&self, col: u32) -> bool {
242			self.indexed_columns.contains(&col)
243		}
244	}
245
246	impl DbAdapter {
247		/// Implementation of of `Database` for parity-db adapter.
248		pub fn new(db: Db, indexed_columns: &[u32]) -> Self {
249			let write_lock = Arc::new(Mutex::new(()));
250			DbAdapter { db, indexed_columns: indexed_columns.iter().cloned().collect(), write_lock }
251		}
252	}
253
254	#[cfg(test)]
255	mod tests {
256		use super::*;
257		use kvdb_shared_tests as st;
258		use std::io;
259		use tempfile::Builder as TempfileBuilder;
260
261		fn create(num_col: u32) -> io::Result<(DbAdapter, tempfile::TempDir)> {
262			let tempdir = TempfileBuilder::new().prefix("").tempdir()?;
263			let mut options = parity_db::Options::with_columns(tempdir.path(), num_col as u8);
264			for i in 0..num_col {
265				options.columns[i as usize].btree_index = true;
266			}
267
268			let db = parity_db::Db::open_or_create(&options)
269				.map_err(|err| io::Error::new(io::ErrorKind::Other, format!("{:?}", err)))?;
270
271			let db = DbAdapter::new(db, &[0]);
272			Ok((db, tempdir))
273		}
274
275		#[test]
276		fn put_and_get() -> io::Result<()> {
277			let (db, _temp_file) = create(1)?;
278			st::test_put_and_get(&db)
279		}
280
281		#[test]
282		fn delete_and_get() -> io::Result<()> {
283			let (db, _temp_file) = create(1)?;
284			st::test_delete_and_get(&db)
285		}
286
287		#[test]
288		fn delete_prefix() -> io::Result<()> {
289			let (db, _temp_file) = create(st::DELETE_PREFIX_NUM_COLUMNS)?;
290			st::test_delete_prefix(&db)
291		}
292
293		#[test]
294		fn iter() -> io::Result<()> {
295			let (db, _temp_file) = create(1)?;
296			st::test_iter(&db)
297		}
298
299		#[test]
300		fn iter_with_prefix() -> io::Result<()> {
301			let (db, _temp_file) = create(1)?;
302			st::test_iter_with_prefix(&db)
303		}
304
305		#[test]
306		fn complex() -> io::Result<()> {
307			let (db, _temp_file) = create(1)?;
308			st::test_complex(&db)
309		}
310	}
311}
polkadot_node_subsystem_util/database.rs

polkadot_node_subsystem_util/
database.rs
