use anyhow::{Context, Result};
use bitcoin_hashes::hex::ToHex;
use bitcoin_hashes::Hash;
use rocksdb::perf::get_memory_usage_stats;
use rocksdb::ColumnFamily;
use std::path::Path;
use std::path::PathBuf;
use std::str::from_utf8;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::{Arc, Condvar, Mutex};
use std::thread;
use std::time::Duration;
use tokio::sync::mpsc;
use crate::def::DATABASE_VERSION;
use crate::indexes::headerindex::HeaderRow;
use crate::indexes::heightindex::HeightIndexRow;
use crate::indexes::inputindex::InputIndexRow;
use crate::indexes::outputindex::OutputIndexRow;
use crate::indexes::outputtokenindex::OutputTokenIndexRow;
use crate::indexes::outtoscriptindex::OutToScripthashIndex;
use crate::indexes::scripthashindex::{OutputFlags, ScriptHashIndexRow};
use crate::indexes::tokenoutputindex::TokenOutputIndexRow;
use crate::indexes::unspentindex::UnspentIndexRow;
use crate::indexes::DBRow;
use crate::metrics::Metrics;
use crate::signal::Waiter;
use crate::util::Bytes;
use crate::writebatch::SpentUpdateBatch;
use crate::writebatch::WriteBatch;
pub const META_CF: &str = "meta";
pub const METADATA_LAST_INDEXED_BLOCK: &[u8] = b"L";
pub const METADATA_DB_VERSION: &[u8] = b"VER";
pub const COLUMN_FAMILIES: &[&str] = &[
HeaderRow::CF,
HeightIndexRow::CF,
TokenOutputIndexRow::CF,
OutputTokenIndexRow::CF,
InputIndexRow::CF,
OutputIndexRow::CF,
ScriptHashIndexRow::CF,
UnspentIndexRow::CF,
OutToScripthashIndex::CF,
META_CF,
];
#[derive(Eq, PartialEq, Clone, Copy, Debug)]
pub enum DBContents {
ConfirmedIndex,
MempoolIndex,
UnspentIndex,
Other,
}
#[derive(Clone)]
pub struct Row {
pub key: Bytes,
pub value: Bytes,
}
pub fn db_encode<T: serde::Serialize + ?Sized>(value: &T) -> Result<Vec<u8>> {
postcard::to_allocvec(value).map_err(anyhow::Error::new)
}
pub fn db_decode<T: serde::de::DeserializeOwned>(value: &[u8]) -> Result<T> {
postcard::from_bytes(value).map_err(anyhow::Error::new)
}
fn default_opts() -> rocksdb::Options {
let mut opts = rocksdb::Options::default();
opts.set_keep_log_file_num(10);
opts.set_max_open_files(16);
opts.set_compaction_style(rocksdb::DBCompactionStyle::Level);
opts.set_compression_type(rocksdb::DBCompressionType::None);
opts.set_target_file_size_base(256 << 20);
opts.set_write_buffer_size(256 << 20);
opts.set_disable_auto_compactions(false);
opts.set_advise_random_on_open(true);
opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(32));
opts
}
pub struct DBStore {
pub contents: DBContents,
db: Arc<rocksdb::DB>,
stats_thread: Option<thread::JoinHandle<()>>,
stats_thread_kill: Arc<(Mutex<bool>, Condvar)>,
path: PathBuf,
}
impl DBStore {
pub fn open(
contents: DBContents,
path: &Path,
metrics: &Metrics,
force_is_new_db: bool,
) -> Result<Self> {
info!("Opening DB at {:?}", path);
let mut db_opts = default_opts();
db_opts.create_if_missing(true);
db_opts.create_missing_column_families(true);
let is_new_db = force_is_new_db || !path.exists();
let db = rocksdb::DB::open_cf_descriptors(&db_opts, path, Self::create_cf_descriptors())?;
let live_files = db.live_files()?;
info!(
"{:?}: {} SST files, {} GB, {} Grows",
path,
live_files.len(),
live_files.iter().map(|f| f.size).sum::<usize>() as f64 / 1e9,
live_files.iter().map(|f| f.num_entries).sum::<u64>() as f64 / 1e9
);
let mut store = DBStore {
contents,
db: Arc::new(db),
stats_thread: None,
stats_thread_kill: Arc::new((Mutex::new(false), Condvar::new())),
path: path.to_path_buf(),
};
let version_marker = version_marker();
if is_new_db {
let b = WriteBatch::new();
b.insert(META_CF, rayon::iter::once(version_marker.clone()));
store.write_batch(&b);
store.flush().expect("Flush on new database failed")
}
if !store.exists_blocking(META_CF, &version_marker.key) {
return Err(anyhow!("Database does not have version marker set."));
}
store.start_stats_thread(metrics);
Ok(store)
}
fn create_cf_descriptors() -> Vec<rocksdb::ColumnFamilyDescriptor> {
COLUMN_FAMILIES
.iter()
.map(|&name| rocksdb::ColumnFamilyDescriptor::new(name, default_opts()))
.collect()
}
fn start_stats_thread(&mut self, metrics: &Metrics) {
static DBINSTANCE_COUNT: AtomicUsize = AtomicUsize::new(0);
let i = DBINSTANCE_COUNT.fetch_add(1, Ordering::Relaxed);
let mem_table_total = metrics.gauge_int(prometheus::Opts::new(
format!("rostrum_rockdb_mem_table_total_{}", i),
"Rockdb approximate memory usage of all the mem-tables".to_string(),
));
let mem_table_unflushed = metrics.gauge_int(prometheus::Opts::new(
format!("rostrum_rockdb_mem_table_unflushed_{}", i),
"Rocksdb approximate usage of un-flushed mem-tables".to_string(),
));
let mem_table_readers_total = metrics.gauge_int(prometheus::Opts::new(
format!("rostrum_rockdb_mem_table_readers_total_{}", i),
"Rocksdb approximate memory usage of all the table readers".to_string(),
));
let dbptr = Arc::clone(&self.db);
let kill = Arc::clone(&self.stats_thread_kill);
self.stats_thread = Some(crate::thread::spawn("dbstats", move || {
let (killthread, cvar) = &*kill;
loop {
Waiter::shutdown_check()?;
let k = killthread.lock().unwrap();
let result = cvar.wait_timeout(k, Duration::from_secs(5)).unwrap();
if *result.0 {
mem_table_total.set(0);
mem_table_unflushed.set(0);
mem_table_readers_total.set(0);
return Ok(());
}
let mem_usage = get_memory_usage_stats(Some(&[&*dbptr]), None);
if let Ok(usage) = mem_usage {
mem_table_total.set(usage.mem_table_total as i64);
mem_table_unflushed.set(usage.mem_table_unflushed as i64);
mem_table_readers_total.set(usage.mem_table_readers_total as i64)
}
}
}));
}
pub fn destroy(path: &Path) {
match rocksdb::DB::destroy(&default_opts(), path) {
Ok(_) => debug!("Database '{}' deleted", path.as_os_str().to_string_lossy()),
Err(err) => info!("Clould not delete database: {}", err),
}
}
pub(crate) fn get_blocking(&self, cf_name: &'static str, key: &[u8]) -> Option<Vec<u8>> {
let cf = self
.db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
self.db.get_cf_opt(cf, key, &opts).expect("get_cf failed")
}
pub(crate) async fn get(
&self,
cf_name: &'static str,
key: Vec<u8>,
) -> (Vec<u8>, Option<Vec<u8>>) {
let db = self.db.clone();
tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
let value = db.get_cf_opt(cf, &key, &opts).expect("get_cf failed");
(key, value)
})
.await
.unwrap()
}
pub(crate) async fn multi_get(
&self,
cf_name: &'static str,
keys: Vec<Vec<u8>>,
) -> Vec<Option<Vec<u8>>> {
let db = self.db.clone();
tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
let keys: Vec<(&ColumnFamily, Vec<u8>)> = keys.into_iter().map(|k| (cf, k)).collect();
db.multi_get_cf_opt(keys, &opts)
.into_iter()
.map(|result| result.expect("multi_get_cf failed"))
.collect::<Vec<Option<Vec<u8>>>>()
})
.await
.unwrap()
}
pub(crate) fn exists_blocking(&self, cf_name: &'static str, key: &[u8]) -> bool {
let cf = self
.db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
self.db.get_pinned_cf_opt(cf, key, &opts).unwrap().is_some()
}
#[cfg(nexa)]
pub(crate) async fn exists(&self, cf_name: &'static str, key: Vec<u8>) -> bool {
let db = self.db.clone();
tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
db.get_pinned_cf_opt(cf, key, &opts).unwrap().is_some()
})
.await
.unwrap()
}
pub(crate) async fn exists_32bit_key(&self, cf_name: &'static str, key: [u8; 32]) -> bool {
let db = self.db.clone();
tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
db.get_pinned_cf_opt(cf, key, &opts).unwrap().is_some()
})
.await
.unwrap()
}
pub(crate) async fn exists_multi_32bit_keys(
&self,
cf_name: &'static str,
keys: Vec<[u8; 32]>,
) -> Vec<bool> {
let db = self.db.clone();
tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
let values = db.batched_multi_get_cf_opt(cf, &keys, false, &opts);
values
.into_iter()
.map(|result| result.unwrap().is_some())
.collect()
})
.await
.unwrap()
}
pub(crate) async fn scan(
&self,
cf_name: &'static str,
prefix: Vec<u8>,
bitmask_scan: Option<(u8 /* bitmask */, usize /* position */)>,
) -> (
tokio::task::JoinHandle<Result<()>>,
impl futures::stream::Stream<Item = Row> + Unpin,
) {
self.scan_inner(cf_name, prefix, bitmask_scan, rocksdb::Direction::Forward)
.await
}
pub(crate) async fn rscan(
&self,
cf_name: &'static str,
prefix: Vec<u8>,
bitmask_scan: Option<(u8 /* bitmask */, usize /* position */)>,
) -> (
tokio::task::JoinHandle<Result<()>>,
impl futures::stream::Stream<Item = Row> + Unpin,
) {
self.scan_inner(cf_name, prefix, bitmask_scan, rocksdb::Direction::Reverse)
.await
}
pub(crate) async fn scan_inner(
&self,
cf_name: &'static str,
prefix: Vec<u8>,
bitmask_scan: Option<(u8 /* bitmask */, usize /* position */)>,
direction: rocksdb::Direction,
) -> (
tokio::task::JoinHandle<Result<()>>,
impl futures::stream::Stream<Item = Row> + Unpin,
) {
let prefix_len = prefix.len();
assert!(
prefix_len >= 32,
"scan uses prefix extractor with keys >= 32 bytes"
);
if let Some((_, pos)) = bitmask_scan {
assert!(
pos >= 32,
"bitmask scan position must be the 33rd byte or later"
);
}
let db = self.db.clone();
let (tx, rx) = mpsc::channel(100);
let hard_timeout = Duration::from_secs(30);
let task = tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let start_key = if matches!(direction, rocksdb::Direction::Reverse) {
let mut upper_bound = prefix.clone();
if upper_bound.len() < 128 {
upper_bound.resize(128, 0xFF);
}
upper_bound
} else {
prefix.clone()
};
let mode = rocksdb::IteratorMode::From(&start_key, direction);
let mut opts = rocksdb::ReadOptions::default();
opts.set_prefix_same_as_start(true);
opts.set_verify_checksums(false);
tokio::runtime::Handle::current().block_on({
let db = Arc::clone(&db);
let prefix = prefix.clone();
async move {
for x in db.iterator_cf_opt(cf, opts, mode) {
let (key, value) = x.expect("failed to read from db");
if let Some((bitmask, position)) = bitmask_scan {
if key[position] & bitmask != 0 {
tx.send_timeout(Row { key, value }, hard_timeout).await?
}
continue;
}
if prefix_len == 32 || key[32..].starts_with(&prefix[32..]) {
tx.send_timeout(Row { key, value }, hard_timeout).await?
} else {
break;
}
}
Ok(())
}
})
});
let stream = futures::stream::unfold(rx, move |mut rx| async move {
rx.recv().await.map(|row| (row, rx))
});
(task, Box::pin(stream))
}
pub async fn scan_values(
&self,
cf_name: &'static str,
prefix: Vec<u8>,
) -> (
tokio::task::JoinHandle<Result<()>>,
impl futures::stream::Stream<Item = Box<[u8]>> + Unpin,
) {
let prefix_len = prefix.len();
assert!(
prefix_len >= 32,
"scan uses prefix extractor with keys >= 32 bytes"
);
let db = self.db.clone();
let (tx, rx) = mpsc::channel(100);
let hard_timeout = Duration::from_secs(30);
let task = tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mode = rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward);
let mut opts = rocksdb::ReadOptions::default();
opts.set_prefix_same_as_start(true);
opts.set_verify_checksums(false);
tokio::runtime::Handle::current().block_on({
let db = Arc::clone(&db);
let prefix = prefix.clone();
async move {
for x in db.iterator_cf_opt(cf, opts, mode) {
let (key, value) = x.expect("failed to read from db");
if prefix_len == 32 || key[32..].starts_with(&prefix[32..]) {
tx.send_timeout(value, hard_timeout).await?
} else {
break;
}
}
Ok(())
}
})
});
let stream = futures::stream::unfold(rx, move |mut rx| async move {
rx.recv().await.map(|value| (value, rx))
});
(task, Box::pin(stream))
}
pub(crate) async fn iter(
&self,
cf_name: &'static str,
fill_cache: bool,
) -> impl futures::stream::Stream<Item = Row> + Unpin {
let db = self.db.clone();
let (tx, rx) = mpsc::channel(1000);
tokio::task::spawn_blocking(move || {
let cf = db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
let mut opts = rocksdb::ReadOptions::default();
opts.fill_cache(fill_cache);
let db_iter = db.iterator_cf_opt(cf, opts, rocksdb::IteratorMode::Start);
for kv in db_iter {
let (key, value) = kv.expect("iterator_cf_opt error");
if tx.blocking_send(Row { key, value }).is_err() {
break;
}
}
});
let stream = futures::stream::unfold(rx, move |mut rx| async move {
rx.recv().await.map(|row| (row, rx))
});
Box::pin(stream)
}
pub(crate) fn write_batch(&self, src: &WriteBatch) {
let mut dst = rocksdb::WriteBatch::default();
src.cf_to_vec_tuples()
.into_iter()
.for_each(|(cf_name, rows)| {
let cf = self
.db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
rows.lock()
.unwrap()
.drain(..)
.for_each(|r| dst.put_cf(cf, r.key, r.value));
});
let mut opts = rocksdb::WriteOptions::new();
opts.set_sync(false);
opts.disable_wal(true);
self.db.write_opt(dst, &opts).expect("write_opt failed");
}
pub(crate) fn erase_batch(&self, src: &WriteBatch) {
let mut dst = rocksdb::WriteBatch::default();
src.cf_to_vec_tuples()
.into_iter()
.for_each(|(cf_name, rows)| {
let cf = self
.db
.cf_handle(cf_name)
.unwrap_or_else(|| panic!("missing cf {}", cf_name));
rows.lock()
.unwrap()
.drain(..)
.for_each(|r| dst.delete_cf(cf, r.key));
});
let mut opts = rocksdb::WriteOptions::new();
opts.set_sync(false);
opts.disable_wal(true);
self.db.write_opt(dst, &opts).expect("write_opt failed");
}
pub fn take_and_delete<K: AsRef<[u8]>>(
&self,
cf: &rocksdb::ColumnFamily,
key: K,
batch: &mut rocksdb::WriteBatch,
) -> Option<Vec<u8>> {
let key_ref = key.as_ref();
if let Ok(Some(value)) = self.db.get_cf(cf, key_ref) {
batch.delete_cf(cf, key_ref);
Some(value)
} else {
None
}
}
pub(crate) fn update_spends(&self, src: &SpentUpdateBatch) {
let mut to_update = rocksdb::WriteBatch::default();
let spending_metadata = src.take_spending_metadata();
if spending_metadata.is_empty() {
return;
}
let scripthash_cf = self.db.cf_handle(OutToScripthashIndex::CF).unwrap();
let unspent_cf = self.db.cf_handle(UnspentIndexRow::CF).unwrap();
let scripthashindex_cf = self.db.cf_handle(ScriptHashIndexRow::CF).unwrap();
let mut entries_needing_lookup = spending_metadata;
let indices_to_remove: Vec<usize> = entries_needing_lookup
.iter()
.enumerate()
.filter_map(|(i, entry)| if entry.4.is_some() { Some(i) } else { None })
.collect();
let mut entries_with_data = Vec::with_capacity(indices_to_remove.len());
for &idx in indices_to_remove.iter().rev() {
entries_with_data.push(entries_needing_lookup.remove(idx));
}
for (oph, txid, vin, height, prefilled) in entries_with_data {
let prefilled = prefilled.unwrap();
let key = UnspentIndexRow::to_key(&prefilled.scripthash, &oph);
to_update.delete_cf(unspent_cf, key);
let out_to_scripthash_key = OutToScripthashIndex::to_key(&oph);
to_update.delete_cf(scripthash_cf, out_to_scripthash_key);
let spending_entry = ScriptHashIndexRow::new_spending(
&prefilled.scripthash,
&oph,
if prefilled.has_token {
OutputFlags::SpentHasTokens
} else {
OutputFlags::SpentNone
},
txid.into_inner(),
vin,
height,
);
if src.is_erase_only() {
to_update.delete_cf(scripthashindex_cf, spending_entry.to_row().key);
} else {
to_update.put_cf(
scripthashindex_cf,
spending_entry.to_row().key,
spending_entry.to_row().value,
);
}
}
if !entries_needing_lookup.is_empty() {
let lookup_keys: Vec<Vec<u8>> = entries_needing_lookup
.iter()
.map(|(oph, _, _, _, _)| OutToScripthashIndex::to_key(oph))
.collect();
let mut opts = rocksdb::ReadOptions::default();
opts.set_verify_checksums(false);
opts.fill_cache(false);
for ((oph, txid, vin, height, _scripthash_and_has_token), result) in
entries_needing_lookup.into_iter().zip(
self.db
.batched_multi_get_cf_opt(scripthash_cf, &lookup_keys, false, &opts)
.into_iter(),
)
{
let out_to_scripthash_opt = result.expect("batched_multi_get_cf failed");
if let Some(out_to_scripthash_value) = out_to_scripthash_opt {
let out_to_scripthash_row = OutToScripthashIndex::from_row(&Row {
key: OutToScripthashIndex::to_key(&oph).into_boxed_slice(),
value: out_to_scripthash_value.to_vec().into_boxed_slice(),
});
let scripthash = out_to_scripthash_row.scripthash();
let has_token = out_to_scripthash_row.has_token();
let key = UnspentIndexRow::to_key(&scripthash, &oph);
to_update.delete_cf(unspent_cf, key);
let out_to_scripthash_key = OutToScripthashIndex::to_key(&oph);
to_update.delete_cf(scripthash_cf, out_to_scripthash_key);
let spending_entry = ScriptHashIndexRow::new_spending(
&scripthash,
&oph,
if has_token {
OutputFlags::SpentHasTokens
} else {
OutputFlags::SpentNone
},
txid.into_inner(),
vin,
height,
);
if src.is_erase_only() {
to_update.delete_cf(scripthashindex_cf, spending_entry.to_row().key);
} else {
to_update.put_cf(
scripthashindex_cf,
spending_entry.to_row().key,
spending_entry.to_row().value,
);
}
} else {
if src.is_erase_only() {
trace!(
"spent update: did not have oph {} in db {:?} (erase_only mode)",
oph.to_hex(),
self.contents
);
} else {
warn!(
"spent update: did not have oph {} in db {:?}",
oph.to_hex(),
self.contents
);
}
}
}
}
let mut opts = rocksdb::WriteOptions::new();
opts.set_sync(false);
opts.disable_wal(true);
self.db
.write_opt(to_update, &opts)
.expect("write_opt failed");
}
}
impl DBStore {
pub fn flush(&self) -> Result<()> {
for name in COLUMN_FAMILIES {
let cf = self.db.cf_handle(name).expect("missing CF");
self.db.flush_cf(cf).context("CF flush failed")?;
}
Ok(())
}
}
impl Drop for DBStore {
fn drop(&mut self) {
info!("Closing DB {:?}", self.path);
let (flag, cvar) = &*self.stats_thread_kill;
*flag.lock().unwrap() = true;
cvar.notify_one();
if let Err(e) = self.flush() {
warn!("Flushing {:?} failed: {e}", self.path)
}
self.stats_thread.take().map(thread::JoinHandle::join);
info!("Done closing DB {:?}", self.path);
}
}
pub fn version_marker() -> Row {
Row {
key: METADATA_DB_VERSION.to_vec().into_boxed_slice(),
value: DATABASE_VERSION.as_bytes().to_vec().into_boxed_slice(),
}
}
pub fn is_compatible_version(path: &Path, metrics: &Metrics) -> bool {
let store = match DBStore::open(DBContents::Other, path, metrics, false) {
Ok(store) => store,
Err(e) => {
if e.to_string().contains("Snappy") {
warn!("Old database using snappy compression ({})", e);
return false;
}
warn!("Unknown database issue: {}", e);
return false;
}
};
if let Err(e) = store.db.flush() {
if e.to_string().contains("ZSTD") {
warn!("Old database using ZSTD compression ({})", e);
return false;
}
panic!("Unknown database issue: {}", e);
}
let version = store.get_blocking(META_CF, &version_marker().key);
match version {
Some(v) => match from_utf8(&v) {
Ok(v) => {
if v != DATABASE_VERSION {
info!("Incompatbile db version. DB is version {v}, Rostrum expects {DATABASE_VERSION}");
false
} else {
info!("DB version is {v}");
true
}
}
Err(e) => {
warn!("utf-8 error reading database version: {e}");
false
}
},
None => {
info!("Version flag not found in database.");
false
}
}
}