Skip to main content

tetsy_kvdb/
lib.rs

1// Copyright 2020 Parity Technologies
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Key-Value store abstraction.
10
11use smallvec::SmallVec;
12use std::io;
13
14mod io_stats;
15
16/// Required length of prefixes.
17pub const PREFIX_LEN: usize = 12;
18
19/// Database value.
20pub type DBValue = Vec<u8>;
21/// Database keys.
22pub type DBKey = SmallVec<[u8; 32]>;
23
24pub use io_stats::{IoStats, Kind as IoStatsKind};
25
26/// Write transaction. Batches a sequence of put/delete operations for efficiency.
27#[derive(Default, Clone, PartialEq)]
28pub struct DBTransaction {
29	/// Database operations.
30	pub ops: Vec<DBOp>,
31}
32
33/// Database operation.
34#[derive(Clone, PartialEq)]
35pub enum DBOp {
36	Insert { col: u32, key: DBKey, value: DBValue },
37	Delete { col: u32, key: DBKey },
38	DeletePrefix { col: u32, prefix: DBKey },
39}
40
41impl DBOp {
42	/// Returns the key associated with this operation.
43	pub fn key(&self) -> &[u8] {
44		match *self {
45			DBOp::Insert { ref key, .. } => key,
46			DBOp::Delete { ref key, .. } => key,
47			DBOp::DeletePrefix { ref prefix, .. } => prefix,
48		}
49	}
50
51	/// Returns the column associated with this operation.
52	pub fn col(&self) -> u32 {
53		match *self {
54			DBOp::Insert { col, .. } => col,
55			DBOp::Delete { col, .. } => col,
56			DBOp::DeletePrefix { col, .. } => col,
57		}
58	}
59}
60
61impl DBTransaction {
62	/// Create new transaction.
63	pub fn new() -> DBTransaction {
64		DBTransaction::with_capacity(256)
65	}
66
67	/// Create new transaction with capacity.
68	pub fn with_capacity(cap: usize) -> DBTransaction {
69		DBTransaction { ops: Vec::with_capacity(cap) }
70	}
71
72	/// Insert a key-value pair in the transaction. Any existing value will be overwritten upon write.
73	pub fn put(&mut self, col: u32, key: &[u8], value: &[u8]) {
74		self.ops.push(DBOp::Insert { col, key: DBKey::from_slice(key), value: value.to_vec() })
75	}
76
77	/// Insert a key-value pair in the transaction. Any existing value will be overwritten upon write.
78	pub fn put_vec(&mut self, col: u32, key: &[u8], value: Vec<u8>) {
79		self.ops.push(DBOp::Insert { col, key: DBKey::from_slice(key), value });
80	}
81
82	/// Delete value by key.
83	pub fn delete(&mut self, col: u32, key: &[u8]) {
84		self.ops.push(DBOp::Delete { col, key: DBKey::from_slice(key) });
85	}
86
87	/// Delete all values with the given key prefix.
88	/// Using an empty prefix here will remove all keys
89	/// (all keys start with the empty prefix).
90	pub fn delete_prefix(&mut self, col: u32, prefix: &[u8]) {
91		self.ops.push(DBOp::DeletePrefix { col, prefix: DBKey::from_slice(prefix) });
92	}
93}
94
95/// Generic key-value database.
96///
97/// The `KeyValueDB` deals with "column families", which can be thought of as distinct
98/// stores within a database. Keys written in one column family will not be accessible from
99/// any other. The number of column families must be specified at initialization, with a
100/// differing interface for each database.
101///
102/// The API laid out here, along with the `Sync` bound implies interior synchronization for
103/// implementation.
104pub trait KeyValueDB: Sync + Send + tetsy_util_mem::MallocSizeOf {
105	/// Helper to create a new transaction.
106	fn transaction(&self) -> DBTransaction {
107		DBTransaction::new()
108	}
109
110	/// Get a value by key.
111	fn get(&self, col: u32, key: &[u8]) -> io::Result<Option<DBValue>>;
112
113	/// Get the first value matching the given prefix.
114	fn get_by_prefix(&self, col: u32, prefix: &[u8]) -> Option<Box<[u8]>>;
115
116	/// Write a transaction of changes to the backing store.
117	fn write(&self, transaction: DBTransaction) -> io::Result<()>;
118
119	/// Iterate over the data for a given column.
120	fn iter<'a>(&'a self, col: u32) -> Box<dyn Iterator<Item = (Box<[u8]>, Box<[u8]>)> + 'a>;
121
122	/// Iterate over the data for a given column, returning all key/value pairs
123	/// where the key starts with the given prefix.
124	fn iter_with_prefix<'a>(
125		&'a self,
126		col: u32,
127		prefix: &'a [u8],
128	) -> Box<dyn Iterator<Item = (Box<[u8]>, Box<[u8]>)> + 'a>;
129
130	/// Attempt to replace this database with a new one located at the given path.
131	fn restore(&self, new_db: &str) -> io::Result<()>;
132
133	/// Query statistics.
134	///
135	/// Not all tetsy-kvdb implementations are able or expected to implement this, so by
136	/// default, empty statistics is returned. Also, not all tetsy-kvdb implementation
137	/// can return every statistic or configured to do so (some statistics gathering
138	/// may impede the performance and might be off by default).
139	fn io_stats(&self, _kind: IoStatsKind) -> IoStats {
140		IoStats::empty()
141	}
142
143	/// Check for the existence of a value by key.
144	fn has_key(&self, col: u32, key: &[u8]) -> io::Result<bool> {
145		self.get(col, key).map(|opt| opt.is_some())
146	}
147
148	/// Check for the existence of a value by prefix.
149	fn has_prefix(&self, col: u32, prefix: &[u8]) -> bool {
150		self.get_by_prefix(col, prefix).is_some()
151	}
152}
153
154/// For a given start prefix (inclusive), returns the correct end prefix (non-inclusive).
155/// This assumes the key bytes are ordered in lexicographical order.
156/// Since key length is not limited, for some case we return `None` because there is
157/// no bounded limit (every keys in the serie `[]`, `[255]`, `[255, 255]` ...).
158pub fn end_prefix(prefix: &[u8]) -> Option<Vec<u8>> {
159	let mut end_range = prefix.to_vec();
160	while let Some(0xff) = end_range.last() {
161		end_range.pop();
162	}
163	if let Some(byte) = end_range.last_mut() {
164		*byte += 1;
165		Some(end_range)
166	} else {
167		None
168	}
169}
170
171#[cfg(test)]
172mod test {
173	use super::end_prefix;
174
175	#[test]
176	fn end_prefix_test() {
177		assert_eq!(end_prefix(&[5, 6, 7]), Some(vec![5, 6, 8]));
178		assert_eq!(end_prefix(&[5, 6, 255]), Some(vec![5, 7]));
179		// This is not equal as the result is before start.
180		assert_ne!(end_prefix(&[5, 255, 255]), Some(vec![5, 255]));
181		// This is equal ([5, 255] will not be deleted because
182		// it is before start).
183		assert_eq!(end_prefix(&[5, 255, 255]), Some(vec![6]));
184		assert_eq!(end_prefix(&[255, 255, 255]), None);
185
186		assert_eq!(end_prefix(&[0x00, 0xff]), Some(vec![0x01]));
187		assert_eq!(end_prefix(&[0xff]), None);
188		assert_eq!(end_prefix(&[]), None);
189		assert_eq!(end_prefix(b"0"), Some(b"1".to_vec()));
190	}
191}