fedimint_rocksdb/
lib.rs

1#![deny(clippy::pedantic)]
2#![allow(clippy::missing_errors_doc)]
3#![allow(clippy::must_use_candidate)]
4#![allow(clippy::needless_lifetimes)]
5
6pub mod envs;
7
8use std::fmt;
9use std::ops::Range;
10use std::path::Path;
11use std::str::FromStr;
12
13use async_trait::async_trait;
14use fedimint_core::db::{
15    DatabaseError, DatabaseResult, IDatabaseTransactionOps, IDatabaseTransactionOpsCore,
16    IRawDatabase, IRawDatabaseTransaction, PrefixStream,
17};
18use fedimint_core::task::block_in_place;
19use fedimint_db_locked::{DbLockError, Locked, LockedBuilder};
20use futures::stream;
21pub use rocksdb;
22use rocksdb::{
23    DBRecoveryMode, OptimisticTransactionDB, OptimisticTransactionOptions, WriteOptions,
24};
25use tracing::debug;
26
27use crate::envs::{FM_ROCKSDB_BLOCK_CACHE_SIZE_ENV, FM_ROCKSDB_WRITE_BUFFER_SIZE_ENV};
28
29// turn an `iter` into a `Stream` where every `next` is ran inside
30// `block_in_place` to offload the blocking calls
31fn convert_to_async_stream<'i, I>(iter: I) -> impl futures::Stream<Item = I::Item> + use<I>
32where
33    I: Iterator + Send + 'i,
34    I::Item: Send,
35{
36    stream::unfold(iter, |mut iter| async {
37        fedimint_core::runtime::block_in_place(|| {
38            let item = iter.next();
39            item.map(|item| (item, iter))
40        })
41    })
42}
43
44#[derive(Debug)]
45pub struct RocksDb(rocksdb::OptimisticTransactionDB);
46
47pub struct RocksDbTransaction<'a>(rocksdb::Transaction<'a, rocksdb::OptimisticTransactionDB>);
48
49#[bon::bon]
50impl RocksDb {
51    /// Open the database using blocking IO
52    #[builder(start_fn = build)]
53    #[builder(finish_fn = open_blocking)]
54    pub fn open_blocking(
55        #[builder(start_fn)] db_path: impl AsRef<Path>,
56    ) -> Result<Locked<RocksDb>, RocksDbOpenError> {
57        let db_path = db_path.as_ref();
58
59        block_in_place(|| {
60            std::fs::create_dir_all(db_path.parent().ok_or(RocksDbOpenError::NoBaseDir)?)?;
61            LockedBuilder::new(db_path)?.with_db(|| Self::open_blocking_unlocked(db_path))
62        })
63    }
64}
65
66impl<I1, S> RocksDbOpenBlockingBuilder<I1, S>
67where
68    S: rocks_db_open_blocking_builder::State,
69    I1: std::convert::AsRef<std::path::Path>,
70{
71    /// Open the database
72    #[allow(clippy::unused_async)]
73    pub async fn open(self) -> Result<Locked<RocksDb>, RocksDbOpenError> {
74        block_in_place(|| self.open_blocking())
75    }
76}
77
78impl RocksDb {
79    fn open_blocking_unlocked(db_path: &Path) -> Result<RocksDb, RocksDbOpenError> {
80        let mut opts = get_default_options()?;
81        // Synchronous writes (set_sync(true)) ensure completed writes are
82        // durable, but a SIGKILL mid-write can still leave a truncated WAL tail
83        // record. TolerateCorruptedTailRecords (RocksDB's own default) discards
84        // only incomplete tail records — no committed data is lost.
85        // AbsoluteConsistency was used previously but made the database
86        // permanently unrecoverable after any unclean shutdown.
87        // See: https://github.com/fedimint/fedimint/issues/8072
88        opts.set_wal_recovery_mode(DBRecoveryMode::TolerateCorruptedTailRecords);
89        let db: rocksdb::OptimisticTransactionDB =
90            rocksdb::OptimisticTransactionDB::<rocksdb::SingleThreaded>::open(&opts, db_path)?;
91        Ok(RocksDb(db))
92    }
93
94    pub fn inner(&self) -> &rocksdb::OptimisticTransactionDB {
95        &self.0
96    }
97}
98
99// TODO: Remove this and inline it in the places where it's used.
100fn is_power_of_two(num: usize) -> bool {
101    num.is_power_of_two()
102}
103
104impl fmt::Debug for RocksDbReadOnlyTransaction<'_> {
105    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
106        f.write_str("RocksDbTransaction")
107    }
108}
109
110impl fmt::Debug for RocksDbTransaction<'_> {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        f.write_str("RocksDbTransaction")
113    }
114}
115
116#[test]
117fn is_power_of_two_sanity() {
118    assert!(!is_power_of_two(0));
119    assert!(is_power_of_two(1));
120    assert!(is_power_of_two(2));
121    assert!(!is_power_of_two(3));
122    assert!(is_power_of_two(4));
123    assert!(!is_power_of_two(5));
124    assert!(is_power_of_two(2 << 10));
125    assert!(!is_power_of_two((2 << 10) + 1));
126}
127
128/// Default write buffer size: 2 MiB (`RocksDB` default is 64 MiB)
129const DEFAULT_WRITE_BUFFER_SIZE: usize = 2 * 1024 * 1024;
130
131/// Default block cache size: 64 MiB (`RocksDB` default is 8 MiB).
132/// Index/filter blocks are placed in this cache too
133/// (`set_cache_index_and_filter_blocks`), so everything is bounded. With
134/// long keys a table's index block reaches several MiB, and every lookup
135/// needs it, so the cache must be able to hold the index blocks of the tables
136/// being read or each lookup reads and decodes them again.
137const DEFAULT_BLOCK_CACHE_SIZE: usize = 64 * 1024 * 1024;
138
139/// Default max open files: 256 (`RocksDB` default is unlimited which
140/// consumes memory for each open file handle and associated metadata)
141const DEFAULT_MAX_OPEN_FILES: i32 = 256;
142
143fn parse_env_size(env_name: &'static str) -> Result<Option<usize>, RocksDbOpenError> {
144    let Ok(var) = std::env::var(env_name) else {
145        return Ok(None);
146    };
147    let size: usize = FromStr::from_str(&var).map_err(|source| RocksDbOpenError::EnvParse {
148        var: env_name,
149        source,
150    })?;
151    if !is_power_of_two(size) {
152        return Err(RocksDbOpenError::EnvNotPowerOfTwo { var: env_name });
153    }
154    Ok(Some(size))
155}
156
157fn get_default_options() -> Result<rocksdb::Options, RocksDbOpenError> {
158    let mut opts = rocksdb::Options::default();
159
160    let write_buffer_size =
161        parse_env_size(FM_ROCKSDB_WRITE_BUFFER_SIZE_ENV)?.unwrap_or(DEFAULT_WRITE_BUFFER_SIZE);
162    opts.set_write_buffer_size(write_buffer_size);
163
164    // Keep at most 2 write buffers (1 active + 1 flushing)
165    opts.set_max_write_buffer_number(2);
166
167    let block_cache_size =
168        parse_env_size(FM_ROCKSDB_BLOCK_CACHE_SIZE_ENV)?.unwrap_or(DEFAULT_BLOCK_CACHE_SIZE);
169    // A sharded cache splits the capacity evenly and an entry larger than its
170    // shard is not retained. The default sharding keeps shards at 512 KiB to
171    // 1 MiB for caches up to 64 MiB, so index blocks of several MiB would never
172    // stay cached.
173    let mut cache_opts = rocksdb::LruCacheOptions::default();
174    cache_opts.set_capacity(block_cache_size);
175    cache_opts.set_num_shard_bits(0);
176    let cache = rocksdb::Cache::new_lru_cache_opts(&cache_opts);
177    let mut block_opts = rocksdb::BlockBasedOptions::default();
178    block_opts.set_block_cache(&cache);
179    // Put index and filter blocks into the block cache so they are
180    // bounded by the same memory budget instead of growing unbounded.
181    block_opts.set_cache_index_and_filter_blocks(true);
182    opts.set_block_based_table_factory(&block_opts);
183
184    opts.set_max_open_files(DEFAULT_MAX_OPEN_FILES);
185
186    debug!(
187        write_buffer_size,
188        block_cache_size,
189        max_open_files = DEFAULT_MAX_OPEN_FILES,
190        "RocksDB memory options"
191    );
192
193    opts.create_if_missing(true);
194    Ok(opts)
195}
196
197#[derive(Debug)]
198pub struct RocksDbReadOnly(rocksdb::DB);
199
200pub struct RocksDbReadOnlyTransaction<'a>(&'a rocksdb::DB);
201
202impl RocksDbReadOnly {
203    #[allow(clippy::unused_async)]
204    pub async fn open_read_only(
205        db_path: impl AsRef<Path>,
206    ) -> Result<RocksDbReadOnly, RocksDbOpenError> {
207        let db_path = db_path.as_ref();
208        block_in_place(|| Self::open_read_only_blocking(db_path))
209    }
210
211    pub fn open_read_only_blocking(db_path: &Path) -> Result<RocksDbReadOnly, RocksDbOpenError> {
212        let opts = get_default_options()?;
213        // Note: rocksdb is OK if one process has write access, and other read-access
214        let db = rocksdb::DB::open_for_read_only(&opts, db_path, false)?;
215        Ok(RocksDbReadOnly(db))
216    }
217}
218
219/// Why a [`RocksDb`] or [`RocksDbReadOnly`] could not be opened.
220#[derive(Debug, thiserror::Error)]
221#[non_exhaustive]
222pub enum RocksDbOpenError {
223    /// The database path has no parent directory to create; only when
224    /// opening for writing.
225    #[error("db path must have a base dir")]
226    NoBaseDir,
227
228    /// The parent directory of the database could not be created; only when
229    /// opening for writing.
230    #[error(transparent)]
231    CreateDir(#[from] std::io::Error),
232
233    /// The lock file next to the database could not be opened or locked;
234    /// only when opening for writing.
235    #[error(transparent)]
236    Lock(#[from] DbLockError),
237
238    /// A size override in the environment is not a number, or is too large
239    /// for a `usize`.
240    #[error("Could not parse {var}")]
241    EnvParse {
242        /// The environment variable.
243        var: &'static str,
244        /// Why its value could not be parsed.
245        #[source]
246        source: std::num::ParseIntError,
247    },
248
249    /// A size override in the environment is not a power of two.
250    #[error("{var} is not a power of 2")]
251    EnvNotPowerOfTwo {
252        /// The environment variable.
253        var: &'static str,
254    },
255
256    /// `RocksDB` could not open the database.
257    #[error(transparent)]
258    RocksDb(#[from] rocksdb::Error),
259}
260
261impl From<rocksdb::OptimisticTransactionDB> for RocksDb {
262    fn from(db: OptimisticTransactionDB) -> Self {
263        RocksDb(db)
264    }
265}
266
267impl From<RocksDb> for rocksdb::OptimisticTransactionDB {
268    fn from(db: RocksDb) -> Self {
269        db.0
270    }
271}
272
273// When finding by prefix iterating in Reverse order, we need to start from
274// "prefix+1" instead of "prefix", using lexicographic ordering. See the tests
275// below.
276// Will return None if there is no next prefix (i.e prefix is already the last
277// possible/max one)
278fn next_prefix(prefix: &[u8]) -> Option<Vec<u8>> {
279    let mut next_prefix = prefix.to_vec();
280    let mut is_last_prefix = true;
281    for i in (0..next_prefix.len()).rev() {
282        next_prefix[i] = next_prefix[i].wrapping_add(1);
283        if next_prefix[i] > 0 {
284            is_last_prefix = false;
285            break;
286        }
287    }
288    if is_last_prefix {
289        // The given prefix is already the last/max prefix, so there is no next prefix,
290        // return None to represent that
291        None
292    } else {
293        Some(next_prefix)
294    }
295}
296
297#[async_trait]
298impl IRawDatabase for RocksDb {
299    type Transaction<'a> = RocksDbTransaction<'a>;
300    async fn begin_transaction<'a>(&'a self) -> RocksDbTransaction {
301        let mut optimistic_options = OptimisticTransactionOptions::default();
302        optimistic_options.set_snapshot(true);
303
304        let mut write_options = WriteOptions::default();
305        // Make sure we never lose data on unclean shutdown
306        write_options.set_sync(true);
307
308        RocksDbTransaction(self.0.transaction_opt(&write_options, &optimistic_options))
309    }
310
311    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
312        let checkpoint =
313            rocksdb::checkpoint::Checkpoint::new(&self.0).map_err(DatabaseError::backend)?;
314        checkpoint
315            .create_checkpoint(backup_path)
316            .map_err(DatabaseError::backend)?;
317        Ok(())
318    }
319}
320
321#[async_trait]
322impl IRawDatabase for RocksDbReadOnly {
323    type Transaction<'a> = RocksDbReadOnlyTransaction<'a>;
324    async fn begin_transaction<'a>(&'a self) -> RocksDbReadOnlyTransaction<'a> {
325        RocksDbReadOnlyTransaction(&self.0)
326    }
327
328    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
329        let checkpoint =
330            rocksdb::checkpoint::Checkpoint::new(&self.0).map_err(DatabaseError::backend)?;
331        checkpoint
332            .create_checkpoint(backup_path)
333            .map_err(DatabaseError::backend)?;
334        Ok(())
335    }
336}
337
338#[async_trait]
339impl IDatabaseTransactionOpsCore for RocksDbTransaction<'_> {
340    async fn raw_insert_bytes(
341        &mut self,
342        key: &[u8],
343        value: &[u8],
344    ) -> DatabaseResult<Option<Vec<u8>>> {
345        fedimint_core::runtime::block_in_place(|| {
346            let val = self.0.snapshot().get(key).unwrap();
347            self.0.put(key, value).map_err(DatabaseError::backend)?;
348            Ok(val)
349        })
350    }
351
352    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
353        fedimint_core::runtime::block_in_place(|| {
354            self.0.snapshot().get(key).map_err(DatabaseError::backend)
355        })
356    }
357
358    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
359        fedimint_core::runtime::block_in_place(|| {
360            let val = self.0.snapshot().get(key).unwrap();
361            self.0.delete(key).map_err(DatabaseError::backend)?;
362            Ok(val)
363        })
364    }
365
366    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
367        Ok(fedimint_core::runtime::block_in_place(|| {
368            let prefix = key_prefix.to_vec();
369            let mut options = rocksdb::ReadOptions::default();
370            options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
371            let iter = self.0.snapshot().iterator_opt(
372                rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward),
373                options,
374            );
375            let rocksdb_iter = iter.map_while(move |res| {
376                let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
377                key_bytes
378                    .starts_with(&prefix)
379                    .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
380            });
381            Box::pin(convert_to_async_stream(rocksdb_iter))
382        }))
383    }
384
385    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
386        Ok(fedimint_core::runtime::block_in_place(|| {
387            let range = Range {
388                start: range.start.to_vec(),
389                end: range.end.to_vec(),
390            };
391            let mut options = rocksdb::ReadOptions::default();
392            options.set_iterate_range(range.clone());
393            let iter = self.0.snapshot().iterator_opt(
394                rocksdb::IteratorMode::From(&range.start, rocksdb::Direction::Forward),
395                options,
396            );
397            let rocksdb_iter = iter.map_while(move |res| {
398                let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
399                (key_bytes.as_ref() < range.end.as_slice())
400                    .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
401            });
402            Box::pin(convert_to_async_stream(rocksdb_iter))
403        }))
404    }
405
406    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
407        fedimint_core::runtime::block_in_place(|| {
408            // Note: delete_range is not supported in Transactions :/
409            let mut options = rocksdb::ReadOptions::default();
410            options.set_iterate_range(rocksdb::PrefixRange(key_prefix.to_owned()));
411            let iter = self
412                .0
413                .snapshot()
414                .iterator_opt(
415                    rocksdb::IteratorMode::From(key_prefix, rocksdb::Direction::Forward),
416                    options,
417                )
418                .map_while(|res| {
419                    res.map(|(key_bytes, _)| {
420                        key_bytes
421                            .starts_with(key_prefix)
422                            .then_some(key_bytes.to_vec())
423                    })
424                    .transpose()
425                });
426
427            for item in iter {
428                let key = item.map_err(DatabaseError::backend)?;
429                self.0.delete(key).map_err(DatabaseError::backend)?;
430            }
431
432            Ok(())
433        })
434    }
435
436    async fn raw_find_by_prefix_sorted_descending(
437        &mut self,
438        key_prefix: &[u8],
439    ) -> DatabaseResult<PrefixStream<'_>> {
440        let prefix = key_prefix.to_vec();
441        let next_prefix = next_prefix(&prefix);
442        let iterator_mode = if let Some(next_prefix) = &next_prefix {
443            rocksdb::IteratorMode::From(next_prefix, rocksdb::Direction::Reverse)
444        } else {
445            rocksdb::IteratorMode::End
446        };
447        Ok(fedimint_core::runtime::block_in_place(|| {
448            let mut options = rocksdb::ReadOptions::default();
449            options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
450            let iter = self.0.snapshot().iterator_opt(iterator_mode, options);
451            let rocksdb_iter = iter.map_while(move |res| {
452                let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
453                key_bytes
454                    .starts_with(&prefix)
455                    .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
456            });
457            Box::pin(convert_to_async_stream(rocksdb_iter))
458        }))
459    }
460}
461
462impl IDatabaseTransactionOps for RocksDbTransaction<'_> {}
463
464#[async_trait]
465impl IRawDatabaseTransaction for RocksDbTransaction<'_> {
466    async fn commit_tx(self) -> DatabaseResult<()> {
467        fedimint_core::runtime::block_in_place(|| {
468            match self.0.commit() {
469                Ok(()) => Ok(()),
470                Err(err) => {
471                    // `Busy` means another transaction wrote a key this one also wrote,
472                    // after our snapshot was taken. `TryAgain` means RocksDB could not
473                    // check for conflicts at all, because our snapshot is older than the
474                    // write history it retains — a different failure with a different
475                    // fix, so it gets its own variant.
476                    //
477                    // Anything else keeps its original kind and message: collapsing
478                    // unrelated failures into a conflict hides what actually went wrong.
479                    // Note that `Database::autocommit` retries on any commit error, so
480                    // nothing here changes which errors it recovers from.
481                    //
482                    // See: https://github.com/fedimint/fedimint/issues/8077
483                    // See: https://github.com/fedimint/fedimint/issues/8872
484                    match err.kind() {
485                        rocksdb::ErrorKind::Busy => Err(DatabaseError::WriteConflict),
486                        rocksdb::ErrorKind::TryAgain => Err(DatabaseError::snapshot_too_old(err)),
487                        _ => Err(DatabaseError::backend(err)),
488                    }
489                }
490            }
491        })
492    }
493}
494
495#[async_trait]
496impl IDatabaseTransactionOpsCore for RocksDbReadOnlyTransaction<'_> {
497    async fn raw_insert_bytes(
498        &mut self,
499        _key: &[u8],
500        _value: &[u8],
501    ) -> DatabaseResult<Option<Vec<u8>>> {
502        panic!("Cannot insert into a read only transaction");
503    }
504
505    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
506        fedimint_core::runtime::block_in_place(|| {
507            self.0.snapshot().get(key).map_err(DatabaseError::backend)
508        })
509    }
510
511    async fn raw_remove_entry(&mut self, _key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
512        panic!("Cannot remove from a read only transaction");
513    }
514
515    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
516        Ok(fedimint_core::runtime::block_in_place(|| {
517            let range = Range {
518                start: range.start.to_vec(),
519                end: range.end.to_vec(),
520            };
521            let mut options = rocksdb::ReadOptions::default();
522            options.set_iterate_range(range.clone());
523            let iter = self.0.snapshot().iterator_opt(
524                rocksdb::IteratorMode::From(&range.start, rocksdb::Direction::Forward),
525                options,
526            );
527            let rocksdb_iter = iter.map_while(move |res| {
528                let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
529                (key_bytes.as_ref() < range.end.as_slice())
530                    .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
531            });
532            Box::pin(convert_to_async_stream(rocksdb_iter))
533        }))
534    }
535
536    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
537        Ok(fedimint_core::runtime::block_in_place(|| {
538            let prefix = key_prefix.to_vec();
539            let mut options = rocksdb::ReadOptions::default();
540            options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
541            let iter = self.0.snapshot().iterator_opt(
542                rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward),
543                options,
544            );
545            let rocksdb_iter = iter.map_while(move |res| {
546                let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
547                key_bytes
548                    .starts_with(&prefix)
549                    .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
550            });
551            Box::pin(convert_to_async_stream(rocksdb_iter))
552        }))
553    }
554
555    async fn raw_remove_by_prefix(&mut self, _key_prefix: &[u8]) -> DatabaseResult<()> {
556        panic!("Cannot remove from a read only transaction");
557    }
558
559    async fn raw_find_by_prefix_sorted_descending(
560        &mut self,
561        key_prefix: &[u8],
562    ) -> DatabaseResult<PrefixStream<'_>> {
563        let prefix = key_prefix.to_vec();
564        let next_prefix = next_prefix(&prefix);
565        let iterator_mode = if let Some(next_prefix) = &next_prefix {
566            rocksdb::IteratorMode::From(next_prefix, rocksdb::Direction::Reverse)
567        } else {
568            rocksdb::IteratorMode::End
569        };
570        Ok(fedimint_core::runtime::block_in_place(|| {
571            let mut options = rocksdb::ReadOptions::default();
572            options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
573            let iter = self.0.snapshot().iterator_opt(iterator_mode, options);
574            let rocksdb_iter = iter.map_while(move |res| {
575                let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
576                key_bytes
577                    .starts_with(&prefix)
578                    .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
579            });
580            Box::pin(stream::iter(rocksdb_iter))
581        }))
582    }
583}
584
585impl IDatabaseTransactionOps for RocksDbReadOnlyTransaction<'_> {}
586
587#[async_trait]
588impl IRawDatabaseTransaction for RocksDbReadOnlyTransaction<'_> {
589    async fn commit_tx(self) -> DatabaseResult<()> {
590        panic!("Cannot commit a read only transaction");
591    }
592}
593
594#[cfg(test)]
595mod fedimint_rocksdb_tests;
fedimint_rocksdb/lib.rs

fedimint_rocksdb/
lib.rs
