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    tempfile::tempdir_in(kaspa_tempdir.as_path()).unwrap()
49}
50
51/// Creates a DB within a temp directory under `<OS SPECIFIC TEMP DIR>/rusty-kaspa`
52/// Callers must keep the `TempDbLifetime` guard for as long as they wish the DB to exist.
53#[macro_export]
54macro_rules! create_temp_db {
55    ($conn_builder: expr) => {{
56        let db_tempdir = $crate::utils::get_kaspa_tempdir();
57        let db_path = db_tempdir.path().to_owned();
58        let db = $conn_builder.with_db_path(db_path).build().unwrap();
59        ($crate::utils::DbLifetime::new(db_tempdir, std::sync::Arc::downgrade(&db)), db)
60    }};
61}
62
63/// Creates a DB within the provided directory path.
64/// Callers must keep the `TempDbLifetime` guard for as long as they wish the DB instance to exist.
65#[macro_export]
66macro_rules! create_permanent_db {
67    ($db_path: expr, $conn_builder: expr) => {{
68        let db_dir = std::path::PathBuf::from($db_path);
69        if let Err(e) = std::fs::create_dir(db_dir.as_path()) {
70            match e.kind() {
71                std::io::ErrorKind::AlreadyExists => panic!("The directory {db_dir:?} already exists"),
72                _ => panic!("{e}"),
73            }
74        }
75        let db = $conn_builder.with_db_path(db_dir).build().unwrap();
76        ($crate::utils::DbLifetime::without_destroy(std::sync::Arc::downgrade(&db)), db)
77    }};
78}
79
80/// Loads an existing DB from the provided directory path.
81/// Callers must keep the `TempDbLifetime` guard for as long as they wish the DB instance to exist.
82#[macro_export]
83macro_rules! load_existing_db {
84    ($db_path: expr, $conn_builder: expr) => {{
85        let db_dir = std::path::PathBuf::from($db_path);
86        let db = $conn_builder.with_db_path(db_dir).with_create_if_missing(false).build().unwrap();
87        ($crate::utils::DbLifetime::without_destroy(std::sync::Arc::downgrade(&db)), db)
88    }};
89}
kaspa_database/utils.rs

kaspa_database/
utils.rs
