kvdb_rocksdb/
iter.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//! This module contains an implementation of a RocksDB iterator
10//! wrapped inside a `RwLock`. Since `RwLock` "owns" the inner data,
11//! we're using `owning_ref` to work around the borrowing rules of Rust.
12//!
13//! Note: this crate does not use "Prefix Seek" mode which means that the prefix iterator
14//! will return keys not starting with the given prefix as well (as long as `key >= prefix`).
15//! To work around this we set an upper bound to the prefix successor.
16//! See https://github.com/facebook/rocksdb/wiki/Prefix-Seek-API-Changes for details.
17
18use crate::{other_io_err, DBAndColumns, DBKeyValue};
19use rocksdb::{DBIterator, Direction, IteratorMode, ReadOptions};
20use std::io;
21
22/// Instantiate iterators yielding `io::Result<DBKeyValue>`s.
23pub trait IterationHandler {
24	type Iterator: Iterator<Item = io::Result<DBKeyValue>>;
25
26	/// Create an `Iterator` over a `ColumnFamily` corresponding to the passed index. Takes
27	/// `ReadOptions` to allow configuration of the new iterator (see
28	/// https://github.com/facebook/rocksdb/blob/master/include/rocksdb/options.h#L1169).
29	fn iter(self, col: u32, read_opts: ReadOptions) -> Self::Iterator;
30	/// Create an `Iterator` over a `ColumnFamily` corresponding to the passed index. Takes
31	/// `ReadOptions` to allow configuration of the new iterator (see
32	/// https://github.com/facebook/rocksdb/blob/master/include/rocksdb/options.h#L1169).
33	/// The `Iterator` iterates over keys which start with the provided `prefix`.
34	fn iter_with_prefix(self, col: u32, prefix: &[u8], read_opts: ReadOptions) -> Self::Iterator;
35}
36
37impl<'a> IterationHandler for &'a DBAndColumns {
38	type Iterator = EitherIter<KvdbAdapter<DBIterator<'a>>, std::iter::Once<io::Result<DBKeyValue>>>;
39
40	fn iter(self, col: u32, read_opts: ReadOptions) -> Self::Iterator {
41		match self.cf(col as usize) {
42			Ok(cf) => EitherIter::A(KvdbAdapter(self.db.iterator_cf_opt(cf, read_opts, IteratorMode::Start))),
43			Err(e) => EitherIter::B(std::iter::once(Err(e))),
44		}
45	}
46
47	fn iter_with_prefix(self, col: u32, prefix: &[u8], read_opts: ReadOptions) -> Self::Iterator {
48		match self.cf(col as usize) {
49			Ok(cf) => EitherIter::A(KvdbAdapter(self.db.iterator_cf_opt(
50				cf,
51				read_opts,
52				IteratorMode::From(prefix, Direction::Forward),
53			))),
54			Err(e) => EitherIter::B(std::iter::once(Err(e))),
55		}
56	}
57}
58
59/// Small enum to avoid boxing iterators.
60pub enum EitherIter<A, B> {
61	A(A),
62	B(B),
63}
64
65impl<A, B, I> Iterator for EitherIter<A, B>
66where
67	A: Iterator<Item = I>,
68	B: Iterator<Item = I>,
69{
70	type Item = I;
71
72	fn next(&mut self) -> Option<Self::Item> {
73		match self {
74			Self::A(a) => a.next(),
75			Self::B(b) => b.next(),
76		}
77	}
78}
79
80/// A simple wrapper that adheres to the `kvdb` interface.
81pub struct KvdbAdapter<T>(T);
82
83impl<T> Iterator for KvdbAdapter<T>
84where
85	T: Iterator<Item = Result<(Box<[u8]>, Box<[u8]>), rocksdb::Error>>,
86{
87	type Item = io::Result<DBKeyValue>;
88
89	fn next(&mut self) -> Option<Self::Item> {
90		self.0
91			.next()
92			.map(|r| r.map_err(other_io_err).map(|(k, v)| (k.into_vec().into(), v.into())))
93	}
94}