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 anyhow::{Context as _, bail};
14use async_trait::async_trait;
15use fedimint_core::db::{
16 DatabaseError, DatabaseResult, IDatabaseTransactionOps, IDatabaseTransactionOpsCore,
17 IRawDatabase, IRawDatabaseTransaction, PrefixStream,
18};
19use fedimint_core::task::block_in_place;
20use fedimint_db_locked::{Locked, LockedBuilder};
21use futures::stream;
22pub use rocksdb;
23use rocksdb::{
24 DBRecoveryMode, OptimisticTransactionDB, OptimisticTransactionOptions, WriteOptions,
25};
26use tracing::debug;
27
28use crate::envs::{FM_ROCKSDB_BLOCK_CACHE_SIZE_ENV, FM_ROCKSDB_WRITE_BUFFER_SIZE_ENV};
29
30fn convert_to_async_stream<'i, I>(iter: I) -> impl futures::Stream<Item = I::Item> + use<I>
33where
34 I: Iterator + Send + 'i,
35 I::Item: Send,
36{
37 stream::unfold(iter, |mut iter| async {
38 fedimint_core::runtime::block_in_place(|| {
39 let item = iter.next();
40 item.map(|item| (item, iter))
41 })
42 })
43}
44
45#[derive(Debug)]
46pub struct RocksDb(rocksdb::OptimisticTransactionDB);
47
48pub struct RocksDbTransaction<'a>(rocksdb::Transaction<'a, rocksdb::OptimisticTransactionDB>);
49
50#[bon::bon]
51impl RocksDb {
52 #[builder(start_fn = build)]
54 #[builder(finish_fn = open_blocking)]
55 pub fn open_blocking(
56 #[builder(start_fn)] db_path: impl AsRef<Path>,
57 ) -> anyhow::Result<Locked<RocksDb>> {
58 let db_path = db_path.as_ref();
59
60 block_in_place(|| {
61 std::fs::create_dir_all(
62 db_path
63 .parent()
64 .ok_or_else(|| anyhow::anyhow!("db path must have a base dir"))?,
65 )?;
66 LockedBuilder::new(db_path)?.with_db(|| Self::open_blocking_unlocked(db_path))
67 })
68 }
69}
70
71impl<I1, S> RocksDbOpenBlockingBuilder<I1, S>
72where
73 S: rocks_db_open_blocking_builder::State,
74 I1: std::convert::AsRef<std::path::Path>,
75{
76 #[allow(clippy::unused_async)]
78 pub async fn open(self) -> anyhow::Result<Locked<RocksDb>> {
79 block_in_place(|| self.open_blocking())
80 }
81}
82
83impl RocksDb {
84 fn open_blocking_unlocked(db_path: &Path) -> anyhow::Result<RocksDb> {
85 let mut opts = get_default_options()?;
86 opts.set_wal_recovery_mode(DBRecoveryMode::TolerateCorruptedTailRecords);
94 let db: rocksdb::OptimisticTransactionDB =
95 rocksdb::OptimisticTransactionDB::<rocksdb::SingleThreaded>::open(&opts, db_path)?;
96 Ok(RocksDb(db))
97 }
98
99 pub fn inner(&self) -> &rocksdb::OptimisticTransactionDB {
100 &self.0
101 }
102}
103
104fn is_power_of_two(num: usize) -> bool {
106 num.is_power_of_two()
107}
108
109impl fmt::Debug for RocksDbReadOnlyTransaction<'_> {
110 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111 f.write_str("RocksDbTransaction")
112 }
113}
114
115impl fmt::Debug for RocksDbTransaction<'_> {
116 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
117 f.write_str("RocksDbTransaction")
118 }
119}
120
121#[test]
122fn is_power_of_two_sanity() {
123 assert!(!is_power_of_two(0));
124 assert!(is_power_of_two(1));
125 assert!(is_power_of_two(2));
126 assert!(!is_power_of_two(3));
127 assert!(is_power_of_two(4));
128 assert!(!is_power_of_two(5));
129 assert!(is_power_of_two(2 << 10));
130 assert!(!is_power_of_two((2 << 10) + 1));
131}
132
133const DEFAULT_WRITE_BUFFER_SIZE: usize = 2 * 1024 * 1024;
135
136const DEFAULT_BLOCK_CACHE_SIZE: usize = 2 * 1024 * 1024;
141
142const DEFAULT_MAX_OPEN_FILES: i32 = 256;
145
146fn parse_env_size(env_name: &str) -> anyhow::Result<Option<usize>> {
147 let Ok(var) = std::env::var(env_name) else {
148 return Ok(None);
149 };
150 let size: usize =
151 FromStr::from_str(&var).with_context(|| format!("Could not parse {env_name}"))?;
152 if !is_power_of_two(size) {
153 bail!("{env_name} is not a power of 2");
154 }
155 Ok(Some(size))
156}
157
158fn get_default_options() -> anyhow::Result<rocksdb::Options> {
159 let mut opts = rocksdb::Options::default();
160
161 let write_buffer_size =
162 parse_env_size(FM_ROCKSDB_WRITE_BUFFER_SIZE_ENV)?.unwrap_or(DEFAULT_WRITE_BUFFER_SIZE);
163 opts.set_write_buffer_size(write_buffer_size);
164
165 opts.set_max_write_buffer_number(2);
167
168 let block_cache_size =
169 parse_env_size(FM_ROCKSDB_BLOCK_CACHE_SIZE_ENV)?.unwrap_or(DEFAULT_BLOCK_CACHE_SIZE);
170 let cache = rocksdb::Cache::new_lru_cache(block_cache_size);
171 let mut block_opts = rocksdb::BlockBasedOptions::default();
172 block_opts.set_block_cache(&cache);
173 block_opts.set_cache_index_and_filter_blocks(true);
176 opts.set_block_based_table_factory(&block_opts);
177
178 opts.set_max_open_files(DEFAULT_MAX_OPEN_FILES);
179
180 debug!(
181 write_buffer_size,
182 block_cache_size,
183 max_open_files = DEFAULT_MAX_OPEN_FILES,
184 "RocksDB memory options"
185 );
186
187 opts.create_if_missing(true);
188 Ok(opts)
189}
190
191#[derive(Debug)]
192pub struct RocksDbReadOnly(rocksdb::DB);
193
194pub struct RocksDbReadOnlyTransaction<'a>(&'a rocksdb::DB);
195
196impl RocksDbReadOnly {
197 #[allow(clippy::unused_async)]
198 pub async fn open_read_only(db_path: impl AsRef<Path>) -> anyhow::Result<RocksDbReadOnly> {
199 let db_path = db_path.as_ref();
200 block_in_place(|| Self::open_read_only_blocking(db_path))
201 }
202
203 pub fn open_read_only_blocking(db_path: &Path) -> anyhow::Result<RocksDbReadOnly> {
204 let opts = get_default_options()?;
205 let db = rocksdb::DB::open_for_read_only(&opts, db_path, false)?;
207 Ok(RocksDbReadOnly(db))
208 }
209}
210
211impl From<rocksdb::OptimisticTransactionDB> for RocksDb {
212 fn from(db: OptimisticTransactionDB) -> Self {
213 RocksDb(db)
214 }
215}
216
217impl From<RocksDb> for rocksdb::OptimisticTransactionDB {
218 fn from(db: RocksDb) -> Self {
219 db.0
220 }
221}
222
223fn next_prefix(prefix: &[u8]) -> Option<Vec<u8>> {
229 let mut next_prefix = prefix.to_vec();
230 let mut is_last_prefix = true;
231 for i in (0..next_prefix.len()).rev() {
232 next_prefix[i] = next_prefix[i].wrapping_add(1);
233 if next_prefix[i] > 0 {
234 is_last_prefix = false;
235 break;
236 }
237 }
238 if is_last_prefix {
239 None
242 } else {
243 Some(next_prefix)
244 }
245}
246
247#[async_trait]
248impl IRawDatabase for RocksDb {
249 type Transaction<'a> = RocksDbTransaction<'a>;
250 async fn begin_transaction<'a>(&'a self) -> RocksDbTransaction {
251 let mut optimistic_options = OptimisticTransactionOptions::default();
252 optimistic_options.set_snapshot(true);
253
254 let mut write_options = WriteOptions::default();
255 write_options.set_sync(true);
257
258 RocksDbTransaction(self.0.transaction_opt(&write_options, &optimistic_options))
259 }
260
261 fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
262 let checkpoint =
263 rocksdb::checkpoint::Checkpoint::new(&self.0).map_err(DatabaseError::backend)?;
264 checkpoint
265 .create_checkpoint(backup_path)
266 .map_err(DatabaseError::backend)?;
267 Ok(())
268 }
269}
270
271#[async_trait]
272impl IRawDatabase for RocksDbReadOnly {
273 type Transaction<'a> = RocksDbReadOnlyTransaction<'a>;
274 async fn begin_transaction<'a>(&'a self) -> RocksDbReadOnlyTransaction<'a> {
275 RocksDbReadOnlyTransaction(&self.0)
276 }
277
278 fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
279 let checkpoint =
280 rocksdb::checkpoint::Checkpoint::new(&self.0).map_err(DatabaseError::backend)?;
281 checkpoint
282 .create_checkpoint(backup_path)
283 .map_err(DatabaseError::backend)?;
284 Ok(())
285 }
286}
287
288#[async_trait]
289impl IDatabaseTransactionOpsCore for RocksDbTransaction<'_> {
290 async fn raw_insert_bytes(
291 &mut self,
292 key: &[u8],
293 value: &[u8],
294 ) -> DatabaseResult<Option<Vec<u8>>> {
295 fedimint_core::runtime::block_in_place(|| {
296 let val = self.0.snapshot().get(key).unwrap();
297 self.0.put(key, value).map_err(DatabaseError::backend)?;
298 Ok(val)
299 })
300 }
301
302 async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
303 fedimint_core::runtime::block_in_place(|| {
304 self.0.snapshot().get(key).map_err(DatabaseError::backend)
305 })
306 }
307
308 async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
309 fedimint_core::runtime::block_in_place(|| {
310 let val = self.0.snapshot().get(key).unwrap();
311 self.0.delete(key).map_err(DatabaseError::backend)?;
312 Ok(val)
313 })
314 }
315
316 async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
317 Ok(fedimint_core::runtime::block_in_place(|| {
318 let prefix = key_prefix.to_vec();
319 let mut options = rocksdb::ReadOptions::default();
320 options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
321 let iter = self.0.snapshot().iterator_opt(
322 rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward),
323 options,
324 );
325 let rocksdb_iter = iter.map_while(move |res| {
326 let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
327 key_bytes
328 .starts_with(&prefix)
329 .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
330 });
331 Box::pin(convert_to_async_stream(rocksdb_iter))
332 }))
333 }
334
335 async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
336 Ok(fedimint_core::runtime::block_in_place(|| {
337 let range = Range {
338 start: range.start.to_vec(),
339 end: range.end.to_vec(),
340 };
341 let mut options = rocksdb::ReadOptions::default();
342 options.set_iterate_range(range.clone());
343 let iter = self.0.snapshot().iterator_opt(
344 rocksdb::IteratorMode::From(&range.start, rocksdb::Direction::Forward),
345 options,
346 );
347 let rocksdb_iter = iter.map_while(move |res| {
348 let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
349 (key_bytes.as_ref() < range.end.as_slice())
350 .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
351 });
352 Box::pin(convert_to_async_stream(rocksdb_iter))
353 }))
354 }
355
356 async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
357 fedimint_core::runtime::block_in_place(|| {
358 let mut options = rocksdb::ReadOptions::default();
360 options.set_iterate_range(rocksdb::PrefixRange(key_prefix.to_owned()));
361 let iter = self
362 .0
363 .snapshot()
364 .iterator_opt(
365 rocksdb::IteratorMode::From(key_prefix, rocksdb::Direction::Forward),
366 options,
367 )
368 .map_while(|res| {
369 res.map(|(key_bytes, _)| {
370 key_bytes
371 .starts_with(key_prefix)
372 .then_some(key_bytes.to_vec())
373 })
374 .transpose()
375 });
376
377 for item in iter {
378 let key = item.map_err(DatabaseError::backend)?;
379 self.0.delete(key).map_err(DatabaseError::backend)?;
380 }
381
382 Ok(())
383 })
384 }
385
386 async fn raw_find_by_prefix_sorted_descending(
387 &mut self,
388 key_prefix: &[u8],
389 ) -> DatabaseResult<PrefixStream<'_>> {
390 let prefix = key_prefix.to_vec();
391 let next_prefix = next_prefix(&prefix);
392 let iterator_mode = if let Some(next_prefix) = &next_prefix {
393 rocksdb::IteratorMode::From(next_prefix, rocksdb::Direction::Reverse)
394 } else {
395 rocksdb::IteratorMode::End
396 };
397 Ok(fedimint_core::runtime::block_in_place(|| {
398 let mut options = rocksdb::ReadOptions::default();
399 options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
400 let iter = self.0.snapshot().iterator_opt(iterator_mode, options);
401 let rocksdb_iter = iter.map_while(move |res| {
402 let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
403 key_bytes
404 .starts_with(&prefix)
405 .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
406 });
407 Box::pin(convert_to_async_stream(rocksdb_iter))
408 }))
409 }
410}
411
412impl IDatabaseTransactionOps for RocksDbTransaction<'_> {}
413
414#[async_trait]
415impl IRawDatabaseTransaction for RocksDbTransaction<'_> {
416 async fn commit_tx(self) -> DatabaseResult<()> {
417 fedimint_core::runtime::block_in_place(|| {
418 match self.0.commit() {
419 Ok(()) => Ok(()),
420 Err(err) => {
421 match err.kind() {
435 rocksdb::ErrorKind::Busy => Err(DatabaseError::WriteConflict),
436 rocksdb::ErrorKind::TryAgain => Err(DatabaseError::snapshot_too_old(err)),
437 _ => Err(DatabaseError::backend(err)),
438 }
439 }
440 }
441 })
442 }
443}
444
445#[async_trait]
446impl IDatabaseTransactionOpsCore for RocksDbReadOnlyTransaction<'_> {
447 async fn raw_insert_bytes(
448 &mut self,
449 _key: &[u8],
450 _value: &[u8],
451 ) -> DatabaseResult<Option<Vec<u8>>> {
452 panic!("Cannot insert into a read only transaction");
453 }
454
455 async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
456 fedimint_core::runtime::block_in_place(|| {
457 self.0.snapshot().get(key).map_err(DatabaseError::backend)
458 })
459 }
460
461 async fn raw_remove_entry(&mut self, _key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
462 panic!("Cannot remove from a read only transaction");
463 }
464
465 async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
466 Ok(fedimint_core::runtime::block_in_place(|| {
467 let range = Range {
468 start: range.start.to_vec(),
469 end: range.end.to_vec(),
470 };
471 let mut options = rocksdb::ReadOptions::default();
472 options.set_iterate_range(range.clone());
473 let iter = self.0.snapshot().iterator_opt(
474 rocksdb::IteratorMode::From(&range.start, rocksdb::Direction::Forward),
475 options,
476 );
477 let rocksdb_iter = iter.map_while(move |res| {
478 let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
479 (key_bytes.as_ref() < range.end.as_slice())
480 .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
481 });
482 Box::pin(convert_to_async_stream(rocksdb_iter))
483 }))
484 }
485
486 async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
487 Ok(fedimint_core::runtime::block_in_place(|| {
488 let prefix = key_prefix.to_vec();
489 let mut options = rocksdb::ReadOptions::default();
490 options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
491 let iter = self.0.snapshot().iterator_opt(
492 rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward),
493 options,
494 );
495 let rocksdb_iter = iter.map_while(move |res| {
496 let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
497 key_bytes
498 .starts_with(&prefix)
499 .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
500 });
501 Box::pin(convert_to_async_stream(rocksdb_iter))
502 }))
503 }
504
505 async fn raw_remove_by_prefix(&mut self, _key_prefix: &[u8]) -> DatabaseResult<()> {
506 panic!("Cannot remove from a read only transaction");
507 }
508
509 async fn raw_find_by_prefix_sorted_descending(
510 &mut self,
511 key_prefix: &[u8],
512 ) -> DatabaseResult<PrefixStream<'_>> {
513 let prefix = key_prefix.to_vec();
514 let next_prefix = next_prefix(&prefix);
515 let iterator_mode = if let Some(next_prefix) = &next_prefix {
516 rocksdb::IteratorMode::From(next_prefix, rocksdb::Direction::Reverse)
517 } else {
518 rocksdb::IteratorMode::End
519 };
520 Ok(fedimint_core::runtime::block_in_place(|| {
521 let mut options = rocksdb::ReadOptions::default();
522 options.set_iterate_range(rocksdb::PrefixRange(prefix.clone()));
523 let iter = self.0.snapshot().iterator_opt(iterator_mode, options);
524 let rocksdb_iter = iter.map_while(move |res| {
525 let (key_bytes, value_bytes) = res.expect("Error reading from RocksDb");
526 key_bytes
527 .starts_with(&prefix)
528 .then_some((key_bytes.to_vec(), value_bytes.to_vec()))
529 });
530 Box::pin(stream::iter(rocksdb_iter))
531 }))
532 }
533}
534
535impl IDatabaseTransactionOps for RocksDbReadOnlyTransaction<'_> {}
536
537#[async_trait]
538impl IRawDatabaseTransaction for RocksDbReadOnlyTransaction<'_> {
539 async fn commit_tx(self) -> DatabaseResult<()> {
540 panic!("Cannot commit a read only transaction");
541 }
542}
543
544#[cfg(test)]
545mod fedimint_rocksdb_tests {
546 use fedimint_core::db::{Database, IDatabaseTransactionOpsCoreTyped};
547 use fedimint_core::encoding::{Decodable, Encodable};
548 use fedimint_core::module::registry::{ModuleDecoderRegistry, ModuleRegistry};
549 use fedimint_core::{impl_db_lookup, impl_db_record};
550 use futures::StreamExt;
551
552 use super::*;
553
554 fn open_temp_db(temp_path: &str) -> Database {
555 let path = tempfile::Builder::new()
556 .prefix(temp_path)
557 .tempdir()
558 .unwrap();
559
560 Database::new(
561 RocksDb::build(path.as_ref()).open_blocking().unwrap(),
562 ModuleDecoderRegistry::default(),
563 )
564 }
565
566 #[tokio::test(flavor = "multi_thread")]
567 async fn test_dbtx_insert_elements() {
568 fedimint_core::db::verify_insert_elements(open_temp_db("fcb-rocksdb-test-insert-elements"))
569 .await;
570 }
571
572 #[tokio::test(flavor = "multi_thread")]
573 async fn test_dbtx_remove_nonexisting() {
574 fedimint_core::db::verify_remove_nonexisting(open_temp_db(
575 "fcb-rocksdb-test-remove-nonexisting",
576 ))
577 .await;
578 }
579
580 #[tokio::test(flavor = "multi_thread")]
581 async fn test_dbtx_remove_existing() {
582 fedimint_core::db::verify_remove_existing(open_temp_db("fcb-rocksdb-test-remove-existing"))
583 .await;
584 }
585
586 #[tokio::test(flavor = "multi_thread")]
587 async fn test_dbtx_read_own_writes() {
588 fedimint_core::db::verify_read_own_writes(open_temp_db("fcb-rocksdb-test-read-own-writes"))
589 .await;
590 }
591
592 #[tokio::test(flavor = "multi_thread")]
593 async fn test_dbtx_prevent_dirty_reads() {
594 fedimint_core::db::verify_prevent_dirty_reads(open_temp_db(
595 "fcb-rocksdb-test-prevent-dirty-reads",
596 ))
597 .await;
598 }
599
600 #[tokio::test(flavor = "multi_thread")]
601 async fn test_dbtx_find_by_range() {
602 fedimint_core::db::verify_find_by_range(open_temp_db("fcb-rocksdb-test-find-by-range"))
603 .await;
604 }
605
606 #[tokio::test(flavor = "multi_thread")]
607 async fn test_dbtx_find_by_prefix() {
608 fedimint_core::db::verify_find_by_prefix(open_temp_db("fcb-rocksdb-test-find-by-prefix"))
609 .await;
610 }
611
612 #[tokio::test(flavor = "multi_thread")]
613 async fn test_dbtx_commit() {
614 fedimint_core::db::verify_commit(open_temp_db("fcb-rocksdb-test-commit")).await;
615 }
616
617 #[tokio::test(flavor = "multi_thread")]
618 async fn test_dbtx_prevent_nonrepeatable_reads() {
619 fedimint_core::db::verify_prevent_nonrepeatable_reads(open_temp_db(
620 "fcb-rocksdb-test-prevent-nonrepeatable-reads",
621 ))
622 .await;
623 }
624
625 #[tokio::test(flavor = "multi_thread")]
626 async fn test_dbtx_snapshot_isolation() {
627 fedimint_core::db::verify_snapshot_isolation(open_temp_db(
628 "fcb-rocksdb-test-snapshot-isolation",
629 ))
630 .await;
631 }
632
633 #[tokio::test(flavor = "multi_thread")]
634 async fn test_dbtx_phantom_entry() {
635 fedimint_core::db::verify_phantom_entry(open_temp_db("fcb-rocksdb-test-phantom-entry"))
636 .await;
637 }
638
639 #[tokio::test(flavor = "multi_thread")]
640 async fn test_dbtx_write_conflict() {
641 fedimint_core::db::expect_write_conflict(open_temp_db("fcb-rocksdb-test-write-conflict"))
642 .await;
643 }
644
645 #[tokio::test(flavor = "multi_thread")]
648 async fn test_concurrent_transaction_conflict_with_autocommit() {
649 use std::sync::Arc;
650
651 let db = Arc::new(open_temp_db("fcb-rocksdb-test-concurrent-conflict"));
652
653 let mut handles = Vec::new();
656
657 for i in 0u64..10 {
658 let db_clone = Arc::clone(&db);
659 let handle =
660 fedimint_core::runtime::spawn("rocksdb-transient-error-test", async move {
661 for j in 0u64..10 {
662 let result = db_clone
664 .autocommit::<_, _, anyhow::Error>(
665 |dbtx, _| {
666 #[allow(clippy::cast_possible_truncation)]
667 let val = (i * 100 + j) as u8;
668 Box::pin(async move {
669 dbtx.insert_entry(&TestKey(vec![0]), &TestVal(vec![val]))
671 .await;
672 Ok(())
673 })
674 },
675 None, )
677 .await;
678
679 assert!(
681 result.is_ok(),
682 "Transaction should succeed after retries, got: {result:?}",
683 );
684 }
685 });
686 handles.push(handle);
687 }
688
689 for handle in handles {
691 handle.await.expect("Task should not panic");
692 }
693 }
694
695 #[tokio::test(flavor = "multi_thread")]
696 async fn test_dbtx_remove_by_prefix() {
697 fedimint_core::db::verify_remove_by_prefix(open_temp_db(
698 "fcb-rocksdb-test-remove-by-prefix",
699 ))
700 .await;
701 }
702
703 #[tokio::test(flavor = "multi_thread")]
704 async fn test_module_dbtx() {
705 fedimint_core::db::verify_module_prefix(open_temp_db("fcb-rocksdb-test-module-prefix"))
706 .await;
707 }
708
709 #[tokio::test(flavor = "multi_thread")]
710 async fn test_module_db() {
711 let module_instance_id = 1;
712 let path = tempfile::Builder::new()
713 .prefix("fcb-rocksdb-test-module-db-prefix")
714 .tempdir()
715 .unwrap();
716
717 let module_db = Database::new(
718 RocksDb::build(path.as_ref()).open_blocking().unwrap(),
719 ModuleDecoderRegistry::default(),
720 );
721
722 fedimint_core::db::verify_module_db(
723 open_temp_db("fcb-rocksdb-test-module-db"),
724 module_db.with_prefix_module_id(module_instance_id).0,
725 )
726 .await;
727 }
728
729 #[test]
730 fn test_next_prefix() {
731 assert_eq!(next_prefix(&[1, 2, 3]).unwrap(), vec![1, 2, 4]);
734 assert_eq!(next_prefix(&[1, 2, 254]).unwrap(), vec![1, 2, 255]);
735 assert_eq!(next_prefix(&[1, 2, 255]).unwrap(), vec![1, 3, 0]);
736 assert_eq!(next_prefix(&[1, 255, 255]).unwrap(), vec![2, 0, 0]);
737 assert!(next_prefix(&[255, 255, 255]).is_none());
739 assert_eq!(next_prefix(&[0]).unwrap(), vec![1]);
741 assert_eq!(next_prefix(&[254]).unwrap(), vec![255]);
742 assert!(next_prefix(&[255]).is_none()); }
744
745 #[repr(u8)]
746 #[derive(Clone)]
747 pub enum TestDbKeyPrefix {
748 Test = 254,
749 MaxTest = 255,
750 }
751
752 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Encodable, Decodable)]
753 pub(super) struct TestKey(pub Vec<u8>);
754
755 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Encodable, Decodable)]
756 pub(super) struct TestVal(pub Vec<u8>);
757
758 #[derive(Debug, Encodable, Decodable)]
759 struct DbPrefixTestPrefix;
760
761 impl_db_record!(
762 key = TestKey,
763 value = TestVal,
764 db_prefix = TestDbKeyPrefix::Test,
765 notify_on_modify = true,
766 );
767 impl_db_lookup!(key = TestKey, query_prefix = DbPrefixTestPrefix);
768
769 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Encodable, Decodable)]
770 pub(super) struct TestKey2(pub Vec<u8>);
771
772 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Encodable, Decodable)]
773 pub(super) struct TestVal2(pub Vec<u8>);
774
775 #[derive(Debug, Encodable, Decodable)]
776 struct DbPrefixTestPrefixMax;
777
778 impl_db_record!(
779 key = TestKey2,
780 value = TestVal2,
781 db_prefix = TestDbKeyPrefix::MaxTest, notify_on_modify = true,
783 );
784 impl_db_lookup!(key = TestKey2, query_prefix = DbPrefixTestPrefixMax);
785
786 #[tokio::test(flavor = "multi_thread")]
787 async fn test_retrieve_descending_order() {
788 let path = tempfile::Builder::new()
789 .prefix("fcb-rocksdb-test-descending-order")
790 .tempdir()
791 .unwrap();
792 {
793 let db = Database::new(
794 RocksDb::build(&path).open().await.unwrap(),
795 ModuleDecoderRegistry::default(),
796 );
797 let mut dbtx = db.begin_transaction().await;
798 dbtx.insert_entry(&TestKey(vec![0]), &TestVal(vec![3]))
799 .await;
800 dbtx.insert_entry(&TestKey(vec![254]), &TestVal(vec![1]))
801 .await;
802 dbtx.insert_entry(&TestKey(vec![255]), &TestVal(vec![2]))
803 .await;
804 dbtx.insert_entry(&TestKey2(vec![0]), &TestVal2(vec![3]))
805 .await;
806 dbtx.insert_entry(&TestKey2(vec![254]), &TestVal2(vec![1]))
807 .await;
808 dbtx.insert_entry(&TestKey2(vec![255]), &TestVal2(vec![2]))
809 .await;
810 let query = dbtx
811 .find_by_prefix_sorted_descending(&DbPrefixTestPrefix)
812 .await
813 .collect::<Vec<_>>()
814 .await;
815 assert_eq!(
816 query,
817 vec![
818 (TestKey(vec![255]), TestVal(vec![2])),
819 (TestKey(vec![254]), TestVal(vec![1])),
820 (TestKey(vec![0]), TestVal(vec![3]))
821 ]
822 );
823 let query = dbtx
824 .find_by_prefix_sorted_descending(&DbPrefixTestPrefixMax)
825 .await
826 .collect::<Vec<_>>()
827 .await;
828 assert_eq!(
829 query,
830 vec![
831 (TestKey2(vec![255]), TestVal2(vec![2])),
832 (TestKey2(vec![254]), TestVal2(vec![1])),
833 (TestKey2(vec![0]), TestVal2(vec![3]))
834 ]
835 );
836 dbtx.commit_tx().await;
837 }
838 let db_readonly = RocksDbReadOnly::open_read_only(path).await.unwrap();
840 let db_readonly = Database::new(db_readonly, ModuleRegistry::default());
841 let mut dbtx = db_readonly.begin_transaction_nc().await;
842 let query = dbtx
843 .find_by_prefix_sorted_descending(&DbPrefixTestPrefix)
844 .await
845 .collect::<Vec<_>>()
846 .await;
847 assert_eq!(
848 query,
849 vec![
850 (TestKey(vec![255]), TestVal(vec![2])),
851 (TestKey(vec![254]), TestVal(vec![1])),
852 (TestKey(vec![0]), TestVal(vec![3]))
853 ]
854 );
855 let query = dbtx
856 .find_by_prefix_sorted_descending(&DbPrefixTestPrefixMax)
857 .await
858 .collect::<Vec<_>>()
859 .await;
860 assert_eq!(
861 query,
862 vec![
863 (TestKey2(vec![255]), TestVal2(vec![2])),
864 (TestKey2(vec![254]), TestVal2(vec![1])),
865 (TestKey2(vec![0]), TestVal2(vec![3]))
866 ]
867 );
868 }
869
870 #[tokio::test(flavor = "multi_thread")]
885 async fn test_long_lived_transaction_fails_without_key_overlap() {
886 let path = tempfile::Builder::new()
887 .prefix("fcb-rocksdb-test-long-lived-transaction")
888 .tempdir()
889 .unwrap();
890
891 let db = Database::new(
892 RocksDb::build(path.as_ref()).open_blocking().unwrap(),
893 ModuleDecoderRegistry::default(),
894 );
895
896 let mut long_lived_dbtx = db.begin_transaction().await;
899 long_lived_dbtx
900 .insert_entry(&TestKey(vec![0]), &TestVal(vec![0]))
901 .await;
902
903 let value = vec![0xab; 64 * 1024];
907
908 for index in 0u32..128 {
909 let mut dbtx = db.begin_transaction().await;
910 dbtx.insert_entry(
911 &TestKey2(index.to_be_bytes().to_vec()),
912 &TestVal2(value.clone()),
913 )
914 .await;
915 dbtx.commit_tx().await;
916 }
917
918 let result = long_lived_dbtx.commit_tx_result().await;
919
920 assert!(
921 matches!(result, Err(DatabaseError::SnapshotTooOld(_))),
922 "expected a stale snapshot, got {result:?}"
923 );
924 }
925
926 #[tokio::test(flavor = "multi_thread")]
929 async fn test_same_key_write_is_a_conflict() {
930 let path = tempfile::Builder::new()
931 .prefix("fcb-rocksdb-test-same-key-write")
932 .tempdir()
933 .unwrap();
934
935 let db = Database::new(
936 RocksDb::build(path.as_ref()).open_blocking().unwrap(),
937 ModuleDecoderRegistry::default(),
938 );
939
940 let mut first_dbtx = db.begin_transaction().await;
941 first_dbtx
942 .insert_entry(&TestKey(vec![0]), &TestVal(vec![1]))
943 .await;
944
945 let mut second_dbtx = db.begin_transaction().await;
946 second_dbtx
947 .insert_entry(&TestKey(vec![0]), &TestVal(vec![2]))
948 .await;
949 second_dbtx.commit_tx().await;
950
951 let result = first_dbtx.commit_tx_result().await;
952
953 assert!(
954 matches!(result, Err(DatabaseError::WriteConflict)),
955 "expected a write conflict, got {result:?}"
956 );
957 }
958}