kaspa_database/
utils.rs

1use crate::prelude::DB;
2use std::sync::Weak;
3use tempfile::TempDir;
4
5#[derive(Default)]
6pub struct DbLifetime {
7    weak_db_ref: Weak<DB>,
8    optional_tempdir: Option<TempDir>,
9}
10
11impl DbLifetime {
12    pub fn new(tempdir: TempDir, weak_db_ref: Weak<DB>) -> Self {
13        Self { optional_tempdir: Some(tempdir), weak_db_ref }
14    }
15
16    /// Tracks the DB reference and makes sure all strong refs are cleaned up
17    /// but does not remove the DB from disk when dropped.
18    pub fn without_destroy(weak_db_ref: Weak<DB>) -> Self {
19        Self { optional_tempdir: None, weak_db_ref }
20    }
21}
22
23impl Drop for DbLifetime {
24    fn drop(&mut self) {
25        for _ in 0..16 {
26            if self.weak_db_ref.strong_count() > 0 {
27                // Sometimes another thread is shuting-down and cleaning resources
28                std::thread::sleep(std::time::Duration::from_millis(1000));
29            } else {
30                break;
31            }
32        }
33        assert_eq!(self.weak_db_ref.strong_count(), 0, "DB is expected to have no strong references when lifetime is dropped");
34        if let Some(dir) = self.optional_tempdir.take() {
35            let options = rocksdb::Options::default();
36            let path_buf = dir.path().to_owned();
37            let path = path_buf.to_str().unwrap();
38            <rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>::destroy(&options, path)
39                .expect("DB is expected to be deletable since there are no references to it");
40        }
41    }
42}
43
44pub fn get_kaspa_tempdir() -> TempDir {
45    let global_tempdir = std::env::temp_dir();
46    let kaspa_tempdir = global_tempdir.join("rusty-kaspa");
47    std::fs::create_dir_all(kaspa_tempdir.as_path()).unwrap();
48    let db_tempdir = tempfile::tempdir_in(kaspa_tempdir.as_path()).unwrap();
49    db_tempdir
50}
51
52/// Creates a DB within a temp directory under `<OS SPECIFIC TEMP DIR>/kaspa-rust`
53/// Callers must keep the `TempDbLifetime` guard for as long as they wish the DB to exist.
54#[macro_export]
55macro_rules! create_temp_db {
56    ($conn_builder: expr) => {{
57        let db_tempdir = $crate::utils::get_kaspa_tempdir();
58        let db_path = db_tempdir.path().to_owned();
59        let db = $conn_builder.with_db_path(db_path).build().unwrap();
60        ($crate::utils::DbLifetime::new(db_tempdir, std::sync::Arc::downgrade(&db)), db)
61    }};
62}
63
64/// Creates a DB within the provided directory path.
65/// Callers must keep the `TempDbLifetime` guard for as long as they wish the DB instance to exist.
66#[macro_export]
67macro_rules! create_permanent_db {
68    ($db_path: expr, $conn_builder: expr) => {{
69        let db_dir = std::path::PathBuf::from($db_path);
70        if let Err(e) = std::fs::create_dir(db_dir.as_path()) {
71            match e.kind() {
72                std::io::ErrorKind::AlreadyExists => panic!("The directory {db_dir:?} already exists"),
73                _ => panic!("{e}"),
74            }
75        }
76        let db = $conn_builder.with_db_path(db_dir).build().unwrap();
77        ($crate::utils::DbLifetime::without_destroy(std::sync::Arc::downgrade(&db)), db)
78    }};
79}
80
81/// Loads an existing DB from the provided directory path.
82/// Callers must keep the `TempDbLifetime` guard for as long as they wish the DB instance to exist.
83#[macro_export]
84macro_rules! load_existing_db {
85    ($db_path: expr, $conn_builder: expr) => {{
86        let db_dir = std::path::PathBuf::from($db_path);
87        let db = $conn_builder.with_db_path(db_dir).with_create_if_missing(false).build().unwrap();
88        ($crate::utils::DbLifetime::without_destroy(std::sync::Arc::downgrade(&db)), db)
89    }};
90}
kaspa_database/utils.rs

kaspa_database/
utils.rs
