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

typed_store/rocks/
mod.rs
