linera_views/backends/
rocks_db.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Implements [`crate::store::KeyValueStore`] for the RocksDB database.
5
6// RocksDB's C API uses `i32` and signed sizes; casts at this boundary are
7// by design.
8#![allow(
9    clippy::cast_possible_truncation,
10    clippy::cast_sign_loss,
11    clippy::cast_possible_wrap
12)]
13
14use std::{
15    ffi::OsString,
16    fmt::Display,
17    path::PathBuf,
18    sync::{
19        atomic::{AtomicBool, Ordering},
20        Arc,
21    },
22};
23
24use linera_base::ensure;
25use rocksdb::{BlockBasedOptions, Cache, DBCompactionStyle, SliceTransform, WriteBufferManager};
26use serde::{Deserialize, Serialize};
27use sysinfo::{MemoryRefreshKind, RefreshKind, System};
28use tempfile::TempDir;
29use thiserror::Error;
30
31#[cfg(with_metrics)]
32use crate::metering::MeteredDatabase;
33#[cfg(with_testing)]
34use crate::store::TestKeyValueDatabase;
35use crate::{
36    batch::{Batch, WriteOperation},
37    common::get_upper_bound_option,
38    lru_caching::{LruCachingConfig, LruCachingDatabase},
39    store::{
40        KeyValueDatabase, KeyValueStoreError, ReadableKeyValueStore, WithError,
41        WritableKeyValueStore,
42    },
43    value_splitting::{ValueSplittingDatabase, ValueSplittingError},
44};
45
46/// The prefixes being used in the system
47static ROOT_KEY_DOMAIN: [u8; 1] = [0];
48static STORED_ROOT_KEYS_PREFIX: u8 = 1;
49
50// The maximum size of values in RocksDB is 3 GiB
51// For offset reasons we decrease by 400
52const MAX_VALUE_SIZE: usize = 3 * 1024 * 1024 * 1024 - 400;
53
54// The maximum size of keys in RocksDB is 8 MiB
55// For offset reasons we decrease by 400
56const MAX_KEY_SIZE: usize = 8 * 1024 * 1024 - 400;
57
58// A small write buffer keeps the memtable flushing even on low-write workloads. A large buffer
59// (e.g. 256 MiB) lets a slowly-filling memtable grow huge and accumulate range tombstones, so every
60// point read has to scan it — which severely amplifies reads during e.g. a long chain
61// synchronization, where blocks are re-executed with little net data written.
62const WRITE_BUFFER_SIZE: usize = 16 * 1024 * 1024; // 16 MiB
63const MAX_WRITE_BUFFER_NUMBER: i32 = 6;
64
65fn get_available_memory(sys: &System) -> usize {
66    sys.cgroup_limits()
67        .map_or_else(|| sys.total_memory() as usize, |c| c.total_memory as usize)
68}
69
70fn get_available_cpus() -> i32 {
71    std::thread::available_parallelism().map_or(1, |p| p.get() as i32)
72}
73
74const HYPER_CLOCK_CACHE_BLOCK_SIZE: usize = 8 * 1024; // 8 KiB
75
76/// The RocksDB client that we use.
77type DB = rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>;
78
79/// The choice of the spawning mode.
80/// `SpawnBlocking` always works and is the safest.
81/// `BlockInPlace` can only be used in multi-threaded environment.
82/// One way to select that is to select BlockInPlace when
83/// `tokio::runtime::Handle::current().metrics().num_workers() > 1`
84/// `BlockInPlace` is documented in <https://docs.rs/tokio/latest/tokio/task/fn.block_in_place.html>
85#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
86pub enum RocksDbSpawnMode {
87    /// This uses the `spawn_blocking` function of Tokio.
88    SpawnBlocking,
89    /// This uses the `block_in_place` function of Tokio.
90    BlockInPlace,
91}
92
93impl RocksDbSpawnMode {
94    /// Obtains the spawning mode from runtime.
95    pub fn get_spawn_mode_from_runtime() -> Self {
96        if tokio::runtime::Handle::current().metrics().num_workers() > 1 {
97            RocksDbSpawnMode::BlockInPlace
98        } else {
99            RocksDbSpawnMode::SpawnBlocking
100        }
101    }
102
103    /// Runs the computation for a function according to the selected policy.
104    #[inline]
105    async fn spawn<F, I, O>(&self, f: F, input: I) -> Result<O, RocksDbStoreInternalError>
106    where
107        F: FnOnce(I) -> Result<O, RocksDbStoreInternalError> + Send + 'static,
108        I: Send + 'static,
109        O: Send + 'static,
110    {
111        Ok(match self {
112            RocksDbSpawnMode::BlockInPlace => tokio::task::block_in_place(move || f(input))?,
113            RocksDbSpawnMode::SpawnBlocking => {
114                tokio::task::spawn_blocking(move || f(input)).await??
115            }
116        })
117    }
118}
119
120impl Display for RocksDbSpawnMode {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        match &self {
123            RocksDbSpawnMode::SpawnBlocking => write!(f, "spawn_blocking"),
124            RocksDbSpawnMode::BlockInPlace => write!(f, "block_in_place"),
125        }
126    }
127}
128
129fn check_key_size(key: &[u8]) -> Result<(), RocksDbStoreInternalError> {
130    ensure!(
131        key.len() <= MAX_KEY_SIZE,
132        RocksDbStoreInternalError::KeyTooLong
133    );
134    Ok(())
135}
136
137#[derive(Clone)]
138struct RocksDbStoreExecutor {
139    db: Arc<DB>,
140    start_key: Vec<u8>,
141}
142
143impl RocksDbStoreExecutor {
144    fn contains_keys_internal(
145        &self,
146        keys: Vec<Vec<u8>>,
147    ) -> Result<Vec<bool>, RocksDbStoreInternalError> {
148        let size = keys.len();
149        let mut results = vec![false; size];
150        let mut indices = Vec::new();
151        let mut keys_red = Vec::new();
152        for (i, key) in keys.into_iter().enumerate() {
153            check_key_size(&key)?;
154            let mut full_key = self.start_key.to_vec();
155            full_key.extend(key);
156            if self.db.key_may_exist(&full_key) {
157                indices.push(i);
158                keys_red.push(full_key);
159            }
160        }
161        let values_red = self.db.multi_get(keys_red);
162        for (index, value) in indices.into_iter().zip(values_red) {
163            results[index] = value?.is_some();
164        }
165        Ok(results)
166    }
167
168    fn read_multi_values_bytes_internal(
169        &self,
170        keys: Vec<Vec<u8>>,
171    ) -> Result<Vec<Option<Vec<u8>>>, RocksDbStoreInternalError> {
172        for key in &keys {
173            check_key_size(key)?;
174        }
175        let full_keys = keys
176            .into_iter()
177            .map(|key| {
178                let mut full_key = self.start_key.to_vec();
179                full_key.extend(key);
180                full_key
181            })
182            .collect::<Vec<_>>();
183        let entries = self.db.multi_get(&full_keys);
184        Ok(entries.into_iter().collect::<Result<_, _>>()?)
185    }
186
187    fn get_find_prefix_iterator(
188        &self,
189        prefix: &[u8],
190    ) -> rocksdb::DBRawIteratorWithThreadMode<'_, DB> {
191        // Configure ReadOptions optimized for SSDs and iterator performance
192        let mut read_opts = rocksdb::ReadOptions::default();
193        // Enable async I/O for better concurrency
194        read_opts.set_async_io(true);
195
196        // Set precise upper bound to minimize key traversal
197        let upper_bound = get_upper_bound_option(prefix);
198        if let Some(upper_bound) = upper_bound {
199            read_opts.set_iterate_upper_bound(upper_bound);
200        }
201
202        let mut iter = self.db.raw_iterator_opt(read_opts);
203        iter.seek(prefix);
204        iter
205    }
206
207    fn find_keys_by_prefix_internal(
208        &self,
209        key_prefix: Vec<u8>,
210    ) -> Result<Vec<Vec<u8>>, RocksDbStoreInternalError> {
211        check_key_size(&key_prefix)?;
212
213        let mut prefix = self.start_key.clone();
214        prefix.extend(key_prefix);
215        let len = prefix.len();
216
217        let mut iter = self.get_find_prefix_iterator(&prefix);
218        let mut keys = Vec::new();
219        while let Some(key) = iter.key() {
220            keys.push(key[len..].to_vec());
221            iter.next();
222        }
223        Ok(keys)
224    }
225
226    #[expect(clippy::type_complexity)]
227    fn find_key_values_by_prefix_internal(
228        &self,
229        key_prefix: Vec<u8>,
230    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, RocksDbStoreInternalError> {
231        check_key_size(&key_prefix)?;
232        let mut prefix = self.start_key.clone();
233        prefix.extend(key_prefix);
234        let len = prefix.len();
235
236        let mut iter = self.get_find_prefix_iterator(&prefix);
237        let mut key_values = Vec::new();
238        while let Some((key, value)) = iter.item() {
239            let key_value = (key[len..].to_vec(), value.to_vec());
240            key_values.push(key_value);
241            iter.next();
242        }
243        Ok(key_values)
244    }
245
246    fn write_batch_internal(
247        &self,
248        batch: Batch,
249        write_root_key: bool,
250    ) -> Result<(), RocksDbStoreInternalError> {
251        let mut inner_batch = rocksdb::WriteBatchWithTransaction::default();
252        for operation in batch.operations {
253            match operation {
254                WriteOperation::Delete { key } => {
255                    check_key_size(&key)?;
256                    let mut full_key = self.start_key.to_vec();
257                    full_key.extend(key);
258                    inner_batch.delete(&full_key)
259                }
260                WriteOperation::Put { key, value } => {
261                    check_key_size(&key)?;
262                    let mut full_key = self.start_key.to_vec();
263                    full_key.extend(key);
264                    inner_batch.put(&full_key, value)
265                }
266                WriteOperation::DeletePrefix { key_prefix } => {
267                    check_key_size(&key_prefix)?;
268                    let mut full_key1 = self.start_key.to_vec();
269                    full_key1.extend(&key_prefix);
270                    let full_key2 =
271                        get_upper_bound_option(&full_key1).expect("the first entry cannot be 255");
272                    inner_batch.delete_range(&full_key1, &full_key2);
273                }
274            }
275        }
276        if write_root_key {
277            let mut full_key = self.start_key.to_vec();
278            full_key[0] = STORED_ROOT_KEYS_PREFIX;
279            inner_batch.put(&full_key, vec![]);
280        }
281        self.db.write(inner_batch)?;
282        Ok(())
283    }
284}
285
286/// The inner client
287#[derive(Clone)]
288pub struct RocksDbStoreInternal {
289    executor: RocksDbStoreExecutor,
290    path_with_guard: PathWithGuard,
291    spawn_mode: RocksDbSpawnMode,
292    root_key_written: Arc<AtomicBool>,
293}
294
295/// Database-level connection to RocksDB for managing namespaces and partitions.
296#[derive(Clone)]
297pub struct RocksDbDatabaseInternal {
298    executor: RocksDbStoreExecutor,
299    path_with_guard: PathWithGuard,
300    spawn_mode: RocksDbSpawnMode,
301}
302
303impl WithError for RocksDbDatabaseInternal {
304    type Error = RocksDbStoreInternalError;
305}
306
307/// The level of detail collected by RocksDB's internal statistics.
308///
309/// This mirrors [`rocksdb::statistics::StatsLevel`]. The levels are nested: each one
310/// collects a superset of the data collected by the previous one, at increasing cost.
311#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize, strum::EnumString)]
312#[serde(rename_all = "kebab-case")]
313#[strum(serialize_all = "kebab-case")]
314pub enum RocksDbStatisticsLevel {
315    /// Collect nothing.
316    DisableAll,
317    /// Collect tickers (counters) only; skip all histograms and timers.
318    #[default]
319    ExceptHistogramOrTimers,
320    /// Collect tickers and histograms, but skip timer statistics.
321    ExceptTimers,
322    /// Collect everything except time spent inside the mutex lock and on compression.
323    ExceptDetailedTimers,
324    /// Collect everything except the counters that require taking time inside the mutex lock.
325    ExceptTimeForMutex,
326    /// Collect everything, including the duration of mutex operations.
327    All,
328}
329
330impl RocksDbStatisticsLevel {
331    fn to_rocksdb(self) -> rocksdb::statistics::StatsLevel {
332        use rocksdb::statistics::StatsLevel;
333        match self {
334            Self::DisableAll => StatsLevel::DisableAll,
335            Self::ExceptHistogramOrTimers => StatsLevel::ExceptHistogramOrTimers,
336            Self::ExceptTimers => StatsLevel::ExceptTimers,
337            Self::ExceptDetailedTimers => StatsLevel::ExceptDetailedTimers,
338            Self::ExceptTimeForMutex => StatsLevel::ExceptTimeForMutex,
339            Self::All => StatsLevel::All,
340        }
341    }
342}
343
344#[cfg(test)]
345mod statistics_level_tests {
346    use std::str::FromStr as _;
347
348    use super::RocksDbStatisticsLevel;
349
350    #[test]
351    fn parses_kebab_case_names() {
352        let cases = [
353            ("disable-all", RocksDbStatisticsLevel::DisableAll),
354            (
355                "except-histogram-or-timers",
356                RocksDbStatisticsLevel::ExceptHistogramOrTimers,
357            ),
358            ("except-timers", RocksDbStatisticsLevel::ExceptTimers),
359            (
360                "except-detailed-timers",
361                RocksDbStatisticsLevel::ExceptDetailedTimers,
362            ),
363            (
364                "except-time-for-mutex",
365                RocksDbStatisticsLevel::ExceptTimeForMutex,
366            ),
367            ("all", RocksDbStatisticsLevel::All),
368        ];
369        for (name, expected) in cases {
370            assert_eq!(RocksDbStatisticsLevel::from_str(name), Ok(expected));
371        }
372        assert!(RocksDbStatisticsLevel::from_str("not-a-level").is_err());
373    }
374}
375
376/// The initial configuration of the system
377#[derive(Clone, Debug, Deserialize, Serialize)]
378pub struct RocksDbStoreInternalConfig {
379    /// The path to the storage containing the namespaces
380    pub path_with_guard: PathWithGuard,
381    /// The chosen spawn mode
382    pub spawn_mode: RocksDbSpawnMode,
383    /// Whether to enable RocksDB's internal statistics collection and export it as
384    /// Prometheus metrics. Disabled by default to avoid overhead in clients that do not
385    /// scrape metrics; enabled explicitly for the workers.
386    #[serde(default)]
387    pub enable_statistics: bool,
388    /// The level of detail collected when `enable_statistics` is set.
389    #[serde(default)]
390    pub statistics_level: RocksDbStatisticsLevel,
391}
392
393impl RocksDbDatabaseInternal {
394    fn check_namespace(namespace: &str) -> Result<(), RocksDbStoreInternalError> {
395        if !namespace
396            .chars()
397            .all(|character| character.is_ascii_alphanumeric() || character == '_')
398        {
399            return Err(RocksDbStoreInternalError::InvalidNamespace);
400        }
401        Ok(())
402    }
403
404    fn build(
405        config: &RocksDbStoreInternalConfig,
406        namespace: &str,
407    ) -> Result<RocksDbDatabaseInternal, RocksDbStoreInternalError> {
408        let start_key = ROOT_KEY_DOMAIN.to_vec();
409        // Create a store to extract its executor and configuration
410        let temp_store = RocksDbStoreInternal::build(config, namespace, start_key)?;
411        Ok(RocksDbDatabaseInternal {
412            executor: temp_store.executor,
413            path_with_guard: temp_store.path_with_guard,
414            spawn_mode: temp_store.spawn_mode,
415        })
416    }
417}
418
419impl RocksDbStoreInternal {
420    fn build(
421        config: &RocksDbStoreInternalConfig,
422        namespace: &str,
423        start_key: Vec<u8>,
424    ) -> Result<RocksDbStoreInternal, RocksDbStoreInternalError> {
425        RocksDbDatabaseInternal::check_namespace(namespace)?;
426        let mut path_buf = config.path_with_guard.path_buf.clone();
427        let mut path_with_guard = config.path_with_guard.clone();
428        path_buf.push(namespace);
429        path_with_guard.path_buf = path_buf.clone();
430        let spawn_mode = config.spawn_mode;
431        if !std::path::Path::exists(&path_buf) {
432            std::fs::create_dir_all(path_buf.clone())?;
433        }
434        let sys = System::new_with_specifics(
435            RefreshKind::nothing().with_memory(MemoryRefreshKind::nothing().with_ram()),
436        );
437        let num_cpus = get_available_cpus();
438        let total_ram = get_available_memory(&sys);
439
440        let mut options = rocksdb::Options::default();
441        options.create_if_missing(true);
442        options.create_missing_column_families(true);
443
444        // Flush in-memory buffer to disk more often
445        options.set_write_buffer_size(WRITE_BUFFER_SIZE);
446        options.set_max_write_buffer_number(MAX_WRITE_BUFFER_NUMBER);
447        options.set_compression_type(rocksdb::DBCompressionType::Lz4);
448        options.set_level_zero_slowdown_writes_trigger(8);
449        options.set_level_zero_stop_writes_trigger(12);
450        options.set_level_zero_file_num_compaction_trigger(2);
451        // We deliberately give RocksDB one background thread *per* CPU so that
452        // flush + (N-1) compactions can hammer the NVMe at full bandwidth while
453        // still leaving enough CPU time for the foreground application threads.
454        options.increase_parallelism(num_cpus);
455        options.set_max_background_jobs(num_cpus);
456        options.set_max_subcompactions(num_cpus as u32);
457        options.set_level_compaction_dynamic_level_bytes(true);
458
459        options.set_compaction_style(DBCompactionStyle::Level);
460        options.set_target_file_size_base(2 * WRITE_BUFFER_SIZE as u64);
461
462        let mut block_options = BlockBasedOptions::default();
463        block_options.set_pin_l0_filter_and_index_blocks_in_cache(true);
464        block_options.set_cache_index_and_filter_blocks(true);
465        // Allocate 1/4 of total RAM for RocksDB block cache, which is a reasonable balance:
466        // - Large enough to significantly improve read performance by caching frequently accessed blocks
467        // - Small enough to leave memory for other system components
468        // - Follows common practice for database caching in server environments
469        // - Prevents excessive memory pressure that could lead to swapping or OOM conditions
470        block_options.set_block_cache(&Cache::new_hyper_clock_cache(
471            total_ram / 4,
472            HYPER_CLOCK_CACHE_BLOCK_SIZE,
473        ));
474
475        // Cap total memtable memory to prevent unbounded growth when multiple column
476        // families are used or many memtables accumulate before flushing.
477        let write_buffer_manager =
478            WriteBufferManager::new_write_buffer_manager(total_ram / 4, true);
479        options.set_write_buffer_manager(&write_buffer_manager);
480
481        // Configure bloom filters for prefix iteration optimization
482        block_options.set_bloom_filter(10.0, false);
483        block_options.set_whole_key_filtering(false);
484
485        // 32KB blocks instead of default 4KB - reduces iterator seeks
486        block_options.set_block_size(32 * 1024);
487        // Use latest format for better compression and performance
488        block_options.set_format_version(5);
489
490        options.set_block_based_table_factory(&block_options);
491
492        // Configure prefix extraction for bloom filter optimization
493        // Use 8 bytes: ROOT_KEY_DOMAIN (1 byte) + BCS variant (1-2 bytes) + identifier start (4-5 bytes)
494        let prefix_extractor = SliceTransform::create_fixed_prefix(8);
495        options.set_prefix_extractor(prefix_extractor);
496
497        // 12.5% of memtable size for bloom filter
498        options.set_memtable_prefix_bloom_ratio(0.125);
499        // Skip bloom filter for memtable when key exists
500        options.set_optimize_filters_for_hits(true);
501        // Use memory-mapped files for faster reads
502        options.set_allow_mmap_reads(true);
503        // Don't use random access pattern since we do prefix scans
504        options.set_advise_random_on_open(false);
505
506        if config.enable_statistics {
507            options.enable_statistics();
508            options.set_statistics_level(config.statistics_level.to_rocksdb());
509        }
510
511        let db = Arc::new(DB::open(&options, path_buf)?);
512        #[cfg(with_metrics)]
513        if config.enable_statistics {
514            statistics_metrics::register(Arc::new(options), db.clone());
515        }
516        let executor = RocksDbStoreExecutor { db, start_key };
517        Ok(RocksDbStoreInternal {
518            executor,
519            path_with_guard,
520            spawn_mode,
521            root_key_written: Arc::new(AtomicBool::new(false)),
522        })
523    }
524}
525
526/// Exports RocksDB's internal statistics as Prometheus metrics.
527///
528/// The collector reads the values lazily at scrape time: cumulative tickers via
529/// `get_ticker_count` and instantaneous LSM state via `GetIntProperty`. Neither requires
530/// the (more expensive) histogram/timer statistics levels.
531#[cfg(with_metrics)]
532mod statistics_metrics {
533    use std::sync::{Arc, OnceLock};
534
535    use prometheus::{
536        core::{Collector, Desc},
537        proto::MetricFamily,
538        IntGauge,
539    };
540    use rocksdb::{statistics::Ticker, Options};
541
542    use super::DB;
543
544    enum Source {
545        Ticker(Ticker),
546        Property(&'static str),
547    }
548
549    struct Entry {
550        source: Source,
551        gauge: IntGauge,
552    }
553
554    fn definitions() -> Vec<(&'static str, &'static str, Source)> {
555        vec![
556            (
557                "linera_rocksdb_block_cache_hit",
558                "Cumulative RocksDB block cache hits since open",
559                Source::Ticker(Ticker::BlockCacheHit),
560            ),
561            (
562                "linera_rocksdb_block_cache_miss",
563                "Cumulative RocksDB block cache misses since open",
564                Source::Ticker(Ticker::BlockCacheMiss),
565            ),
566            (
567                "linera_rocksdb_compact_read_bytes",
568                "Cumulative bytes read during compaction since open",
569                Source::Ticker(Ticker::CompactReadBytes),
570            ),
571            (
572                "linera_rocksdb_compact_write_bytes",
573                "Cumulative bytes written during compaction since open",
574                Source::Ticker(Ticker::CompactWriteBytes),
575            ),
576            (
577                "linera_rocksdb_flush_write_bytes",
578                "Cumulative bytes written during flushes since open",
579                Source::Ticker(Ticker::FlushWriteBytes),
580            ),
581            (
582                "linera_rocksdb_stall_micros",
583                "Cumulative write-stall time in microseconds since open",
584                Source::Ticker(Ticker::StallMicros),
585            ),
586            (
587                "linera_rocksdb_bytes_written",
588                "Cumulative user bytes written since open",
589                Source::Ticker(Ticker::BytesWritten),
590            ),
591            (
592                "linera_rocksdb_bytes_read",
593                "Cumulative user bytes read since open",
594                Source::Ticker(Ticker::BytesRead),
595            ),
596            (
597                "linera_rocksdb_wal_bytes",
598                "Cumulative bytes written to the write-ahead log since open",
599                Source::Ticker(Ticker::WalFileBytes),
600            ),
601            (
602                "linera_rocksdb_bloom_filter_useful",
603                "Cumulative count of reads avoided by the bloom filter since open",
604                Source::Ticker(Ticker::BloomFilterUseful),
605            ),
606            (
607                "linera_rocksdb_memtable_hit",
608                "Cumulative memtable hits since open",
609                Source::Ticker(Ticker::MemtableHit),
610            ),
611            (
612                "linera_rocksdb_memtable_miss",
613                "Cumulative memtable misses since open",
614                Source::Ticker(Ticker::MemtableMiss),
615            ),
616            (
617                "linera_rocksdb_number_keys_written",
618                "Cumulative number of keys written since open",
619                Source::Ticker(Ticker::NumberKeysWritten),
620            ),
621            (
622                "linera_rocksdb_num_files_at_level0",
623                "Number of files at level 0",
624                Source::Property("rocksdb.num-files-at-level0"),
625            ),
626            (
627                "linera_rocksdb_estimate_pending_compaction_bytes",
628                "Estimated bytes pending compaction",
629                Source::Property("rocksdb.estimate-pending-compaction-bytes"),
630            ),
631            (
632                "linera_rocksdb_num_running_compactions",
633                "Number of currently running compactions",
634                Source::Property("rocksdb.num-running-compactions"),
635            ),
636            (
637                "linera_rocksdb_num_running_flushes",
638                "Number of currently running flushes",
639                Source::Property("rocksdb.num-running-flushes"),
640            ),
641            (
642                "linera_rocksdb_is_write_stopped",
643                "Whether writes are currently stopped (1) or not (0)",
644                Source::Property("rocksdb.is-write-stopped"),
645            ),
646            (
647                "linera_rocksdb_actual_delayed_write_rate",
648                "Current delayed write rate in bytes/s (0 when not delayed)",
649                Source::Property("rocksdb.actual-delayed-write-rate"),
650            ),
651            (
652                "linera_rocksdb_cur_size_all_mem_tables",
653                "Approximate size in bytes of all active and unflushed memtables",
654                Source::Property("rocksdb.cur-size-all-mem-tables"),
655            ),
656            (
657                "linera_rocksdb_num_immutable_mem_table",
658                "Number of immutable memtables not yet flushed",
659                Source::Property("rocksdb.num-immutable-mem-table"),
660            ),
661            (
662                "linera_rocksdb_live_sst_files_size",
663                "Total size in bytes of all live SST files",
664                Source::Property("rocksdb.live-sst-files-size"),
665            ),
666            (
667                "linera_rocksdb_total_sst_files_size",
668                "Total size in bytes of all SST files including obsolete ones",
669                Source::Property("rocksdb.total-sst-files-size"),
670            ),
671            (
672                "linera_rocksdb_estimate_num_keys",
673                "Estimated number of keys in the database",
674                Source::Property("rocksdb.estimate-num-keys"),
675            ),
676            (
677                "linera_rocksdb_block_cache_usage",
678                "Memory in bytes used by the block cache",
679                Source::Property("rocksdb.block-cache-usage"),
680            ),
681            (
682                "linera_rocksdb_block_cache_capacity",
683                "Capacity in bytes of the block cache",
684                Source::Property("rocksdb.block-cache-capacity"),
685            ),
686        ]
687    }
688
689    struct RocksDbStatisticsCollector {
690        options: Arc<Options>,
691        db: Arc<DB>,
692        entries: Vec<Entry>,
693    }
694
695    impl RocksDbStatisticsCollector {
696        fn new(options: Arc<Options>, db: Arc<DB>) -> Self {
697            let entries = definitions()
698                .into_iter()
699                .map(|(name, help, source)| Entry {
700                    source,
701                    gauge: IntGauge::new(name, help)
702                        .expect("RocksDB statistics metric name is valid"),
703                })
704                .collect();
705            Self {
706                options,
707                db,
708                entries,
709            }
710        }
711    }
712
713    impl Collector for RocksDbStatisticsCollector {
714        fn desc(&self) -> Vec<&Desc> {
715            self.entries
716                .iter()
717                .flat_map(|entry| entry.gauge.desc())
718                .collect()
719        }
720
721        fn collect(&self) -> Vec<MetricFamily> {
722            self.entries
723                .iter()
724                .flat_map(|entry| {
725                    let value = match &entry.source {
726                        Source::Ticker(ticker) => self.options.get_ticker_count(*ticker) as i64,
727                        Source::Property(property) => {
728                            self.db
729                                .property_int_value(*property)
730                                .ok()
731                                .flatten()
732                                .unwrap_or(0) as i64
733                        }
734                    };
735                    entry.gauge.set(value);
736                    entry.gauge.collect()
737                })
738                .collect()
739        }
740    }
741
742    pub(super) fn register(options: Arc<Options>, db: Arc<DB>) {
743        static REGISTERED: OnceLock<()> = OnceLock::new();
744        if REGISTERED.set(()).is_err() {
745            tracing::warn!(
746                "RocksDB statistics collector is already registered; skipping additional store"
747            );
748            return;
749        }
750        let collector = RocksDbStatisticsCollector::new(options, db);
751        if let Err(error) = prometheus::register(Box::new(collector)) {
752            tracing::warn!("failed to register the RocksDB statistics collector: {error}");
753        }
754    }
755
756    #[cfg(test)]
757    mod tests {
758        use std::collections::HashSet;
759
760        use super::{definitions, IntGauge};
761
762        #[test]
763        fn definitions_build_unique_valid_gauges() {
764            let definitions = definitions();
765            assert!(!definitions.is_empty());
766            let mut names = HashSet::new();
767            for (name, help, _source) in &definitions {
768                assert!(!help.is_empty(), "metric {name} has empty help text");
769                assert!(names.insert(*name), "duplicate metric name: {name}");
770                IntGauge::new(*name, *help).expect("metric definition should be valid");
771            }
772        }
773    }
774}
775
776impl WithError for RocksDbStoreInternal {
777    type Error = RocksDbStoreInternalError;
778}
779
780impl ReadableKeyValueStore for RocksDbStoreInternal {
781    const MAX_KEY_SIZE: usize = MAX_KEY_SIZE;
782
783    fn root_key(&self) -> Result<Vec<u8>, RocksDbStoreInternalError> {
784        assert!(self.executor.start_key.starts_with(&ROOT_KEY_DOMAIN));
785        let root_key = bcs::from_bytes(&self.executor.start_key[ROOT_KEY_DOMAIN.len()..])?;
786        Ok(root_key)
787    }
788
789    async fn read_value_bytes(
790        &self,
791        key: &[u8],
792    ) -> Result<Option<Vec<u8>>, RocksDbStoreInternalError> {
793        check_key_size(key)?;
794        let db = self.executor.db.clone();
795        let mut full_key = self.executor.start_key.to_vec();
796        full_key.extend(key);
797        self.spawn_mode
798            .spawn(move |x| Ok(db.get(&x)?), full_key)
799            .await
800    }
801
802    async fn contains_key(&self, key: &[u8]) -> Result<bool, RocksDbStoreInternalError> {
803        check_key_size(key)?;
804        let db = self.executor.db.clone();
805        let mut full_key = self.executor.start_key.to_vec();
806        full_key.extend(key);
807        self.spawn_mode
808            .spawn(
809                move |x| {
810                    if !db.key_may_exist(&x) {
811                        return Ok(false);
812                    }
813                    Ok(db.get(&x)?.is_some())
814                },
815                full_key,
816            )
817            .await
818    }
819
820    async fn contains_keys(
821        &self,
822        keys: &[Vec<u8>],
823    ) -> Result<Vec<bool>, RocksDbStoreInternalError> {
824        let executor = self.executor.clone();
825        self.spawn_mode
826            .spawn(move |x| executor.contains_keys_internal(x), keys.to_vec())
827            .await
828    }
829
830    async fn read_multi_values_bytes(
831        &self,
832        keys: &[Vec<u8>],
833    ) -> Result<Vec<Option<Vec<u8>>>, RocksDbStoreInternalError> {
834        let executor = self.executor.clone();
835        self.spawn_mode
836            .spawn(
837                move |x| executor.read_multi_values_bytes_internal(x),
838                keys.to_vec(),
839            )
840            .await
841    }
842
843    async fn find_keys_by_prefix(
844        &self,
845        key_prefix: &[u8],
846    ) -> Result<Vec<Vec<u8>>, RocksDbStoreInternalError> {
847        let executor = self.executor.clone();
848        let key_prefix = key_prefix.to_vec();
849        self.spawn_mode
850            .spawn(
851                move |x| executor.find_keys_by_prefix_internal(x),
852                key_prefix,
853            )
854            .await
855    }
856
857    async fn find_key_values_by_prefix(
858        &self,
859        key_prefix: &[u8],
860    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, RocksDbStoreInternalError> {
861        let executor = self.executor.clone();
862        let key_prefix = key_prefix.to_vec();
863        self.spawn_mode
864            .spawn(
865                move |x| executor.find_key_values_by_prefix_internal(x),
866                key_prefix,
867            )
868            .await
869    }
870}
871
872impl WritableKeyValueStore for RocksDbStoreInternal {
873    const MAX_VALUE_SIZE: usize = MAX_VALUE_SIZE;
874
875    async fn write_batch(&self, batch: Batch) -> Result<(), RocksDbStoreInternalError> {
876        let write_root_key = !self.root_key_written.fetch_or(true, Ordering::SeqCst);
877        let executor = self.executor.clone();
878        self.spawn_mode
879            .spawn(
880                move |x| executor.write_batch_internal(x, write_root_key),
881                batch,
882            )
883            .await
884    }
885
886    async fn clear_journal(&self) -> Result<(), RocksDbStoreInternalError> {
887        Ok(())
888    }
889}
890
891impl KeyValueDatabase for RocksDbDatabaseInternal {
892    type Config = RocksDbStoreInternalConfig;
893    type Store = RocksDbStoreInternal;
894
895    fn get_name() -> String {
896        "rocksdb internal".to_string()
897    }
898
899    async fn connect(
900        config: &Self::Config,
901        namespace: &str,
902    ) -> Result<Self, RocksDbStoreInternalError> {
903        Self::build(config, namespace)
904    }
905
906    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, RocksDbStoreInternalError> {
907        let mut start_key = ROOT_KEY_DOMAIN.to_vec();
908        start_key.extend(bcs::to_bytes(root_key)?);
909        let mut executor = self.executor.clone();
910        executor.start_key = start_key;
911        Ok(RocksDbStoreInternal {
912            executor,
913            path_with_guard: self.path_with_guard.clone(),
914            spawn_mode: self.spawn_mode,
915            root_key_written: Arc::new(AtomicBool::new(false)),
916        })
917    }
918
919    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, RocksDbStoreInternalError> {
920        self.open_shared(root_key)
921    }
922
923    async fn list_all(config: &Self::Config) -> Result<Vec<String>, RocksDbStoreInternalError> {
924        let entries = std::fs::read_dir(config.path_with_guard.path_buf.clone())?;
925        let mut namespaces = Vec::new();
926        for entry in entries {
927            let entry = entry?;
928            if !entry.file_type()?.is_dir() {
929                return Err(RocksDbStoreInternalError::NonDirectoryNamespace);
930            }
931            let namespace = match entry.file_name().into_string() {
932                Err(error) => {
933                    return Err(RocksDbStoreInternalError::IntoStringError(error));
934                }
935                Ok(namespace) => namespace,
936            };
937            namespaces.push(namespace);
938        }
939        Ok(namespaces)
940    }
941
942    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, RocksDbStoreInternalError> {
943        let mut store = self.open_shared(&[])?;
944        store.executor.start_key = vec![STORED_ROOT_KEYS_PREFIX];
945        let bcs_root_keys = store.find_keys_by_prefix(&[]).await?;
946        let mut root_keys = Vec::new();
947        for bcs_root_key in bcs_root_keys {
948            let root_key = bcs::from_bytes::<Vec<u8>>(&bcs_root_key)?;
949            root_keys.push(root_key);
950        }
951        Ok(root_keys)
952    }
953
954    async fn delete_all(config: &Self::Config) -> Result<(), RocksDbStoreInternalError> {
955        let namespaces = Self::list_all(config).await?;
956        for namespace in namespaces {
957            let mut path_buf = config.path_with_guard.path_buf.clone();
958            path_buf.push(&namespace);
959            std::fs::remove_dir_all(path_buf.as_path())?;
960        }
961        Ok(())
962    }
963
964    async fn exists(
965        config: &Self::Config,
966        namespace: &str,
967    ) -> Result<bool, RocksDbStoreInternalError> {
968        Self::check_namespace(namespace)?;
969        let mut path_buf = config.path_with_guard.path_buf.clone();
970        path_buf.push(namespace);
971        let test = std::path::Path::exists(&path_buf);
972        Ok(test)
973    }
974
975    async fn create(
976        config: &Self::Config,
977        namespace: &str,
978    ) -> Result<(), RocksDbStoreInternalError> {
979        Self::check_namespace(namespace)?;
980        let mut path_buf = config.path_with_guard.path_buf.clone();
981        path_buf.push(namespace);
982        if std::path::Path::exists(&path_buf) {
983            return Err(RocksDbStoreInternalError::StoreAlreadyExists);
984        }
985        std::fs::create_dir_all(path_buf)?;
986        Ok(())
987    }
988
989    async fn delete(
990        config: &Self::Config,
991        namespace: &str,
992    ) -> Result<(), RocksDbStoreInternalError> {
993        Self::check_namespace(namespace)?;
994        let mut path_buf = config.path_with_guard.path_buf.clone();
995        path_buf.push(namespace);
996        let path = path_buf.as_path();
997        std::fs::remove_dir_all(path)?;
998        Ok(())
999    }
1000}
1001
1002#[cfg(with_testing)]
1003impl TestKeyValueDatabase for RocksDbDatabaseInternal {
1004    async fn new_test_config() -> Result<RocksDbStoreInternalConfig, RocksDbStoreInternalError> {
1005        let path_with_guard = PathWithGuard::new_testing();
1006        let spawn_mode = RocksDbSpawnMode::get_spawn_mode_from_runtime();
1007        Ok(RocksDbStoreInternalConfig {
1008            path_with_guard,
1009            spawn_mode,
1010            enable_statistics: false,
1011            statistics_level: RocksDbStatisticsLevel::default(),
1012        })
1013    }
1014}
1015
1016/// The error type for [`RocksDbStoreInternal`]
1017#[derive(Error, Debug)]
1018pub enum RocksDbStoreInternalError {
1019    /// Store already exists
1020    #[error("Store already exists")]
1021    StoreAlreadyExists,
1022
1023    /// Tokio join error in RocksDB.
1024    #[error("tokio join error: {0}")]
1025    TokioJoinError(#[from] tokio::task::JoinError),
1026
1027    /// RocksDB error.
1028    #[error("RocksDB error: {0}")]
1029    RocksDb(#[from] rocksdb::Error),
1030
1031    /// The database contains a file which is not a directory
1032    #[error("Namespaces should be directories")]
1033    NonDirectoryNamespace,
1034
1035    /// Error converting `OsString` to `String`
1036    #[error("error in the conversion from OsString: {0:?}")]
1037    IntoStringError(OsString),
1038
1039    /// The key must have at most 8 MiB
1040    #[error("The key must have at most 8 MiB")]
1041    KeyTooLong,
1042
1043    /// Namespace contains forbidden characters
1044    #[error("Namespace contains forbidden characters")]
1045    InvalidNamespace,
1046
1047    /// Filesystem error
1048    #[error("Filesystem error: {0}")]
1049    FsError(#[from] std::io::Error),
1050
1051    /// BCS serialization error.
1052    #[error(transparent)]
1053    BcsError(#[from] bcs::Error),
1054}
1055
1056/// A path and the guard for the temporary directory if needed
1057#[derive(Clone, Debug, Deserialize, Serialize)]
1058pub struct PathWithGuard {
1059    /// The path to the data
1060    pub path_buf: PathBuf,
1061    /// The guard for the directory if one is needed
1062    #[serde(skip)]
1063    _dir_guard: Option<Arc<TempDir>>,
1064}
1065
1066impl PathWithGuard {
1067    /// Creates a `PathWithGuard` from an existing path.
1068    pub fn new(path_buf: PathBuf) -> Self {
1069        Self {
1070            path_buf,
1071            _dir_guard: None,
1072        }
1073    }
1074
1075    /// Returns the test path for RocksDB without common config.
1076    #[cfg(with_testing)]
1077    fn new_testing() -> PathWithGuard {
1078        let dir = TempDir::new().unwrap();
1079        let path_buf = dir.path().to_path_buf();
1080        let dir_guard = Some(Arc::new(dir));
1081        PathWithGuard {
1082            path_buf,
1083            _dir_guard: dir_guard,
1084        }
1085    }
1086}
1087
1088impl PartialEq for PathWithGuard {
1089    fn eq(&self, other: &Self) -> bool {
1090        self.path_buf == other.path_buf
1091    }
1092}
1093impl Eq for PathWithGuard {}
1094
1095impl KeyValueStoreError for RocksDbStoreInternalError {
1096    const BACKEND: &'static str = "rocks_db";
1097}
1098
1099/// The composed error type for the `RocksDbStore`
1100pub type RocksDbStoreError = ValueSplittingError<RocksDbStoreInternalError>;
1101
1102/// The composed config type for the `RocksDbStore`
1103pub type RocksDbStoreConfig = LruCachingConfig<RocksDbStoreInternalConfig>;
1104
1105/// The `RocksDbDatabase` composed type with metrics
1106#[cfg(with_metrics)]
1107pub type RocksDbDatabase = MeteredDatabase<
1108    LruCachingDatabase<
1109        MeteredDatabase<ValueSplittingDatabase<MeteredDatabase<RocksDbDatabaseInternal>>>,
1110    >,
1111>;
1112/// The `RocksDbDatabase` composed type
1113#[cfg(not(with_metrics))]
1114pub type RocksDbDatabase = LruCachingDatabase<ValueSplittingDatabase<RocksDbDatabaseInternal>>;
1115
1116#[cfg(with_testing)]
1117impl crate::backends::DatabaseBackup for RocksDbDatabaseInternal {
1118    fn backup_to(&self, dir: &std::path::Path) -> anyhow::Result<()> {
1119        use rocksdb::{
1120            backup::{BackupEngine, BackupEngineOptions},
1121            Env,
1122        };
1123        let opts = BackupEngineOptions::new(dir)?;
1124        let env = Env::new()?;
1125        let mut engine = BackupEngine::open(&opts, &env)?;
1126        engine.create_new_backup_flush(&*self.executor.db, true)?;
1127        Ok(())
1128    }
1129}
linera_views/backends/rocks_db.rs

linera_views/backends/
rocks_db.rs
