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

typed_store/rocks/
mod.rs
