typed_store/rocks/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5pub mod errors;
6pub(crate) mod options;
7pub(crate) mod rocks_util;
8pub(crate) mod safe_iter;
9
10use std::{
11    collections::HashSet,
12    ffi::CStr,
13    path::{Path, PathBuf},
14    sync::Arc,
15    time::Duration,
16};
17
18use backoff::backoff::Backoff;
19use iota_macros::nondeterministic;
20use rocksdb::{
21    AsColumnFamilyRef, ColumnFamilyDescriptor, Error, MultiThreaded, properties,
22    properties::num_files_at_level,
23};
24use tracing::warn;
25use typed_store_error::TypedStoreError;
26
27pub use crate::{
28    database::{DBBatch, DBMap, MetricConf},
29    rocks::options::{
30        BulkIngestionOptions, DBMapTableConfigMap, DBOptions, ReadWriteOptions,
31        bulk_ingestion_options, bulk_ingestion_write_options, default_db_options, list_tables,
32        read_size_from_env,
33    },
34};
35use crate::{
36    database::{Database, Storage},
37    metrics::DBMetrics,
38    rocks::errors::typed_store_err_from_rocks_err,
39};
40
41// TODO: remove this after Rust rocksdb has the TOTAL_BLOB_FILES_SIZE property
42// built-in. From https://github.com/facebook/rocksdb/blob/bd80433c73691031ba7baa65c16c63a83aef201a/include/rocksdb/db.h#L1169
43const ROCKSDB_PROPERTY_TOTAL_BLOB_FILES_SIZE: &CStr =
44    unsafe { CStr::from_bytes_with_nul_unchecked("rocksdb.total-blob-file-size\0".as_bytes()) };
45
46const METRICS_ERROR: i64 = -1;
47
48const DB_CORRUPTED_KEY: &[u8] = b"db_corrupted";
49
50#[cfg(test)]
51mod tests;
52
53#[derive(Debug)]
54pub(crate) struct RocksDB {
55    pub(crate) underlying: rocksdb::DBWithThreadMode<MultiThreaded>,
56    /// Names of all column families opened on this database.
57    pub(crate) cf_names: Vec<String>,
58}
59
60impl Drop for RocksDB {
61    fn drop(&mut self) {
62        self.underlying.cancel_all_background_work(/* wait */ true);
63    }
64}
65
66pub(crate) fn rocks_cf<'a>(
67    rocks_db: &'a RocksDB,
68    cf_name: &str,
69) -> Arc<rocksdb::BoundColumnFamily<'a>> {
70    rocks_db
71        .underlying
72        .cf_handle(cf_name)
73        .expect("Map-keying column family should have been checked at DB creation")
74}
75
76// Check if the database is corrupted, and if so, panic.
77// If the corrupted key is not set, we set it to [1].
78pub fn check_and_mark_db_corruption(path: &Path) -> Result<(), String> {
79    // rocksdb spawns its background threads when the DB is opened; under the
80    // simulator that open must run off the test thread, or the threads are
81    // scheduled against the simulated clock and then run against the real one,
82    // aborting periodic-task registration. Only the open needs the guard — the
83    // get/put below spawn no threads. See `open_cf_opts`. No-op outside msim.
84    let db = nondeterministic!(rocksdb::DB::open_default(path)).map_err(|e| e.to_string())?;
85
86    db.get(DB_CORRUPTED_KEY)
87        .map_err(|e| format!("Failed to open database: {e}"))
88        .and_then(|value| match value {
89            Some(v) if v[0] == 1 => Err(
90                "Database is corrupted, please remove the current database and start clean!"
91                    .to_string(),
92            ),
93            Some(_) => Ok(()),
94            None => db
95                .put(DB_CORRUPTED_KEY, [1])
96                .map_err(|e| format!("Failed to set corrupted key in database: {e}")),
97        })?;
98
99    Ok(())
100}
101
102pub fn unmark_db_corruption(path: &Path) -> Result<(), Error> {
103    // See `check_and_mark_db_corruption` for why the open runs off the test thread.
104    nondeterministic!(rocksdb::DB::open_default(path))?.put(DB_CORRUPTED_KEY, [0])
105}
106
107/// Opens a database with options, and a number of column families with
108/// individual options that are created if they do not exist.
109#[tracing::instrument(level="debug", skip_all, fields(path = ?path.as_ref()), err)]
110pub fn open_cf_opts<P: AsRef<Path>>(
111    path: P,
112    db_options: Option<rocksdb::Options>,
113    metric_conf: MetricConf,
114    opt_cfs: &[(&str, rocksdb::Options)],
115) -> Result<Arc<Database>, TypedStoreError> {
116    let path = path.as_ref();
117    // In the simulator, we intercept the wall clock in the test thread only. This
118    // causes problems because rocksdb uses the simulated clock when creating
119    // its background threads, but then those threads see the real wall clock
120    // (because they are not the test thread), which causes rocksdb to panic.
121    // The `nondeterministic` macro evaluates expressions in new threads, which
122    // resolves the issue.
123    //
124    // This is a no-op in non-simulator builds.
125
126    let cfs = populate_missing_cfs(opt_cfs, path).map_err(typed_store_err_from_rocks_err)?;
127    nondeterministic!({
128        let mut options = db_options.unwrap_or_else(|| default_db_options().options);
129        options.create_if_missing(true);
130        options.create_missing_column_families(true);
131        let cf_names: Vec<String> = cfs.iter().map(|(name, _)| name.clone()).collect();
132        let rocksdb = {
133            rocksdb::DBWithThreadMode::<MultiThreaded>::open_cf_descriptors(
134                &options,
135                path,
136                cfs.into_iter()
137                    .map(|(name, opts)| ColumnFamilyDescriptor::new(name, opts)),
138            )
139            .map_err(typed_store_err_from_rocks_err)?
140        };
141        Ok(Arc::new(Database::new(
142            Storage::Rocks(RocksDB {
143                underlying: rocksdb,
144                cf_names,
145            }),
146            metric_conf,
147        )))
148    })
149}
150
151/// Opens a database with options, and a number of column families with
152/// individual options that are created if they do not exist.
153pub fn open_cf_opts_secondary<P: AsRef<Path>>(
154    primary_path: P,
155    secondary_path: Option<P>,
156    db_options: Option<rocksdb::Options>,
157    metric_conf: MetricConf,
158    opt_cfs: &[(&str, rocksdb::Options)],
159) -> Result<Arc<Database>, TypedStoreError> {
160    let primary_path = primary_path.as_ref();
161    let secondary_path = secondary_path.as_ref().map(|p| p.as_ref());
162    // See comment above for explanation of why nondeterministic is necessary here.
163    nondeterministic!({
164        // Customize database options
165        let mut options = db_options.unwrap_or_else(|| default_db_options().options);
166
167        fdlimit::raise_fd_limit();
168        // This is a requirement by RocksDB when opening as secondary
169        options.set_max_open_files(-1);
170
171        let mut opt_cfs: std::collections::HashMap<_, _> = opt_cfs.iter().cloned().collect();
172        let cfs = rocksdb::DBWithThreadMode::<MultiThreaded>::list_cf(&options, primary_path)
173            .ok()
174            .unwrap_or_default();
175
176        let default_db_options = default_db_options();
177        // Add CFs not explicitly listed
178        for cf_key in cfs.iter() {
179            if !opt_cfs.contains_key(&cf_key[..]) {
180                opt_cfs.insert(cf_key, default_db_options.options.clone());
181            }
182        }
183
184        let primary_path = primary_path.to_path_buf();
185        let secondary_path = secondary_path.map(|q| q.to_path_buf()).unwrap_or_else(|| {
186            let mut s = primary_path.clone();
187            s.pop();
188            s.push("SECONDARY");
189            s.as_path().to_path_buf()
190        });
191
192        let rocksdb = {
193            options.create_if_missing(true);
194            options.create_missing_column_families(true);
195            let db = rocksdb::DBWithThreadMode::<MultiThreaded>::open_cf_descriptors_as_secondary(
196                &options,
197                &primary_path,
198                &secondary_path,
199                opt_cfs
200                    .iter()
201                    .map(|(name, opts)| ColumnFamilyDescriptor::new(*name, (*opts).clone())),
202            )
203            .map_err(typed_store_err_from_rocks_err)?;
204            db.try_catch_up_with_primary()
205                .map_err(typed_store_err_from_rocks_err)?;
206            db
207        };
208        let cf_names: Vec<String> = opt_cfs.keys().map(|name| name.to_string()).collect();
209        Ok(Arc::new(Database::new(
210            Storage::Rocks(RocksDB {
211                underlying: rocksdb,
212                cf_names,
213            }),
214            metric_conf,
215        )))
216    })
217}
218
219// Drops a database if there is no other handle to it, with retries and timeout.
220#[cfg(msim)]
221pub async fn safe_drop_db(path: PathBuf, timeout: Duration) -> Result<(), rocksdb::Error> {
222    // The destroy fails until rocksdb's background threads release the file
223    // lock, which happens on the real clock. Retrying on the simulated clock
224    // would consume a machine-load-dependent amount of simulated time (and rng,
225    // through the backoff jitter), breaking simtest determinism, so retry on a
226    // real thread with real sleeps instead.
227    nondeterministic!({
228        let deadline = std::time::Instant::now() + timeout;
229        loop {
230            match rocksdb::DB::destroy(&rocksdb::Options::default(), path.clone()) {
231                Ok(()) => return Ok(()),
232                Err(err) => {
233                    if std::time::Instant::now() >= deadline {
234                        return Err(err);
235                    }
236                    std::thread::sleep(Duration::from_millis(100));
237                }
238            }
239        }
240    })
241}
242
243// Drops a database if there is no other handle to it, with retries and timeout.
244#[cfg(not(msim))]
245pub async fn safe_drop_db(path: PathBuf, timeout: Duration) -> Result<(), rocksdb::Error> {
246    let mut backoff = backoff::ExponentialBackoff {
247        max_elapsed_time: Some(timeout),
248        ..Default::default()
249    };
250    loop {
251        match rocksdb::DB::destroy(&rocksdb::Options::default(), path.clone()) {
252            Ok(()) => return Ok(()),
253            Err(err) => match backoff.next_backoff() {
254                Some(duration) => tokio::time::sleep(duration).await,
255                None => return Err(err),
256            },
257        }
258    }
259}
260
261fn populate_missing_cfs(
262    input_cfs: &[(&str, rocksdb::Options)],
263    path: &Path,
264) -> Result<Vec<(String, rocksdb::Options)>, rocksdb::Error> {
265    let mut cfs = vec![];
266    let input_cf_index: HashSet<_> = input_cfs.iter().map(|(name, _)| *name).collect();
267    let existing_cfs =
268        rocksdb::DBWithThreadMode::<MultiThreaded>::list_cf(&rocksdb::Options::default(), path)
269            .ok()
270            .unwrap_or_default();
271
272    for cf_name in existing_cfs {
273        if !input_cf_index.contains(&cf_name[..]) {
274            cfs.push((cf_name, rocksdb::Options::default()));
275        }
276    }
277    cfs.extend(
278        input_cfs
279            .iter()
280            .map(|(name, opts)| (name.to_string(), (*opts).clone())),
281    );
282    Ok(cfs)
283}
284
285/// RocksDB-specific methods on `DBMap`. These are kept separate from the
286/// generic impl in `database.rs` because they directly access RocksDB
287/// internals and have no meaning for other storage backends.
288impl<K, V> DBMap<K, V> {
289    fn get_rocksdb_int_property(
290        rocksdb: &RocksDB,
291        cf: &impl AsColumnFamilyRef,
292        property_name: &CStr,
293    ) -> Result<i64, TypedStoreError> {
294        match rocksdb.underlying.property_int_value_cf(cf, property_name) {
295            Ok(Some(value)) => Ok(value.min(i64::MAX as u64).try_into().unwrap_or_default()),
296            Ok(None) => Ok(0),
297            Err(e) => Err(TypedStoreError::RocksDB(e.into_string())),
298        }
299    }
300
301    pub(crate) fn report_rocksdb_metrics(
302        database: &Arc<Database>,
303        cf_name: &str,
304        db_metrics: &Arc<DBMetrics>,
305    ) {
306        let Storage::Rocks(rocksdb) = &database.storage else {
307            return;
308        };
309
310        let Some(cf) = rocksdb.underlying.cf_handle(cf_name) else {
311            warn!(
312                "unable to report metrics for cf {cf_name:?} in db {:?}",
313                database.db_name()
314            );
315            return;
316        };
317
318        db_metrics
319            .cf_metrics
320            .rocksdb_total_sst_files_size
321            .with_label_values(&[cf_name])
322            .set(
323                Self::get_rocksdb_int_property(rocksdb, &cf, properties::TOTAL_SST_FILES_SIZE)
324                    .unwrap_or(METRICS_ERROR),
325            );
326        db_metrics
327            .cf_metrics
328            .rocksdb_total_blob_files_size
329            .with_label_values(&[cf_name])
330            .set(
331                Self::get_rocksdb_int_property(
332                    rocksdb,
333                    &cf,
334                    ROCKSDB_PROPERTY_TOTAL_BLOB_FILES_SIZE,
335                )
336                .unwrap_or(METRICS_ERROR),
337            );
338        // 7 is the default number of levels in RocksDB. If we ever change the number of
339        // levels using `set_num_levels`, we need to update here as well. Note
340        // that there isn't an API to query the DB to get the number of levels (yet).
341        let total_num_files: i64 = (0..=6)
342            .map(|level| {
343                Self::get_rocksdb_int_property(rocksdb, &cf, &num_files_at_level(level))
344                    .unwrap_or(METRICS_ERROR)
345            })
346            .sum();
347        db_metrics
348            .cf_metrics
349            .rocksdb_total_num_files
350            .with_label_values(&[cf_name])
351            .set(total_num_files);
352        db_metrics
353            .cf_metrics
354            .rocksdb_num_level0_files
355            .with_label_values(&[cf_name])
356            .set(
357                Self::get_rocksdb_int_property(rocksdb, &cf, &num_files_at_level(0))
358                    .unwrap_or(METRICS_ERROR),
359            );
360        db_metrics
361            .cf_metrics
362            .rocksdb_current_size_active_mem_tables
363            .with_label_values(&[cf_name])
364            .set(
365                Self::get_rocksdb_int_property(rocksdb, &cf, properties::CUR_SIZE_ACTIVE_MEM_TABLE)
366                    .unwrap_or(METRICS_ERROR),
367            );
368        db_metrics
369            .cf_metrics
370            .rocksdb_size_all_mem_tables
371            .with_label_values(&[cf_name])
372            .set(
373                Self::get_rocksdb_int_property(rocksdb, &cf, properties::SIZE_ALL_MEM_TABLES)
374                    .unwrap_or(METRICS_ERROR),
375            );
376        db_metrics
377            .cf_metrics
378            .rocksdb_num_snapshots
379            .with_label_values(&[cf_name])
380            .set(
381                Self::get_rocksdb_int_property(rocksdb, &cf, properties::NUM_SNAPSHOTS)
382                    .unwrap_or(METRICS_ERROR),
383            );
384        db_metrics
385            .cf_metrics
386            .rocksdb_oldest_snapshot_time
387            .with_label_values(&[cf_name])
388            .set(
389                Self::get_rocksdb_int_property(rocksdb, &cf, properties::OLDEST_SNAPSHOT_TIME)
390                    .unwrap_or(METRICS_ERROR),
391            );
392        db_metrics
393            .cf_metrics
394            .rocksdb_actual_delayed_write_rate
395            .with_label_values(&[cf_name])
396            .set(
397                Self::get_rocksdb_int_property(rocksdb, &cf, properties::ACTUAL_DELAYED_WRITE_RATE)
398                    .unwrap_or(METRICS_ERROR),
399            );
400        db_metrics
401            .cf_metrics
402            .rocksdb_is_write_stopped
403            .with_label_values(&[cf_name])
404            .set(
405                Self::get_rocksdb_int_property(rocksdb, &cf, properties::IS_WRITE_STOPPED)
406                    .unwrap_or(METRICS_ERROR),
407            );
408        db_metrics
409            .cf_metrics
410            .rocksdb_block_cache_capacity
411            .with_label_values(&[cf_name])
412            .set(
413                Self::get_rocksdb_int_property(rocksdb, &cf, properties::BLOCK_CACHE_CAPACITY)
414                    .unwrap_or(METRICS_ERROR),
415            );
416        db_metrics
417            .cf_metrics
418            .rocksdb_block_cache_usage
419            .with_label_values(&[cf_name])
420            .set(
421                Self::get_rocksdb_int_property(rocksdb, &cf, properties::BLOCK_CACHE_USAGE)
422                    .unwrap_or(METRICS_ERROR),
423            );
424        db_metrics
425            .cf_metrics
426            .rocksdb_block_cache_pinned_usage
427            .with_label_values(&[cf_name])
428            .set(
429                Self::get_rocksdb_int_property(rocksdb, &cf, properties::BLOCK_CACHE_PINNED_USAGE)
430                    .unwrap_or(METRICS_ERROR),
431            );
432        db_metrics
433            .cf_metrics
434            .rocksdb_estimate_table_readers_mem
435            .with_label_values(&[cf_name])
436            .set(
437                Self::get_rocksdb_int_property(
438                    rocksdb,
439                    &cf,
440                    properties::ESTIMATE_TABLE_READERS_MEM,
441                )
442                .unwrap_or(METRICS_ERROR),
443            );
444        db_metrics
445            .cf_metrics
446            .rocksdb_estimated_num_keys
447            .with_label_values(&[cf_name])
448            .set(
449                Self::get_rocksdb_int_property(rocksdb, &cf, properties::ESTIMATE_NUM_KEYS)
450                    .unwrap_or(METRICS_ERROR),
451            );
452        db_metrics
453            .cf_metrics
454            .rocksdb_num_immutable_mem_tables
455            .with_label_values(&[cf_name])
456            .set(
457                Self::get_rocksdb_int_property(rocksdb, &cf, properties::NUM_IMMUTABLE_MEM_TABLE)
458                    .unwrap_or(METRICS_ERROR),
459            );
460        db_metrics
461            .cf_metrics
462            .rocksdb_mem_table_flush_pending
463            .with_label_values(&[cf_name])
464            .set(
465                Self::get_rocksdb_int_property(rocksdb, &cf, properties::MEM_TABLE_FLUSH_PENDING)
466                    .unwrap_or(METRICS_ERROR),
467            );
468        db_metrics
469            .cf_metrics
470            .rocksdb_compaction_pending
471            .with_label_values(&[cf_name])
472            .set(
473                Self::get_rocksdb_int_property(rocksdb, &cf, properties::COMPACTION_PENDING)
474                    .unwrap_or(METRICS_ERROR),
475            );
476        db_metrics
477            .cf_metrics
478            .rocksdb_estimate_pending_compaction_bytes
479            .with_label_values(&[cf_name])
480            .set(
481                Self::get_rocksdb_int_property(
482                    rocksdb,
483                    &cf,
484                    properties::ESTIMATE_PENDING_COMPACTION_BYTES,
485                )
486                .unwrap_or(METRICS_ERROR),
487            );
488        db_metrics
489            .cf_metrics
490            .rocksdb_num_running_compactions
491            .with_label_values(&[cf_name])
492            .set(
493                Self::get_rocksdb_int_property(rocksdb, &cf, properties::NUM_RUNNING_COMPACTIONS)
494                    .unwrap_or(METRICS_ERROR),
495            );
496        db_metrics
497            .cf_metrics
498            .rocksdb_num_running_flushes
499            .with_label_values(&[cf_name])
500            .set(
501                Self::get_rocksdb_int_property(rocksdb, &cf, properties::NUM_RUNNING_FLUSHES)
502                    .unwrap_or(METRICS_ERROR),
503            );
504        db_metrics
505            .cf_metrics
506            .rocksdb_estimate_oldest_key_time
507            .with_label_values(&[cf_name])
508            .set(
509                Self::get_rocksdb_int_property(rocksdb, &cf, properties::ESTIMATE_OLDEST_KEY_TIME)
510                    .unwrap_or(METRICS_ERROR),
511            );
512        db_metrics
513            .cf_metrics
514            .rocksdb_background_errors
515            .with_label_values(&[cf_name])
516            .set(
517                Self::get_rocksdb_int_property(rocksdb, &cf, properties::BACKGROUND_ERRORS)
518                    .unwrap_or(METRICS_ERROR),
519            );
520        db_metrics
521            .cf_metrics
522            .rocksdb_base_level
523            .with_label_values(&[cf_name])
524            .set(
525                Self::get_rocksdb_int_property(rocksdb, &cf, properties::BASE_LEVEL)
526                    .unwrap_or(METRICS_ERROR),
527            );
528    }
529}
typed_store/rocks/mod.rs

typed_store/rocks/
mod.rs
