utxo_scanner/
database.rs

1extern crate rocksdb;
2
3use crate::BITCOIN_DATADIR;
4use rocksdb::{DBCompressionType, DB, DBCompactionStyle};
5use std::fs;
6use std::path::Path;
7use std::sync::Arc;
8
9pub trait Database {
10    fn new() -> Self;
11    fn put(&self, header: &String, sequence: &String);
12    fn get(&self, header: &String) -> Option<String>;
13    fn delete(&self, header: &String) -> bool;
14}
15
16#[derive(Clone)]
17pub struct RocksDB {
18    pub db: Arc<DB>,
19}
20
21impl Database for RocksDB {
22    fn new() -> Self {
23        // Create directory, if needed
24        let rocksdb_path = format!("{}/{}", BITCOIN_DATADIR.lock().unwrap(), "rocksdb");
25        if !Path::new(&rocksdb_path).exists() {
26            match fs::create_dir_all(&rocksdb_path) {
27                Ok(_) => {}
28                Err(_error) => panic!("Unable to create directory at {}.", rocksdb_path),
29            };
30        }
31
32        // Set options
33        let mut opts = rocksdb::Options::default();
34        opts.create_if_missing(true);
35        opts.set_compression_type(DBCompressionType::None);
36        //opts.set_compaction_style(DBCompactionStyle::Universal);
37        opts.increase_parallelism(4);
38
39        // Connect to database
40        let database = match DB::open(&opts, rocksdb_path.as_str()) {
41            Ok(r) => r,
42            Err(e) => panic!("Unable to open RocksDB at {}, error: {}", rocksdb_path, e),
43        };
44
45        // Return
46        RocksDB {
47            db: Arc::new(database),
48        }
49    }
50
51    fn put(&self, header: &String, sequence: &String) {
52        self.db.put(header.as_bytes(), sequence.as_bytes()).unwrap();
53    }
54
55    fn get(&self, header: &String) -> Option<String> {
56        let sequence = match self.db.get(header.as_bytes()) {
57            Ok(Some(r)) => String::from_utf8(r).unwrap(),
58            Ok(None) => return None,
59            Err(e) => panic!(
60                "Received database error when trying to retrieve sequence, error: {}",
61                e
62            ),
63        };
64
65        Some(sequence)
66    }
67
68    fn delete(&self, header: &String) -> bool {
69        self.db.delete(header.as_bytes()).is_ok()
70    }
71}
utxo_scanner/database.rs

utxo_scanner/
database.rs
