Skip to main content

fedimint_core/db/
mod.rs

1//! Core Fedimint database traits and types
2//!
3//! This module provides the core key-value database for Fedimint.
4//!
5//! # Usage
6//!
7//! To use the database, you typically follow these steps:
8//!
9//! 1. Create a `Database` instance
10//! 2. Begin a transaction
11//! 3. Perform operations within the transaction
12//! 4. Commit the transaction
13//!
14//! ## Example
15//!
16//! ```rust
17//! use fedimint_core::db::mem_impl::MemDatabase;
18//! use fedimint_core::db::{Database, DatabaseTransaction, IDatabaseTransactionOpsCoreTyped};
19//! use fedimint_core::encoding::{Decodable, Encodable};
20//! use fedimint_core::impl_db_record;
21//! use fedimint_core::module::registry::ModuleDecoderRegistry;
22//!
23//! #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Encodable, Decodable)]
24//! pub struct TestKey(pub u64);
25//! #[derive(Debug, Encodable, Decodable, Eq, PartialEq, PartialOrd, Ord)]
26//! pub struct TestVal(pub u64);
27//!
28//! #[repr(u8)]
29//! #[derive(Clone)]
30//! pub enum TestDbKeyPrefix {
31//!     Test = 0x42,
32//! }
33//!
34//! impl_db_record!(
35//!     key = TestKey,
36//!     value = TestVal,
37//!     db_prefix = TestDbKeyPrefix::Test,
38//! );
39//!
40//! # async fn example() {
41//! // Create a new in-memory database
42//! let db = Database::new(MemDatabase::new(), ModuleDecoderRegistry::default());
43//!
44//! // Begin a transaction
45//! let mut tx = db.begin_transaction().await;
46//!
47//! // Perform operations
48//! tx.insert_entry(&TestKey(1), &TestVal(100)).await;
49//! let value = tx.get_value(&TestKey(1)).await;
50//!
51//! // Commit the transaction
52//! tx.commit_tx().await;
53//!
54//! // For operations that may need to be retried due to conflicts, use the
55//! // `autocommit` function:
56//!
57//! db.autocommit(
58//!     |dbtx, _| {
59//!         Box::pin(async move {
60//!             dbtx.insert_entry(&TestKey(1), &TestVal(100)).await;
61//!             anyhow::Ok(())
62//!         })
63//!     },
64//!     None,
65//! )
66//! .await
67//! .unwrap();
68//! # }
69//! ```
70//!
71//! # Isolation of database transactions
72//!
73//! Fedimint requires that the database implementation implement Snapshot
74//! Isolation. Snapshot Isolation is a database isolation level that guarantees
75//! consistent reads from the time that the snapshot was created (at transaction
76//! creation time). Transactions with Snapshot Isolation level will only commit
77//! if there has been no write to the modified keys since the snapshot (i.e.
78//! write-write conflicts are prevented).
79//!
80//! Specifically, Fedimint expects the database implementation to prevent the
81//! following anomalies:
82//!
83//! Non-Readable Write: TX1 writes (K1, V1) at time t but cannot read (K1, V1)
84//! at time (t + i)
85//!
86//! Dirty Read: TX1 is able to read TX2's uncommitted writes.
87//!
88//! Non-Repeatable Read: TX1 reads (K1, V1) at time t and retrieves (K1, V2) at
89//! time (t + i) where V1 != V2.
90//!
91//! Phantom Record: TX1 retrieves X number of records for a prefix at time t and
92//! retrieves Y number of records for the same prefix at time (t + i).
93//!
94//! Lost Writes: TX1 writes (K1, V1) at the same time as TX2 writes (K1, V2). V2
95//! overwrites V1 as the value for K1 (write-write conflict).
96//!
97//! | Type     | Non-Readable Write | Dirty Read | Non-Repeatable Read | Phantom
98//! Record | Lost Writes | | -------- | ------------------ | ---------- |
99//! ------------------- | -------------- | ----------- | | MemoryDB | Prevented
100//! | Prevented  | Prevented           | Prevented      | Possible    |
101//! | RocksDB  | Prevented          | Prevented  | Prevented           |
102//! Prevented      | Prevented   | | Sqlite   | Prevented          | Prevented
103//! | Prevented           | Prevented      | Prevented   |
104
105use std::any;
106use std::collections::{BTreeMap, BTreeSet};
107use std::error::Error;
108use std::fmt::{self, Debug};
109use std::marker::{self, PhantomData};
110use std::ops::{self, DerefMut, Range};
111use std::path::Path;
112use std::pin::Pin;
113use std::sync::Arc;
114use std::time::Duration;
115
116use bitcoin::hex::DisplayHex as _;
117use fedimint_core::util::BoxFuture;
118use fedimint_logging::LOG_DB;
119use fedimint_util_error::FmtCompact as _;
120use futures::{Stream, StreamExt};
121use macro_rules_attribute::apply;
122use rand::Rng;
123use serde::Serialize;
124use strum_macros::EnumIter;
125use thiserror::Error;
126use tracing::{debug, info, instrument, trace, warn};
127
128use crate::core::{ModuleInstanceId, ModuleKind};
129use crate::encoding::{Decodable, DecodeError, Encodable};
130use crate::fmt_utils::AbbreviateHexBytes;
131use crate::task::{MaybeSend, MaybeSync};
132use crate::{async_trait_maybe_send, maybe_add_send, maybe_add_send_sync, timing};
133
134pub mod mem_impl;
135pub mod notifications;
136
137pub use test_utils::*;
138
139use self::notifications::{Notifications, NotifyQueue};
140use crate::module::registry::{ModuleDecoderRegistry, ModuleRegistry};
141
142pub const MODULE_GLOBAL_PREFIX: u8 = 0xff;
143
144/// Result type for database operations
145pub type DatabaseResult<T> = std::result::Result<T, DatabaseError>;
146
147pub trait DatabaseKeyPrefix: Debug {
148    fn to_bytes(&self) -> Vec<u8>;
149}
150
151/// A key + value pair in the database with a unique prefix
152/// Extends `DatabaseKeyPrefix` to prepend the key's prefix.
153pub trait DatabaseRecord: DatabaseKeyPrefix {
154    const DB_PREFIX: u8;
155    const NOTIFY_ON_MODIFY: bool = false;
156    type Key: DatabaseKey + Debug;
157    type Value: DatabaseValue + Debug;
158}
159
160/// A key that can be used to query one or more `DatabaseRecord`
161/// Extends `DatabaseKeyPrefix` to prepend the key's prefix.
162pub trait DatabaseLookup: DatabaseKeyPrefix {
163    type Record: DatabaseRecord;
164}
165
166// Every `DatabaseRecord` is automatically a `DatabaseLookup`
167impl<Record> DatabaseLookup for Record
168where
169    Record: DatabaseRecord + Debug + Decodable + Encodable,
170{
171    type Record = Record;
172}
173
174/// `DatabaseKey` that represents the lookup structure for retrieving key/value
175/// pairs from the database.
176pub trait DatabaseKey: Sized {
177    /// Send a notification to tasks waiting to be notified if the value of
178    /// `DatabaseKey` is modified
179    ///
180    /// For instance, this can be used to be notified when a key in the
181    /// database is created. It is also possible to run a closure with the
182    /// value of the `DatabaseKey` as parameter to verify some changes to
183    /// that value.
184    const NOTIFY_ON_MODIFY: bool = false;
185    fn from_bytes(
186        data: &[u8],
187        modules: &ModuleDecoderRegistry,
188    ) -> std::result::Result<Self, DecodingError>;
189}
190
191/// Marker trait for `DatabaseKey`s where `NOTIFY` is true
192pub trait DatabaseKeyWithNotify {}
193
194/// `DatabaseValue` that represents the value structure of database records.
195pub trait DatabaseValue: Sized + Debug {
196    fn from_bytes(
197        data: &[u8],
198        modules: &ModuleDecoderRegistry,
199    ) -> std::result::Result<Self, DecodingError>;
200    fn to_bytes(&self) -> Vec<u8>;
201}
202
203pub type PrefixStream<'a> = Pin<Box<maybe_add_send!(dyn Stream<Item = (Vec<u8>, Vec<u8>)> + 'a)>>;
204
205/// Just ignore this type, it's only there to make compiler happy
206///
207/// See <https://users.rust-lang.org/t/argument-requires-that-is-borrowed-for-static/66503/2?u=yandros> for details.
208pub type PhantomBound<'big, 'small> = PhantomData<&'small &'big ()>;
209
210/// Error returned when the autocommit function fails
211#[derive(Debug, Error)]
212pub enum AutocommitError<E> {
213    /// Committing the transaction failed too many times, giving up
214    #[error("Commit Failed: {last_error}")]
215    CommitFailed {
216        /// Number of attempts
217        attempts: usize,
218        /// Last error on commit
219        last_error: DatabaseError,
220    },
221    /// Error returned by the closure provided to `autocommit`. If returned no
222    /// commit was attempted in that round
223    #[error("Closure error: {error}")]
224    ClosureError {
225        /// The attempt on which the closure returned an error
226        ///
227        /// Values other than 0 typically indicate a logic error since the
228        /// closure given to `autocommit` should not have side effects
229        /// and thus keep succeeding if it succeeded once.
230        attempts: usize,
231        /// Error returned by the closure
232        error: E,
233    },
234}
235
236pub trait AutocommitResultExt<T, E> {
237    /// Unwraps the "commit failed" error variant. Use this in cases where
238    /// autocommit is instructed to run indefinitely and commit will thus never
239    /// fail.
240    fn unwrap_autocommit(self) -> std::result::Result<T, E>;
241}
242
243impl<T, E> AutocommitResultExt<T, E> for std::result::Result<T, AutocommitError<E>> {
244    fn unwrap_autocommit(self) -> std::result::Result<T, E> {
245        match self {
246            Ok(value) => Ok(value),
247            Err(AutocommitError::CommitFailed { .. }) => {
248                panic!("`unwrap_autocommit` called on a autocommit result with finite retries");
249            }
250            Err(AutocommitError::ClosureError { error, .. }) => Err(error),
251        }
252    }
253}
254
255/// Raw database implementation
256///
257/// This and [`IRawDatabaseTransaction`] are meant to be implemented
258/// by crates like `fedimint-rocksdb` to provide a concrete implementation
259/// of a database to be used by Fedimint.
260///
261/// This is in contrast of [`IDatabase`] which includes extra
262/// functionality that Fedimint needs (and adds) on top of it.
263#[apply(async_trait_maybe_send!)]
264pub trait IRawDatabase: Debug + MaybeSend + MaybeSync + 'static {
265    /// A raw database transaction type
266    type Transaction<'a>: IRawDatabaseTransaction + Debug;
267
268    /// Start a database transaction
269    async fn begin_transaction<'a>(&'a self) -> Self::Transaction<'a>;
270
271    // Checkpoint the database to a backup directory
272    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()>;
273}
274
275#[apply(async_trait_maybe_send!)]
276impl<T> IRawDatabase for Box<T>
277where
278    T: IRawDatabase,
279{
280    type Transaction<'a> = <T as IRawDatabase>::Transaction<'a>;
281
282    async fn begin_transaction<'a>(&'a self) -> Self::Transaction<'a> {
283        (**self).begin_transaction().await
284    }
285
286    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
287        (**self).checkpoint(backup_path)
288    }
289}
290
291/// An extension trait with convenience operations on [`IRawDatabase`]
292pub trait IRawDatabaseExt: IRawDatabase + Sized {
293    /// Convert to type implementing [`IRawDatabase`] into [`Database`].
294    ///
295    /// When type inference is not an issue, [`Into::into`] can be used instead.
296    fn into_database(self) -> Database {
297        Database::new(self, ModuleRegistry::default())
298    }
299}
300
301impl<T> IRawDatabaseExt for T where T: IRawDatabase {}
302
303impl<T> From<T> for Database
304where
305    T: IRawDatabase,
306{
307    fn from(raw: T) -> Self {
308        Self::new(raw, ModuleRegistry::default())
309    }
310}
311
312/// A database that on top of a raw database operation, implements
313/// key notification system.
314#[apply(async_trait_maybe_send!)]
315pub trait IDatabase: Debug + MaybeSend + MaybeSync + 'static {
316    /// Start a database transaction
317    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a>;
318    /// Register (and wait) for `key` updates
319    async fn register(&self, key: &[u8]);
320    /// Notify about `key` update (creation, modification, deletion)
321    async fn notify(&self, key: &[u8]);
322
323    /// The prefix len of this database refers to the global (as opposed to
324    /// module-isolated) key space
325    fn is_global(&self) -> bool;
326
327    /// Checkpoints the database to a backup directory
328    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()>;
329}
330
331#[apply(async_trait_maybe_send!)]
332impl<T> IDatabase for Arc<T>
333where
334    T: IDatabase + ?Sized,
335{
336    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a> {
337        (**self).begin_transaction().await
338    }
339    async fn register(&self, key: &[u8]) {
340        (**self).register(key).await;
341    }
342    async fn notify(&self, key: &[u8]) {
343        (**self).notify(key).await;
344    }
345
346    fn is_global(&self) -> bool {
347        (**self).is_global()
348    }
349
350    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
351        (**self).checkpoint(backup_path)
352    }
353}
354
355/// Base functionality around [`IRawDatabase`] to make it a [`IDatabase`]
356///
357/// Mostly notification system, but also run-time single-commit handling.
358struct BaseDatabase<RawDatabase> {
359    notifications: Arc<Notifications>,
360    raw: RawDatabase,
361}
362
363impl<RawDatabase> fmt::Debug for BaseDatabase<RawDatabase> {
364    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
365        f.write_str("BaseDatabase")
366    }
367}
368
369#[apply(async_trait_maybe_send!)]
370impl<RawDatabase: IRawDatabase + MaybeSend + 'static> IDatabase for BaseDatabase<RawDatabase> {
371    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a> {
372        Box::new(BaseDatabaseTransaction::new(
373            self.raw.begin_transaction().await,
374            self.notifications.clone(),
375        ))
376    }
377    async fn register(&self, key: &[u8]) {
378        self.notifications.register(key).await;
379    }
380    async fn notify(&self, key: &[u8]) {
381        self.notifications.notify(key);
382    }
383
384    fn is_global(&self) -> bool {
385        true
386    }
387
388    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
389        self.raw.checkpoint(backup_path)
390    }
391}
392
393/// A public-facing newtype over `IDatabase`
394///
395/// Notably carries set of module decoders (`ModuleDecoderRegistry`)
396/// and implements common utility function for auto-commits, db isolation,
397/// and other.
398#[derive(Clone, Debug)]
399pub struct Database {
400    inner: Arc<dyn IDatabase + 'static>,
401    module_decoders: ModuleDecoderRegistry,
402}
403
404impl Database {
405    pub fn strong_count(&self) -> usize {
406        Arc::strong_count(&self.inner)
407    }
408
409    pub fn into_inner(self) -> Arc<dyn IDatabase + 'static> {
410        self.inner
411    }
412}
413
414impl Database {
415    /// Creates a new Fedimint database from any object implementing
416    /// [`IDatabase`].
417    ///
418    /// See also [`Database::new_from_arc`].
419    pub fn new(raw: impl IRawDatabase + 'static, module_decoders: ModuleDecoderRegistry) -> Self {
420        let inner = BaseDatabase {
421            raw,
422            notifications: Arc::new(Notifications::new()),
423        };
424        Self::new_from_arc(
425            Arc::new(inner) as Arc<dyn IDatabase + 'static>,
426            module_decoders,
427        )
428    }
429
430    /// Create [`Database`] from an already typed-erased `IDatabase`.
431    pub fn new_from_arc(
432        inner: Arc<dyn IDatabase + 'static>,
433        module_decoders: ModuleDecoderRegistry,
434    ) -> Self {
435        Self {
436            inner,
437            module_decoders,
438        }
439    }
440
441    /// Create [`Database`] isolated to a partition with a given `prefix`
442    pub fn with_prefix(&self, prefix: Vec<u8>) -> Self {
443        Self {
444            inner: Arc::new(PrefixDatabase {
445                inner: self.inner.clone(),
446                global_dbtx_access_token: None,
447                prefix,
448            }),
449            module_decoders: self.module_decoders.clone(),
450        }
451    }
452
453    /// Create [`Database`] isolated to a partition with a prefix for a given
454    /// `module_instance_id`, allowing the module to access `global_dbtx` with
455    /// the right `access_token`
456    pub fn with_prefix_module_id(
457        &self,
458        module_instance_id: ModuleInstanceId,
459    ) -> (Self, GlobalDBTxAccessToken) {
460        let prefix = module_instance_id_to_byte_prefix(module_instance_id);
461        let global_dbtx_access_token = GlobalDBTxAccessToken::from_prefix(&prefix);
462        (
463            Self {
464                inner: Arc::new(PrefixDatabase {
465                    inner: self.inner.clone(),
466                    global_dbtx_access_token: Some(global_dbtx_access_token),
467                    prefix,
468                }),
469                module_decoders: self.module_decoders.clone(),
470            },
471            global_dbtx_access_token,
472        )
473    }
474
475    pub fn with_decoders(&self, module_decoders: ModuleDecoderRegistry) -> Self {
476        Self {
477            inner: self.inner.clone(),
478            module_decoders,
479        }
480    }
481
482    /// Is this `Database` a global, unpartitioned `Database`
483    pub fn is_global(&self) -> bool {
484        self.inner.is_global()
485    }
486
487    /// `Err` if [`Self::is_global`] is not true
488    pub fn ensure_global(&self) -> DatabaseResult<()> {
489        if !self.is_global() {
490            return Err(DatabaseError::NotGlobal);
491        }
492
493        Ok(())
494    }
495
496    /// `Err` if [`Self::is_global`] is true
497    pub fn ensure_isolated(&self) -> DatabaseResult<()> {
498        if self.is_global() {
499            return Err(DatabaseError::NotIsolated);
500        }
501
502        Ok(())
503    }
504
505    /// Begin a new committable database transaction
506    pub async fn begin_transaction<'s, 'tx>(&'s self) -> DatabaseTransaction<'tx, Committable>
507    where
508        's: 'tx,
509    {
510        DatabaseTransaction::<Committable>::new(
511            self.inner.begin_transaction().await,
512            self.module_decoders.clone(),
513        )
514    }
515
516    /// Begin a new non-committable database transaction
517    pub async fn begin_transaction_nc<'s, 'tx>(&'s self) -> DatabaseTransaction<'tx, NonCommittable>
518    where
519        's: 'tx,
520    {
521        self.begin_transaction().await.into_nc()
522    }
523
524    pub fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
525        self.inner.checkpoint(backup_path)
526    }
527
528    /// Runs a closure with a reference to a database transaction and tries to
529    /// commit the transaction if the closure returns `Ok` and rolls it back
530    /// otherwise. If committing fails the closure is run for up to
531    /// `max_attempts` times. If `max_attempts` is `None` it will run
532    /// `usize::MAX` times which is close enough to infinite times.
533    ///
534    /// The closure `tx_fn` provided should not have side effects outside of the
535    /// database transaction provided, or if it does these should be
536    /// idempotent, since the closure might be run multiple times.
537    ///
538    /// # Lifetime Parameters
539    ///
540    /// The higher rank trait bound (HRTB) `'a` that is applied to the the
541    /// mutable reference to the database transaction ensures that the
542    /// reference lives as least as long as the returned future of the
543    /// closure.
544    ///
545    /// Further, the reference to self (`'s`) must outlive the
546    /// `DatabaseTransaction<'dt>`. In other words, the
547    /// `DatabaseTransaction` must live as least as long as `self` and that is
548    /// true as the `DatabaseTransaction` is only dropped at the end of the
549    /// `loop{}`.
550    ///
551    /// # Panics
552    ///
553    /// This function panics when the given number of maximum attempts is zero.
554    /// `max_attempts` must be greater or equal to one.
555    pub async fn autocommit<'s, 'dbtx, F, T, E>(
556        &'s self,
557        tx_fn: F,
558        max_attempts: Option<usize>,
559    ) -> std::result::Result<T, AutocommitError<E>>
560    where
561        's: 'dbtx,
562        for<'r, 'o> F: Fn(
563            &'r mut DatabaseTransaction<'o>,
564            PhantomBound<'dbtx, 'o>,
565        ) -> BoxFuture<'r, std::result::Result<T, E>>,
566    {
567        assert_ne!(max_attempts, Some(0));
568        let mut curr_attempts: usize = 0;
569
570        loop {
571            // The `checked_add()` function is used to catch the `usize` overflow.
572            // With `usize=32bit` and an assumed time of 1ms per iteration, this would crash
573            // after ~50 days. But if that's the case, something else must be wrong.
574            // With `usize=64bit` it would take much longer, obviously.
575            curr_attempts = curr_attempts
576                .checked_add(1)
577                .expect("db autocommit attempt counter overflowed");
578
579            let mut dbtx = self.begin_transaction().await;
580
581            let tx_fn_res = tx_fn(&mut dbtx.to_ref_nc(), PhantomData).await;
582            let val = match tx_fn_res {
583                Ok(val) => val,
584                Err(err) => {
585                    dbtx.ignore_uncommitted();
586                    return Err(AutocommitError::ClosureError {
587                        attempts: curr_attempts,
588                        error: err,
589                    });
590                }
591            };
592
593            let _timing /* logs on drop */ = timing::TimeReporter::new("autocommit - commit_tx");
594
595            match dbtx.commit_tx_result().await {
596                Ok(()) => {
597                    return Ok(val);
598                }
599                Err(err) => {
600                    if max_attempts.is_some_and(|max_att| max_att <= curr_attempts) {
601                        warn!(
602                            target: LOG_DB,
603                            curr_attempts,
604                            err = %err.fmt_compact(),
605                            "Database commit failed in an autocommit block - terminating"
606                        );
607                        return Err(AutocommitError::CommitFailed {
608                            attempts: curr_attempts,
609                            last_error: err,
610                        });
611                    }
612
613                    let delay = (2u64.pow(curr_attempts.min(7) as u32) * 10).min(1000);
614                    let delay = rand::thread_rng().gen_range(delay..(2 * delay));
615                    warn!(
616                        target: LOG_DB,
617                        curr_attempts,
618                        err = %err.fmt_compact(),
619                        delay_ms = %delay,
620                        "Database commit failed in an autocommit block - retrying"
621                    );
622                    crate::runtime::sleep(Duration::from_millis(delay)).await;
623                }
624            }
625        }
626    }
627
628    /// Waits for key to be notified.
629    ///
630    /// Calls the `checker` when value of the key may have changed.
631    /// Returns the value when `checker` returns a `Some(T)`.
632    pub async fn wait_key_check<'a, K, T>(
633        &'a self,
634        key: &K,
635        checker: impl Fn(Option<K::Value>) -> Option<T>,
636    ) -> (T, DatabaseTransaction<'a, Committable>)
637    where
638        K: DatabaseKey + DatabaseRecord + DatabaseKeyWithNotify,
639    {
640        let key_bytes = key.to_bytes();
641        loop {
642            // register for notification
643            let notify = self.inner.register(&key_bytes);
644
645            // check for value in db
646            let mut tx = self.inner.begin_transaction().await;
647
648            let maybe_value_bytes = tx
649                .raw_get_bytes(&key_bytes)
650                .await
651                .expect("Unrecoverable error when reading from database")
652                .map(|value_bytes| {
653                    decode_value_expect(&value_bytes, &self.module_decoders, &key_bytes)
654                });
655
656            if let Some(value) = checker(maybe_value_bytes) {
657                return (
658                    value,
659                    DatabaseTransaction::new(tx, self.module_decoders.clone()),
660                );
661            }
662
663            // key not found, try again
664            notify.await;
665            // if miss a notification between await and next register, it is
666            // fine. because we are going check the database
667        }
668    }
669
670    /// Waits for key to be present in database.
671    pub async fn wait_key_exists<K>(&self, key: &K) -> K::Value
672    where
673        K: DatabaseKey + DatabaseRecord + DatabaseKeyWithNotify,
674    {
675        self.wait_key_check(key, std::convert::identity).await.0
676    }
677}
678
679fn module_instance_id_to_byte_prefix(module_instance_id: u16) -> Vec<u8> {
680    let mut bytes = vec![MODULE_GLOBAL_PREFIX];
681    bytes.append(&mut module_instance_id.consensus_encode_to_vec());
682    bytes
683}
684
685/// A database that wraps an `inner` one and adds a prefix to all operations,
686/// effectively creating an isolated partition.
687#[derive(Clone, Debug)]
688struct PrefixDatabase<Inner>
689where
690    Inner: Debug,
691{
692    prefix: Vec<u8>,
693    global_dbtx_access_token: Option<GlobalDBTxAccessToken>,
694    inner: Inner,
695}
696
697impl<Inner> PrefixDatabase<Inner>
698where
699    Inner: Debug,
700{
701    // TODO: we should optimize these concatenations, maybe by having an internal
702    // `key: &[&[u8]]` that we flatten once, when passing to lowest layer, or
703    // something
704    fn get_full_key(&self, key: &[u8]) -> Vec<u8> {
705        let mut full_key = self.prefix.clone();
706        full_key.extend_from_slice(key);
707        full_key
708    }
709}
710
711#[apply(async_trait_maybe_send!)]
712impl<Inner> IDatabase for PrefixDatabase<Inner>
713where
714    Inner: Debug + MaybeSend + MaybeSync + 'static + IDatabase,
715{
716    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a> {
717        Box::new(PrefixDatabaseTransaction {
718            inner: self.inner.begin_transaction().await,
719            global_dbtx_access_token: self.global_dbtx_access_token,
720            prefix: self.prefix.clone(),
721        })
722    }
723    async fn register(&self, key: &[u8]) {
724        self.inner.register(&self.get_full_key(key)).await;
725    }
726
727    async fn notify(&self, key: &[u8]) {
728        self.inner.notify(&self.get_full_key(key)).await;
729    }
730
731    fn is_global(&self) -> bool {
732        if self.global_dbtx_access_token.is_some() {
733            false
734        } else {
735            self.inner.is_global()
736        }
737    }
738
739    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
740        self.inner.checkpoint(backup_path)
741    }
742}
743
744/// A database transactions that wraps an `inner` one and adds a prefix to all
745/// operations, effectively creating an isolated partition.
746///
747/// Produced by [`PrefixDatabase`].
748#[derive(Debug)]
749struct PrefixDatabaseTransaction<Inner> {
750    inner: Inner,
751    global_dbtx_access_token: Option<GlobalDBTxAccessToken>,
752    prefix: Vec<u8>,
753}
754
755impl<Inner> PrefixDatabaseTransaction<Inner> {
756    // TODO: we should optimize these concatenations, maybe by having an internal
757    // `key: &[&[u8]]` that we flatten once, when passing to lowest layer, or
758    // something
759    fn get_full_key(&self, key: &[u8]) -> Vec<u8> {
760        let mut full_key = self.prefix.clone();
761        full_key.extend_from_slice(key);
762        full_key
763    }
764
765    fn get_full_range(&self, range: Range<&[u8]>) -> Range<Vec<u8>> {
766        Range {
767            start: self.get_full_key(range.start),
768            end: self.get_full_key(range.end),
769        }
770    }
771
772    fn adapt_prefix_stream(stream: PrefixStream<'_>, prefix_len: usize) -> PrefixStream<'_> {
773        Box::pin(stream.map(move |(k, v)| (k[prefix_len..].to_owned(), v)))
774    }
775}
776
777#[apply(async_trait_maybe_send!)]
778impl<Inner> IDatabaseTransaction for PrefixDatabaseTransaction<Inner>
779where
780    Inner: IDatabaseTransaction,
781{
782    async fn commit_tx(&mut self) -> DatabaseResult<()> {
783        self.inner.commit_tx().await
784    }
785
786    fn is_global(&self) -> bool {
787        if self.global_dbtx_access_token.is_some() {
788            false
789        } else {
790            self.inner.is_global()
791        }
792    }
793
794    fn global_dbtx(
795        &mut self,
796        access_token: GlobalDBTxAccessToken,
797    ) -> &mut dyn IDatabaseTransaction {
798        if let Some(self_global_dbtx_access_token) = self.global_dbtx_access_token {
799            assert_eq!(
800                access_token, self_global_dbtx_access_token,
801                "Invalid access key used to access global_dbtx"
802            );
803            &mut self.inner
804        } else {
805            self.inner.global_dbtx(access_token)
806        }
807    }
808}
809
810#[apply(async_trait_maybe_send!)]
811impl<Inner> IDatabaseTransactionOpsCore for PrefixDatabaseTransaction<Inner>
812where
813    Inner: IDatabaseTransactionOpsCore,
814{
815    async fn raw_insert_bytes(
816        &mut self,
817        key: &[u8],
818        value: &[u8],
819    ) -> DatabaseResult<Option<Vec<u8>>> {
820        let key = self.get_full_key(key);
821        self.inner.raw_insert_bytes(&key, value).await
822    }
823
824    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
825        let key = self.get_full_key(key);
826        self.inner.raw_get_bytes(&key).await
827    }
828
829    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
830        let key = self.get_full_key(key);
831        self.inner.raw_remove_entry(&key).await
832    }
833
834    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
835        let key = self.get_full_key(key_prefix);
836        let stream = self.inner.raw_find_by_prefix(&key).await?;
837        Ok(Self::adapt_prefix_stream(stream, self.prefix.len()))
838    }
839
840    async fn raw_find_by_prefix_sorted_descending(
841        &mut self,
842        key_prefix: &[u8],
843    ) -> DatabaseResult<PrefixStream<'_>> {
844        let key = self.get_full_key(key_prefix);
845        let stream = self
846            .inner
847            .raw_find_by_prefix_sorted_descending(&key)
848            .await?;
849        Ok(Self::adapt_prefix_stream(stream, self.prefix.len()))
850    }
851
852    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
853        let range = self.get_full_range(range);
854        let stream = self
855            .inner
856            .raw_find_by_range(Range {
857                start: &range.start,
858                end: &range.end,
859            })
860            .await?;
861        Ok(Self::adapt_prefix_stream(stream, self.prefix.len()))
862    }
863
864    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
865        let key = self.get_full_key(key_prefix);
866        self.inner.raw_remove_by_prefix(&key).await
867    }
868}
869
870impl<Inner> IDatabaseTransactionOps for PrefixDatabaseTransaction<Inner> where
871    Inner: IDatabaseTransactionOps
872{
873}
874
875/// Core raw a operations database transactions supports
876///
877/// Used to enforce the same signature on all types supporting it
878#[apply(async_trait_maybe_send!)]
879pub trait IDatabaseTransactionOpsCore: MaybeSend {
880    /// Insert entry
881    async fn raw_insert_bytes(
882        &mut self,
883        key: &[u8],
884        value: &[u8],
885    ) -> DatabaseResult<Option<Vec<u8>>>;
886
887    /// Get key value
888    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>>;
889
890    /// Remove entry by `key`
891    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>>;
892
893    /// Returns an stream of key-value pairs with keys that start with
894    /// `key_prefix`, sorted by key.
895    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>>;
896
897    /// Same as [`Self::raw_find_by_prefix`] but the order is descending by key.
898    async fn raw_find_by_prefix_sorted_descending(
899        &mut self,
900        key_prefix: &[u8],
901    ) -> DatabaseResult<PrefixStream<'_>>;
902
903    /// Returns an stream of key-value pairs with keys within a `range`, sorted
904    /// by key. [`Range`] is an (half-open) range bounded inclusively below and
905    /// exclusively above.
906    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>>;
907
908    /// Delete keys matching prefix
909    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()>;
910}
911
912#[apply(async_trait_maybe_send!)]
913impl<T> IDatabaseTransactionOpsCore for Box<T>
914where
915    T: IDatabaseTransactionOpsCore + ?Sized,
916{
917    async fn raw_insert_bytes(
918        &mut self,
919        key: &[u8],
920        value: &[u8],
921    ) -> DatabaseResult<Option<Vec<u8>>> {
922        (**self).raw_insert_bytes(key, value).await
923    }
924
925    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
926        (**self).raw_get_bytes(key).await
927    }
928
929    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
930        (**self).raw_remove_entry(key).await
931    }
932
933    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
934        (**self).raw_find_by_prefix(key_prefix).await
935    }
936
937    async fn raw_find_by_prefix_sorted_descending(
938        &mut self,
939        key_prefix: &[u8],
940    ) -> DatabaseResult<PrefixStream<'_>> {
941        (**self)
942            .raw_find_by_prefix_sorted_descending(key_prefix)
943            .await
944    }
945
946    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
947        (**self).raw_find_by_range(range).await
948    }
949
950    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
951        (**self).raw_remove_by_prefix(key_prefix).await
952    }
953}
954
955#[apply(async_trait_maybe_send!)]
956impl<T> IDatabaseTransactionOpsCore for &mut T
957where
958    T: IDatabaseTransactionOpsCore + ?Sized,
959{
960    async fn raw_insert_bytes(
961        &mut self,
962        key: &[u8],
963        value: &[u8],
964    ) -> DatabaseResult<Option<Vec<u8>>> {
965        (**self).raw_insert_bytes(key, value).await
966    }
967
968    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
969        (**self).raw_get_bytes(key).await
970    }
971
972    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
973        (**self).raw_remove_entry(key).await
974    }
975
976    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
977        (**self).raw_find_by_prefix(key_prefix).await
978    }
979
980    async fn raw_find_by_prefix_sorted_descending(
981        &mut self,
982        key_prefix: &[u8],
983    ) -> DatabaseResult<PrefixStream<'_>> {
984        (**self)
985            .raw_find_by_prefix_sorted_descending(key_prefix)
986            .await
987    }
988
989    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
990        (**self).raw_find_by_range(range).await
991    }
992
993    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
994        (**self).raw_remove_by_prefix(key_prefix).await
995    }
996}
997
998/// Additional operations (only some) database transactions expose, on top of
999/// [`IDatabaseTransactionOpsCore`]
1000///
1001/// In certain contexts exposing these operations would be a problem, so they
1002/// are moved to a separate trait.
1003pub trait IDatabaseTransactionOps: IDatabaseTransactionOpsCore + MaybeSend {}
1004
1005impl<T> IDatabaseTransactionOps for Box<T> where T: IDatabaseTransactionOps + ?Sized {}
1006
1007impl<T> IDatabaseTransactionOps for &mut T where T: IDatabaseTransactionOps + ?Sized {}
1008
1009/// Like [`IDatabaseTransactionOpsCore`], but typed
1010///
1011/// Implemented via blanket impl for everything that implements
1012/// [`IDatabaseTransactionOpsCore`] that has decoders (implements
1013/// [`WithDecoders`]).
1014#[apply(async_trait_maybe_send!)]
1015pub trait IDatabaseTransactionOpsCoreTyped<'a> {
1016    async fn get_value<K>(&mut self, key: &K) -> Option<K::Value>
1017    where
1018        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync;
1019
1020    async fn insert_entry<K>(&mut self, key: &K, value: &K::Value) -> Option<K::Value>
1021    where
1022        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1023        K::Value: MaybeSend + MaybeSync;
1024
1025    async fn insert_new_entry<K>(&mut self, key: &K, value: &K::Value)
1026    where
1027        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1028        K::Value: MaybeSend + MaybeSync;
1029
1030    async fn find_by_range<K>(
1031        &mut self,
1032        key_range: Range<K>,
1033    ) -> Pin<Box<maybe_add_send!(dyn Stream<Item = (K, K::Value)> + '_)>>
1034    where
1035        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1036        K::Value: MaybeSend + MaybeSync;
1037
1038    async fn find_by_prefix<KP>(
1039        &mut self,
1040        key_prefix: &KP,
1041    ) -> Pin<
1042        Box<
1043            maybe_add_send!(
1044                dyn Stream<
1045                        Item = (
1046                            KP::Record,
1047                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1048                        ),
1049                    > + '_
1050            ),
1051        >,
1052    >
1053    where
1054        KP: DatabaseLookup + MaybeSend + MaybeSync,
1055        KP::Record: DatabaseKey;
1056
1057    async fn find_by_prefix_sorted_descending<KP>(
1058        &mut self,
1059        key_prefix: &KP,
1060    ) -> Pin<
1061        Box<
1062            maybe_add_send!(
1063                dyn Stream<
1064                        Item = (
1065                            KP::Record,
1066                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1067                        ),
1068                    > + '_
1069            ),
1070        >,
1071    >
1072    where
1073        KP: DatabaseLookup + MaybeSend + MaybeSync,
1074        KP::Record: DatabaseKey;
1075
1076    async fn remove_entry<K>(&mut self, key: &K) -> Option<K::Value>
1077    where
1078        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync;
1079
1080    async fn remove_by_prefix<KP>(&mut self, key_prefix: &KP)
1081    where
1082        KP: DatabaseLookup + MaybeSend + MaybeSync;
1083}
1084
1085// blanket implementation of typed ops for anything that implements raw ops and
1086// has decoders
1087#[apply(async_trait_maybe_send!)]
1088impl<T> IDatabaseTransactionOpsCoreTyped<'_> for T
1089where
1090    T: IDatabaseTransactionOpsCore + WithDecoders,
1091{
1092    async fn get_value<K>(&mut self, key: &K) -> Option<K::Value>
1093    where
1094        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1095    {
1096        let key_bytes = key.to_bytes();
1097        let raw = self
1098            .raw_get_bytes(&key_bytes)
1099            .await
1100            .expect("Unrecoverable error occurred while reading and entry from the database");
1101        raw.map(|value_bytes| {
1102            decode_value_expect::<K::Value>(&value_bytes, self.decoders(), &key_bytes)
1103        })
1104    }
1105
1106    async fn insert_entry<K>(&mut self, key: &K, value: &K::Value) -> Option<K::Value>
1107    where
1108        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1109        K::Value: MaybeSend + MaybeSync,
1110    {
1111        let key_bytes = key.to_bytes();
1112        self.raw_insert_bytes(&key_bytes, &value.to_bytes())
1113            .await
1114            .expect("Unrecoverable error occurred while inserting entry into the database")
1115            .map(|value_bytes| {
1116                decode_value_expect::<K::Value>(&value_bytes, self.decoders(), &key_bytes)
1117            })
1118    }
1119
1120    async fn insert_new_entry<K>(&mut self, key: &K, value: &K::Value)
1121    where
1122        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1123        K::Value: MaybeSend + MaybeSync,
1124    {
1125        if let Some(prev) = self.insert_entry(key, value).await {
1126            panic!(
1127                "Database overwriting element when expecting insertion of new entry. Key: {key:?} Prev Value: {prev:?}"
1128            );
1129        }
1130    }
1131
1132    async fn find_by_range<K>(
1133        &mut self,
1134        key_range: Range<K>,
1135    ) -> Pin<Box<maybe_add_send!(dyn Stream<Item = (K, K::Value)> + '_)>>
1136    where
1137        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1138        K::Value: MaybeSend + MaybeSync,
1139    {
1140        let decoders = self.decoders().clone();
1141        Box::pin(
1142            self.raw_find_by_range(Range {
1143                start: &key_range.start.to_bytes(),
1144                end: &key_range.end.to_bytes(),
1145            })
1146            .await
1147            .expect("Unrecoverable error occurred while listing entries from the database")
1148            .map(move |(key_bytes, value_bytes)| {
1149                let key = decode_key_expect(&key_bytes, &decoders);
1150                let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
1151                (key, value)
1152            }),
1153        )
1154    }
1155
1156    async fn find_by_prefix<KP>(
1157        &mut self,
1158        key_prefix: &KP,
1159    ) -> Pin<
1160        Box<
1161            maybe_add_send!(
1162                dyn Stream<
1163                        Item = (
1164                            KP::Record,
1165                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1166                        ),
1167                    > + '_
1168            ),
1169        >,
1170    >
1171    where
1172        KP: DatabaseLookup + MaybeSend + MaybeSync,
1173        KP::Record: DatabaseKey,
1174    {
1175        let decoders = self.decoders().clone();
1176        Box::pin(
1177            self.raw_find_by_prefix(&key_prefix.to_bytes())
1178                .await
1179                .expect("Unrecoverable error occurred while listing entries from the database")
1180                .map(move |(key_bytes, value_bytes)| {
1181                    let key = decode_key_expect(&key_bytes, &decoders);
1182                    let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
1183                    (key, value)
1184                }),
1185        )
1186    }
1187
1188    async fn find_by_prefix_sorted_descending<KP>(
1189        &mut self,
1190        key_prefix: &KP,
1191    ) -> Pin<
1192        Box<
1193            maybe_add_send!(
1194                dyn Stream<
1195                        Item = (
1196                            KP::Record,
1197                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1198                        ),
1199                    > + '_
1200            ),
1201        >,
1202    >
1203    where
1204        KP: DatabaseLookup + MaybeSend + MaybeSync,
1205        KP::Record: DatabaseKey,
1206    {
1207        let decoders = self.decoders().clone();
1208        Box::pin(
1209            self.raw_find_by_prefix_sorted_descending(&key_prefix.to_bytes())
1210                .await
1211                .expect("Unrecoverable error occurred while listing entries from the database")
1212                .map(move |(key_bytes, value_bytes)| {
1213                    let key = decode_key_expect(&key_bytes, &decoders);
1214                    let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
1215                    (key, value)
1216                }),
1217        )
1218    }
1219    async fn remove_entry<K>(&mut self, key: &K) -> Option<K::Value>
1220    where
1221        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1222    {
1223        let key_bytes = key.to_bytes();
1224        self.raw_remove_entry(&key_bytes)
1225            .await
1226            .expect("Unrecoverable error occurred while inserting removing entry from the database")
1227            .map(|value_bytes| {
1228                decode_value_expect::<K::Value>(&value_bytes, self.decoders(), &key_bytes)
1229            })
1230    }
1231    async fn remove_by_prefix<KP>(&mut self, key_prefix: &KP)
1232    where
1233        KP: DatabaseLookup + MaybeSend + MaybeSync,
1234    {
1235        self.raw_remove_by_prefix(&key_prefix.to_bytes())
1236            .await
1237            .expect("Unrecoverable error when removing entries from the database");
1238    }
1239}
1240
1241/// A database type that has decoders, which allows it to implement
1242/// [`IDatabaseTransactionOpsCoreTyped`]
1243pub trait WithDecoders {
1244    fn decoders(&self) -> &ModuleDecoderRegistry;
1245}
1246
1247/// Raw database transaction (e.g. rocksdb implementation)
1248#[apply(async_trait_maybe_send!)]
1249pub trait IRawDatabaseTransaction: MaybeSend + IDatabaseTransactionOps {
1250    async fn commit_tx(self) -> DatabaseResult<()>;
1251}
1252
1253/// Fedimint database transaction
1254///
1255/// See [`IDatabase`] for more info.
1256#[apply(async_trait_maybe_send!)]
1257pub trait IDatabaseTransaction: MaybeSend + IDatabaseTransactionOps + fmt::Debug {
1258    /// Commit the transaction
1259    async fn commit_tx(&mut self) -> DatabaseResult<()>;
1260
1261    /// Is global database
1262    fn is_global(&self) -> bool;
1263
1264    /// Get the global database tx from a module-prefixed database transaction
1265    ///
1266    /// Meant to be called only by core internals, and module developers should
1267    /// not call it directly.
1268    #[doc(hidden)]
1269    fn global_dbtx(&mut self, access_token: GlobalDBTxAccessToken)
1270    -> &mut dyn IDatabaseTransaction;
1271}
1272
1273#[apply(async_trait_maybe_send!)]
1274impl<T> IDatabaseTransaction for Box<T>
1275where
1276    T: IDatabaseTransaction + ?Sized,
1277{
1278    async fn commit_tx(&mut self) -> DatabaseResult<()> {
1279        (**self).commit_tx().await
1280    }
1281
1282    fn is_global(&self) -> bool {
1283        (**self).is_global()
1284    }
1285
1286    fn global_dbtx(
1287        &mut self,
1288        access_token: GlobalDBTxAccessToken,
1289    ) -> &mut dyn IDatabaseTransaction {
1290        (**self).global_dbtx(access_token)
1291    }
1292}
1293
1294#[apply(async_trait_maybe_send!)]
1295impl<'a, T> IDatabaseTransaction for &'a mut T
1296where
1297    T: IDatabaseTransaction + ?Sized,
1298{
1299    async fn commit_tx(&mut self) -> DatabaseResult<()> {
1300        (**self).commit_tx().await
1301    }
1302
1303    fn is_global(&self) -> bool {
1304        (**self).is_global()
1305    }
1306
1307    fn global_dbtx(&mut self, access_key: GlobalDBTxAccessToken) -> &mut dyn IDatabaseTransaction {
1308        (**self).global_dbtx(access_key)
1309    }
1310}
1311
1312/// Struct that implements `IRawDatabaseTransaction` and can be wrapped
1313/// easier in other structs since it does not consumed `self` by move.
1314struct BaseDatabaseTransaction<Tx> {
1315    // TODO: merge options
1316    raw: Option<Tx>,
1317    notify_queue: Option<NotifyQueue>,
1318    notifications: Arc<Notifications>,
1319}
1320
1321impl<Tx> fmt::Debug for BaseDatabaseTransaction<Tx>
1322where
1323    Tx: fmt::Debug,
1324{
1325    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1326        f.write_fmt(format_args!(
1327            "BaseDatabaseTransaction{{ raw={:?} }}",
1328            self.raw
1329        ))
1330    }
1331}
1332impl<Tx> BaseDatabaseTransaction<Tx>
1333where
1334    Tx: IRawDatabaseTransaction,
1335{
1336    fn new(dbtx: Tx, notifications: Arc<Notifications>) -> Self {
1337        Self {
1338            raw: Some(dbtx),
1339            notifications,
1340            notify_queue: Some(NotifyQueue::new()),
1341        }
1342    }
1343
1344    fn add_notification_key(&mut self, key: &[u8]) -> DatabaseResult<()> {
1345        self.notify_queue
1346            .as_mut()
1347            .ok_or(DatabaseError::TransactionConsumed)?
1348            .add(key);
1349        Ok(())
1350    }
1351}
1352
1353#[apply(async_trait_maybe_send!)]
1354impl<Tx: IRawDatabaseTransaction> IDatabaseTransactionOpsCore for BaseDatabaseTransaction<Tx> {
1355    async fn raw_insert_bytes(
1356        &mut self,
1357        key: &[u8],
1358        value: &[u8],
1359    ) -> DatabaseResult<Option<Vec<u8>>> {
1360        self.add_notification_key(key)?;
1361        self.raw
1362            .as_mut()
1363            .ok_or(DatabaseError::TransactionConsumed)?
1364            .raw_insert_bytes(key, value)
1365            .await
1366    }
1367
1368    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1369        self.raw
1370            .as_mut()
1371            .ok_or(DatabaseError::TransactionConsumed)?
1372            .raw_get_bytes(key)
1373            .await
1374    }
1375
1376    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1377        self.add_notification_key(key)?;
1378        self.raw
1379            .as_mut()
1380            .ok_or(DatabaseError::TransactionConsumed)?
1381            .raw_remove_entry(key)
1382            .await
1383    }
1384
1385    async fn raw_find_by_range(
1386        &mut self,
1387        key_range: Range<&[u8]>,
1388    ) -> DatabaseResult<PrefixStream<'_>> {
1389        self.raw
1390            .as_mut()
1391            .ok_or(DatabaseError::TransactionConsumed)?
1392            .raw_find_by_range(key_range)
1393            .await
1394    }
1395
1396    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
1397        self.raw
1398            .as_mut()
1399            .ok_or(DatabaseError::TransactionConsumed)?
1400            .raw_find_by_prefix(key_prefix)
1401            .await
1402    }
1403
1404    async fn raw_find_by_prefix_sorted_descending(
1405        &mut self,
1406        key_prefix: &[u8],
1407    ) -> DatabaseResult<PrefixStream<'_>> {
1408        self.raw
1409            .as_mut()
1410            .ok_or(DatabaseError::TransactionConsumed)?
1411            .raw_find_by_prefix_sorted_descending(key_prefix)
1412            .await
1413    }
1414
1415    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
1416        self.raw
1417            .as_mut()
1418            .ok_or(DatabaseError::TransactionConsumed)?
1419            .raw_remove_by_prefix(key_prefix)
1420            .await
1421    }
1422}
1423
1424impl<Tx: IRawDatabaseTransaction> IDatabaseTransactionOps for BaseDatabaseTransaction<Tx> {}
1425
1426#[apply(async_trait_maybe_send!)]
1427impl<Tx: IRawDatabaseTransaction + fmt::Debug> IDatabaseTransaction
1428    for BaseDatabaseTransaction<Tx>
1429{
1430    async fn commit_tx(&mut self) -> DatabaseResult<()> {
1431        self.raw
1432            .take()
1433            .ok_or(DatabaseError::TransactionConsumed)?
1434            .commit_tx()
1435            .await?;
1436        self.notifications.submit_queue(
1437            &self
1438                .notify_queue
1439                .take()
1440                .expect("commit must be called only once"),
1441        );
1442        Ok(())
1443    }
1444
1445    fn is_global(&self) -> bool {
1446        true
1447    }
1448
1449    fn global_dbtx(
1450        &mut self,
1451        _access_token: GlobalDBTxAccessToken,
1452    ) -> &mut dyn IDatabaseTransaction {
1453        panic!("Illegal to call global_dbtx on BaseDatabaseTransaction");
1454    }
1455}
1456
1457/// A helper for tracking and logging on `Drop` any instances of uncommitted
1458/// writes
1459#[derive(Clone)]
1460struct CommitTracker {
1461    /// Is the dbtx committed
1462    is_committed: bool,
1463    /// Does the dbtx have any writes
1464    has_writes: bool,
1465    /// Don't warn-log uncommitted writes
1466    ignore_uncommitted: bool,
1467}
1468
1469impl Drop for CommitTracker {
1470    fn drop(&mut self) {
1471        if self.has_writes && !self.is_committed {
1472            if self.ignore_uncommitted {
1473                trace!(
1474                    target: LOG_DB,
1475                    "DatabaseTransaction has writes and has not called commit, but that's expected."
1476                );
1477            } else {
1478                warn!(
1479                    target: LOG_DB,
1480                    location = ?backtrace::Backtrace::new(),
1481                    "DatabaseTransaction has writes and has not called commit."
1482                );
1483            }
1484        }
1485    }
1486}
1487
1488enum MaybeRef<'a, T> {
1489    Owned(T),
1490    Borrowed(&'a mut T),
1491}
1492
1493impl<T> ops::Deref for MaybeRef<'_, T> {
1494    type Target = T;
1495
1496    fn deref(&self) -> &Self::Target {
1497        match self {
1498            MaybeRef::Owned(o) => o,
1499            MaybeRef::Borrowed(r) => r,
1500        }
1501    }
1502}
1503
1504impl<T> ops::DerefMut for MaybeRef<'_, T> {
1505    fn deref_mut(&mut self) -> &mut Self::Target {
1506        match self {
1507            MaybeRef::Owned(o) => o,
1508            MaybeRef::Borrowed(r) => r,
1509        }
1510    }
1511}
1512
1513/// Session type for [`DatabaseTransaction`] that is allowed to commit
1514///
1515/// Opposite of [`NonCommittable`].
1516pub struct Committable;
1517
1518/// Session type for a [`DatabaseTransaction`] that is not allowed to commit
1519///
1520/// Opposite of [`Committable`].
1521pub struct NonCommittable;
1522
1523/// A high level database transaction handle
1524///
1525/// `Cap` is a session type
1526pub struct DatabaseTransaction<'tx, Cap = NonCommittable> {
1527    tx: Box<dyn IDatabaseTransaction + 'tx>,
1528    decoders: ModuleDecoderRegistry,
1529    commit_tracker: MaybeRef<'tx, CommitTracker>,
1530    on_commit_hooks: MaybeRef<'tx, Vec<Box<maybe_add_send!(dyn FnOnce())>>>,
1531    capability: marker::PhantomData<Cap>,
1532}
1533
1534impl<Cap> fmt::Debug for DatabaseTransaction<'_, Cap> {
1535    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1536        f.write_fmt(format_args!(
1537            "DatabaseTransaction {{ tx: {:?}, decoders={:?} }}",
1538            self.tx, self.decoders
1539        ))
1540    }
1541}
1542
1543impl<Cap> WithDecoders for DatabaseTransaction<'_, Cap> {
1544    fn decoders(&self) -> &ModuleDecoderRegistry {
1545        &self.decoders
1546    }
1547}
1548
1549#[instrument(target = LOG_DB, level = "trace", skip_all, fields(value_type = std::any::type_name::<V>()), err)]
1550fn decode_value<V: DatabaseValue>(
1551    value_bytes: &[u8],
1552    decoders: &ModuleDecoderRegistry,
1553) -> std::result::Result<V, DecodingError> {
1554    trace!(
1555        bytes = %AbbreviateHexBytes(value_bytes),
1556        "decoding value",
1557    );
1558    V::from_bytes(value_bytes, decoders)
1559}
1560
1561#[track_caller]
1562fn decode_value_expect<V: DatabaseValue>(
1563    value_bytes: &[u8],
1564    decoders: &ModuleDecoderRegistry,
1565    key_bytes: &[u8],
1566) -> V {
1567    decode_value(value_bytes, decoders).unwrap_or_else(|err| {
1568        panic!(
1569            "Unrecoverable decoding DatabaseValue as {}; err={}, key_bytes={}, val_bytes={}",
1570            any::type_name::<V>(),
1571            err.fmt_compact(),
1572            AbbreviateHexBytes(key_bytes),
1573            AbbreviateHexBytes(value_bytes),
1574        )
1575    })
1576}
1577
1578#[track_caller]
1579fn decode_key_expect<K: DatabaseKey>(key_bytes: &[u8], decoders: &ModuleDecoderRegistry) -> K {
1580    trace!(
1581        bytes = %AbbreviateHexBytes(key_bytes),
1582        "decoding key",
1583    );
1584    K::from_bytes(key_bytes, decoders).unwrap_or_else(|err| {
1585        panic!(
1586            "Unrecoverable decoding DatabaseKey as {}; err={}; bytes={}",
1587            any::type_name::<K>(),
1588            err.fmt_compact(),
1589            AbbreviateHexBytes(key_bytes)
1590        )
1591    })
1592}
1593
1594impl<'tx, Cap> DatabaseTransaction<'tx, Cap> {
1595    /// Convert into a non-committable version
1596    pub fn into_nc(self) -> DatabaseTransaction<'tx, NonCommittable> {
1597        DatabaseTransaction {
1598            tx: self.tx,
1599            decoders: self.decoders,
1600            commit_tracker: self.commit_tracker,
1601            on_commit_hooks: self.on_commit_hooks,
1602            capability: PhantomData::<NonCommittable>,
1603        }
1604    }
1605
1606    /// Get a reference to a non-committeable version
1607    pub fn to_ref_nc<'s, 'a>(&'s mut self) -> DatabaseTransaction<'a, NonCommittable>
1608    where
1609        's: 'a,
1610    {
1611        self.to_ref().into_nc()
1612    }
1613
1614    /// Get [`DatabaseTransaction`] isolated to a `prefix`
1615    pub fn with_prefix<'a: 'tx>(self, prefix: Vec<u8>) -> DatabaseTransaction<'a, Cap>
1616    where
1617        'tx: 'a,
1618    {
1619        DatabaseTransaction {
1620            tx: Box::new(PrefixDatabaseTransaction {
1621                inner: self.tx,
1622                global_dbtx_access_token: None,
1623                prefix,
1624            }),
1625            decoders: self.decoders,
1626            commit_tracker: self.commit_tracker,
1627            on_commit_hooks: self.on_commit_hooks,
1628            capability: self.capability,
1629        }
1630    }
1631
1632    /// Get [`DatabaseTransaction`] isolated to a prefix of a given
1633    /// `module_instance_id`, allowing the module to access global_dbtx
1634    /// with the right access token.
1635    pub fn with_prefix_module_id<'a: 'tx>(
1636        self,
1637        module_instance_id: ModuleInstanceId,
1638    ) -> (DatabaseTransaction<'a, Cap>, GlobalDBTxAccessToken)
1639    where
1640        'tx: 'a,
1641    {
1642        let prefix = module_instance_id_to_byte_prefix(module_instance_id);
1643        let global_dbtx_access_token = GlobalDBTxAccessToken::from_prefix(&prefix);
1644        (
1645            DatabaseTransaction {
1646                tx: Box::new(PrefixDatabaseTransaction {
1647                    inner: self.tx,
1648                    global_dbtx_access_token: Some(global_dbtx_access_token),
1649                    prefix,
1650                }),
1651                decoders: self.decoders,
1652                commit_tracker: self.commit_tracker,
1653                on_commit_hooks: self.on_commit_hooks,
1654                capability: self.capability,
1655            },
1656            global_dbtx_access_token,
1657        )
1658    }
1659
1660    /// Get [`DatabaseTransaction`] to `self`
1661    pub fn to_ref<'s, 'a>(&'s mut self) -> DatabaseTransaction<'a, Cap>
1662    where
1663        's: 'a,
1664    {
1665        let decoders = self.decoders.clone();
1666
1667        DatabaseTransaction {
1668            tx: Box::new(&mut self.tx),
1669            decoders,
1670            commit_tracker: match self.commit_tracker {
1671                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1672                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1673            },
1674            on_commit_hooks: match self.on_commit_hooks {
1675                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1676                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1677            },
1678            capability: self.capability,
1679        }
1680    }
1681
1682    /// Get [`DatabaseTransaction`] isolated to a `prefix` of `self`
1683    pub fn to_ref_with_prefix<'a>(&'a mut self, prefix: Vec<u8>) -> DatabaseTransaction<'a, Cap>
1684    where
1685        'tx: 'a,
1686    {
1687        DatabaseTransaction {
1688            tx: Box::new(PrefixDatabaseTransaction {
1689                inner: &mut self.tx,
1690                global_dbtx_access_token: None,
1691                prefix,
1692            }),
1693            decoders: self.decoders.clone(),
1694            commit_tracker: match self.commit_tracker {
1695                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1696                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1697            },
1698            on_commit_hooks: match self.on_commit_hooks {
1699                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1700                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1701            },
1702            capability: self.capability,
1703        }
1704    }
1705
1706    pub fn to_ref_with_prefix_module_id<'a>(
1707        &'a mut self,
1708        module_instance_id: ModuleInstanceId,
1709    ) -> (DatabaseTransaction<'a, Cap>, GlobalDBTxAccessToken)
1710    where
1711        'tx: 'a,
1712    {
1713        let prefix = module_instance_id_to_byte_prefix(module_instance_id);
1714        let global_dbtx_access_token = GlobalDBTxAccessToken::from_prefix(&prefix);
1715        (
1716            DatabaseTransaction {
1717                tx: Box::new(PrefixDatabaseTransaction {
1718                    inner: &mut self.tx,
1719                    global_dbtx_access_token: Some(global_dbtx_access_token),
1720                    prefix,
1721                }),
1722                decoders: self.decoders.clone(),
1723                commit_tracker: match self.commit_tracker {
1724                    MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1725                    MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1726                },
1727                on_commit_hooks: match self.on_commit_hooks {
1728                    MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1729                    MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1730                },
1731                capability: self.capability,
1732            },
1733            global_dbtx_access_token,
1734        )
1735    }
1736
1737    /// Is this `Database` a global, unpartitioned `Database`
1738    pub fn is_global(&self) -> bool {
1739        self.tx.is_global()
1740    }
1741
1742    /// `Err` if [`Self::is_global`] is not true
1743    pub fn ensure_global(&self) -> DatabaseResult<()> {
1744        if !self.is_global() {
1745            return Err(DatabaseError::NotGlobal);
1746        }
1747
1748        Ok(())
1749    }
1750
1751    /// `Err` if [`Self::is_global`] is true
1752    pub fn ensure_isolated(&self) -> DatabaseResult<()> {
1753        if self.is_global() {
1754            return Err(DatabaseError::NotIsolated);
1755        }
1756
1757        Ok(())
1758    }
1759
1760    /// Cancel the tx to avoid debugging warnings about uncommitted writes
1761    pub fn ignore_uncommitted(&mut self) -> &mut Self {
1762        self.commit_tracker.ignore_uncommitted = true;
1763        self
1764    }
1765
1766    /// Create warnings about uncommitted writes
1767    pub fn warn_uncommitted(&mut self) -> &mut Self {
1768        self.commit_tracker.ignore_uncommitted = false;
1769        self
1770    }
1771
1772    /// Register a hook that will be run after commit succeeds.
1773    #[instrument(target = LOG_DB, level = "trace", skip_all)]
1774    pub fn on_commit(&mut self, f: maybe_add_send!(impl FnOnce() + 'static)) {
1775        self.on_commit_hooks.push(Box::new(f));
1776    }
1777
1778    pub fn global_dbtx<'a>(
1779        &'a mut self,
1780        access_token: GlobalDBTxAccessToken,
1781    ) -> DatabaseTransaction<'a, Cap>
1782    where
1783        'tx: 'a,
1784    {
1785        let decoders = self.decoders.clone();
1786
1787        DatabaseTransaction {
1788            tx: Box::new(self.tx.global_dbtx(access_token)),
1789            decoders,
1790            commit_tracker: match self.commit_tracker {
1791                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1792                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1793            },
1794            on_commit_hooks: match self.on_commit_hooks {
1795                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1796                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1797            },
1798            capability: self.capability,
1799        }
1800    }
1801}
1802
1803/// Code used to access `global_dbtx`
1804#[derive(Copy, Clone, Debug, PartialEq, Eq)]
1805pub struct GlobalDBTxAccessToken(u32);
1806
1807impl GlobalDBTxAccessToken {
1808    /// Calculate an access code for accessing global_dbtx from a prefixed
1809    /// database tx
1810    ///
1811    /// Since we need to do it at runtime, we want the user modules not to be
1812    /// able to call `global_dbtx` too easily. But at the same time we don't
1813    /// need to be paranoid.
1814    ///
1815    /// This must be deterministic during whole instance of the software running
1816    /// (because it's being rederived independently in multiple codepahs) , but
1817    /// it could be somewhat randomized between different runs and releases.
1818    fn from_prefix(prefix: &[u8]) -> Self {
1819        Self(prefix.iter().fold(0, |acc, b| acc + u32::from(*b)) + 513)
1820    }
1821}
1822
1823impl<'tx> DatabaseTransaction<'tx, Committable> {
1824    pub fn new(dbtx: Box<dyn IDatabaseTransaction + 'tx>, decoders: ModuleDecoderRegistry) -> Self {
1825        Self {
1826            tx: dbtx,
1827            decoders,
1828            commit_tracker: MaybeRef::Owned(CommitTracker {
1829                is_committed: false,
1830                has_writes: false,
1831                ignore_uncommitted: false,
1832            }),
1833            on_commit_hooks: MaybeRef::Owned(vec![]),
1834            capability: PhantomData,
1835        }
1836    }
1837
1838    pub async fn commit_tx_result(mut self) -> DatabaseResult<()> {
1839        self.commit_tracker.is_committed = true;
1840        let commit_result = self.tx.commit_tx().await;
1841
1842        // Run commit hooks in case commit was successful
1843        if commit_result.is_ok() {
1844            for hook in self.on_commit_hooks.deref_mut().drain(..) {
1845                hook();
1846            }
1847        }
1848
1849        commit_result
1850    }
1851
1852    pub async fn commit_tx(mut self) {
1853        self.commit_tracker.is_committed = true;
1854        self.commit_tx_result()
1855            .await
1856            .expect("Unrecoverable error occurred while committing to the database.");
1857    }
1858}
1859
1860#[apply(async_trait_maybe_send!)]
1861impl<Cap> IDatabaseTransactionOpsCore for DatabaseTransaction<'_, Cap>
1862where
1863    Cap: Send,
1864{
1865    async fn raw_insert_bytes(
1866        &mut self,
1867        key: &[u8],
1868        value: &[u8],
1869    ) -> DatabaseResult<Option<Vec<u8>>> {
1870        self.commit_tracker.has_writes = true;
1871        self.tx.raw_insert_bytes(key, value).await
1872    }
1873
1874    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1875        self.tx.raw_get_bytes(key).await
1876    }
1877
1878    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1879        self.tx.raw_remove_entry(key).await
1880    }
1881
1882    async fn raw_find_by_range(
1883        &mut self,
1884        key_range: Range<&[u8]>,
1885    ) -> DatabaseResult<PrefixStream<'_>> {
1886        self.tx.raw_find_by_range(key_range).await
1887    }
1888
1889    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
1890        self.tx.raw_find_by_prefix(key_prefix).await
1891    }
1892
1893    async fn raw_find_by_prefix_sorted_descending(
1894        &mut self,
1895        key_prefix: &[u8],
1896    ) -> DatabaseResult<PrefixStream<'_>> {
1897        self.tx
1898            .raw_find_by_prefix_sorted_descending(key_prefix)
1899            .await
1900    }
1901
1902    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
1903        self.commit_tracker.has_writes = true;
1904        self.tx.raw_remove_by_prefix(key_prefix).await
1905    }
1906}
1907impl IDatabaseTransactionOps for DatabaseTransaction<'_, Committable> {}
1908
1909impl<T> DatabaseKeyPrefix for T
1910where
1911    T: DatabaseLookup + crate::encoding::Encodable + Debug,
1912{
1913    fn to_bytes(&self) -> Vec<u8> {
1914        let mut data = vec![<Self as DatabaseLookup>::Record::DB_PREFIX];
1915        data.append(&mut self.consensus_encode_to_vec());
1916        data
1917    }
1918}
1919
1920impl<T> DatabaseKey for T
1921where
1922    // Note: key can only be `T` that can be decoded without modules (even if
1923    // module type is `()`)
1924    T: DatabaseRecord + crate::encoding::Decodable + Sized,
1925{
1926    const NOTIFY_ON_MODIFY: bool = <T as DatabaseRecord>::NOTIFY_ON_MODIFY;
1927    fn from_bytes(
1928        data: &[u8],
1929        modules: &ModuleDecoderRegistry,
1930    ) -> std::result::Result<Self, DecodingError> {
1931        if data.is_empty() {
1932            // TODO: build better coding errors, pretty useless right now
1933            return Err(DecodingError::wrong_length(1, 0));
1934        }
1935
1936        if data[0] != Self::DB_PREFIX {
1937            return Err(DecodingError::wrong_prefix(Self::DB_PREFIX, data[0]));
1938        }
1939
1940        <Self as crate::encoding::Decodable>::consensus_decode_whole(&data[1..], modules)
1941            .map_err(DecodingError::from)
1942    }
1943}
1944
1945impl<T> DatabaseValue for T
1946where
1947    T: Debug + Encodable + Decodable,
1948{
1949    fn from_bytes(
1950        data: &[u8],
1951        modules: &ModuleDecoderRegistry,
1952    ) -> std::result::Result<Self, DecodingError> {
1953        T::consensus_decode_whole(data, modules).map_err(DecodingError::from)
1954    }
1955
1956    fn to_bytes(&self) -> Vec<u8> {
1957        self.consensus_encode_to_vec()
1958    }
1959}
1960
1961/// This is a helper macro that generates the implementations of
1962/// `DatabaseRecord` necessary for reading/writing to the
1963/// database and fetching by prefix.
1964///
1965/// - `key`: This is the type of struct that will be used as the key into the
1966///   database
1967/// - `value`: This is the type of struct that will be used as the value into
1968///   the database
1969/// - `db_prefix`: Required enum expression that is represented as a `u8` and is
1970///   prepended to this key
1971/// - `query_prefix`: Optional type of struct that can be passed zero or more
1972///   times. Every query prefix can be used to query the database via
1973///   `find_by_prefix`
1974///
1975/// # Examples
1976///
1977/// ```
1978/// use fedimint_core::encoding::{Decodable, Encodable};
1979/// use fedimint_core::impl_db_record;
1980///
1981/// #[derive(Debug, Encodable, Decodable)]
1982/// struct MyKey;
1983///
1984/// #[derive(Debug, Encodable, Decodable)]
1985/// struct MyValue;
1986///
1987/// #[repr(u8)]
1988/// #[derive(Clone, Debug)]
1989/// pub enum DbKeyPrefix {
1990///     MyKey = 0x50,
1991/// }
1992///
1993/// impl_db_record!(key = MyKey, value = MyValue, db_prefix = DbKeyPrefix::MyKey);
1994/// ```
1995///
1996/// Use the required parameters and specify one `query_prefix`
1997///
1998/// ```
1999/// use fedimint_core::encoding::{Decodable, Encodable};
2000/// use fedimint_core::{impl_db_lookup, impl_db_record};
2001///
2002/// #[derive(Debug, Encodable, Decodable)]
2003/// struct MyKey;
2004///
2005/// #[derive(Debug, Encodable, Decodable)]
2006/// struct MyValue;
2007///
2008/// #[repr(u8)]
2009/// #[derive(Clone, Debug)]
2010/// pub enum DbKeyPrefix {
2011///     MyKey = 0x50,
2012/// }
2013///
2014/// #[derive(Debug, Encodable, Decodable)]
2015/// struct MyKeyPrefix;
2016///
2017/// impl_db_record!(key = MyKey, value = MyValue, db_prefix = DbKeyPrefix::MyKey,);
2018///
2019/// impl_db_lookup!(key = MyKey, query_prefix = MyKeyPrefix);
2020/// ```
2021#[macro_export]
2022macro_rules! impl_db_record {
2023    (key = $key:ty, value = $val:ty, db_prefix = $db_prefix:expr_2021 $(, notify_on_modify = $notify:tt)? $(,)?) => {
2024        impl $crate::db::DatabaseRecord for $key {
2025            const DB_PREFIX: u8 = $db_prefix as u8;
2026            $(const NOTIFY_ON_MODIFY: bool = $notify;)?
2027            type Key = Self;
2028            type Value = $val;
2029        }
2030        $(
2031            impl_db_record! {
2032                @impl_notify_marker key = $key, notify_on_modify = $notify
2033            }
2034        )?
2035    };
2036    // if notify is set to true
2037    (@impl_notify_marker key = $key:ty, notify_on_modify = true) => {
2038        impl $crate::db::DatabaseKeyWithNotify for $key {}
2039    };
2040    // if notify is set to false
2041    (@impl_notify_marker key = $key:ty, notify_on_modify = false) => {};
2042}
2043
2044#[macro_export]
2045macro_rules! impl_db_lookup{
2046    (key = $key:ty $(, query_prefix = $query_prefix:ty)* $(,)?) => {
2047        $(
2048            impl $crate::db::DatabaseLookup for $query_prefix {
2049                type Record = $key;
2050            }
2051        )*
2052    };
2053}
2054
2055/// Deprecated: Use `DatabaseVersionKey(ModuleInstanceId)` instead.
2056#[derive(Debug, Encodable, Decodable, Serialize)]
2057pub struct DatabaseVersionKeyV0;
2058
2059#[derive(Debug, Encodable, Decodable, Serialize)]
2060pub struct DatabaseVersionKey(pub ModuleInstanceId);
2061
2062#[derive(Debug, Encodable, Decodable, Serialize, Clone, PartialOrd, Ord, PartialEq, Eq, Copy)]
2063pub struct DatabaseVersion(pub u64);
2064
2065impl_db_record!(
2066    key = DatabaseVersionKeyV0,
2067    value = DatabaseVersion,
2068    db_prefix = DbKeyPrefix::DatabaseVersion
2069);
2070
2071impl_db_record!(
2072    key = DatabaseVersionKey,
2073    value = DatabaseVersion,
2074    db_prefix = DbKeyPrefix::DatabaseVersion
2075);
2076
2077impl std::fmt::Display for DatabaseVersion {
2078    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2079        write!(f, "{}", self.0)
2080    }
2081}
2082
2083impl DatabaseVersion {
2084    pub fn increment(&self) -> Self {
2085        Self(self.0 + 1)
2086    }
2087}
2088
2089impl std::fmt::Display for DbKeyPrefix {
2090    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2091        write!(f, "{self:?}")
2092    }
2093}
2094
2095#[repr(u8)]
2096#[derive(Clone, EnumIter, Debug)]
2097pub enum DbKeyPrefix {
2098    DatabaseVersion = 0x50,
2099    ClientBackup = 0x51,
2100}
2101
2102#[derive(Debug, Error)]
2103#[non_exhaustive]
2104pub enum DecodingError {
2105    #[error("Key had a wrong prefix, expected {expected} but got {found}")]
2106    WrongPrefix { expected: u8, found: u8 },
2107    #[error("Key had a wrong length, expected {expected} but got {found}")]
2108    WrongLength { expected: usize, found: usize },
2109    #[error("Other decoding error")]
2110    Other(#[source] Box<dyn Error + Send + Sync>),
2111    /// The bytes are not a valid consensus encoding of the record type.
2112    #[error("Invalid consensus encoding")]
2113    Decode(#[from] DecodeError),
2114}
2115
2116impl DecodingError {
2117    pub fn other<E: Error + Send + Sync + 'static>(error: E) -> Self {
2118        Self::Other(Box::new(error))
2119    }
2120
2121    pub fn wrong_prefix(expected: u8, found: u8) -> Self {
2122        Self::WrongPrefix { expected, found }
2123    }
2124
2125    pub fn wrong_length(expected: usize, found: usize) -> Self {
2126        Self::WrongLength { expected, found }
2127    }
2128}
2129
2130/// Error type for database operations
2131#[derive(Debug, Error)]
2132#[non_exhaustive]
2133#[cfg_attr(feature = "uniffi", derive(uniffi::Error))]
2134#[cfg_attr(feature = "uniffi", uniffi(flat_error))]
2135pub enum DatabaseError {
2136    /// Write-write conflict during optimistic transaction commit.
2137    /// This occurs when two transactions attempt to modify the same key.
2138    #[error("Write-write conflict detected")]
2139    WriteConflict,
2140
2141    /// The backend could not tell whether this transaction conflicted with
2142    /// another one, because its snapshot is older than the write history the
2143    /// backend still keeps around for conflict detection.
2144    ///
2145    /// This is *not* a conflict: it says nothing about whether any other
2146    /// transaction touched the same keys, only that the transaction was open
2147    /// for too many intervening writes to check. Nothing was written, and
2148    /// rerunning the whole operation against a fresh transaction is the only
2149    /// supported recovery — the failed transaction cannot be committed again.
2150    ///
2151    /// The usual cause is a transaction held open across something slow.
2152    #[error("Transaction snapshot is older than the retained write history: {0}")]
2153    SnapshotTooOld(Box<dyn Error + Send + Sync>),
2154
2155    /// The transaction has already been consumed (committed or dropped).
2156    /// Operations cannot be performed on a consumed transaction.
2157    #[error("Transaction already consumed")]
2158    TransactionConsumed,
2159
2160    /// Error from the underlying database backend (e.g., RocksDB I/O errors).
2161    #[error("Database backend error: {0}")]
2162    DatabaseBackend(#[from] Box<dyn Error + Send + Sync>),
2163
2164    /// The operation requires a global (non module-isolated) database instance.
2165    #[error("Database instance is not global")]
2166    NotGlobal,
2167
2168    /// The operation requires a module-isolated database instance.
2169    #[error("Database instance is not isolated")]
2170    NotIsolated,
2171}
2172
2173impl DatabaseError {
2174    /// Create a DatabaseBackend error from any error type
2175    pub fn backend<E: Error + Send + Sync + 'static>(error: E) -> Self {
2176        Self::DatabaseBackend(Box::new(error))
2177    }
2178
2179    /// Create a `SnapshotTooOld` error, preserving the backend's own error
2180    pub fn snapshot_too_old<E: Error + Send + Sync + 'static>(error: E) -> Self {
2181        Self::SnapshotTooOld(Box::new(error))
2182    }
2183}
2184
2185#[macro_export]
2186macro_rules! push_db_pair_items {
2187    ($dbtx:ident, $prefix_type:expr_2021, $key_type:ty, $value_type:ty, $map:ident, $key_literal:literal) => {
2188        let db_items =
2189            $crate::db::IDatabaseTransactionOpsCoreTyped::find_by_prefix($dbtx, &$prefix_type)
2190                .await
2191                .map(|(key, val)| {
2192                    (
2193                        $crate::encoding::Encodable::consensus_encode_to_hex(&key),
2194                        val,
2195                    )
2196                })
2197                .collect::<BTreeMap<String, $value_type>>()
2198                .await;
2199
2200        $map.insert($key_literal.to_string(), Box::new(db_items));
2201    };
2202}
2203
2204#[macro_export]
2205macro_rules! push_db_key_items {
2206    ($dbtx:ident, $prefix_type:expr_2021, $key_type:ty, $map:ident, $key_literal:literal) => {
2207        let db_items =
2208            $crate::db::IDatabaseTransactionOpsCoreTyped::find_by_prefix($dbtx, &$prefix_type)
2209                .await
2210                .map(|(key, _)| key)
2211                .collect::<Vec<$key_type>>()
2212                .await;
2213
2214        $map.insert($key_literal.to_string(), Box::new(db_items));
2215    };
2216}
2217
2218/// Context passed to the db migration _functions_ (pay attention to `Fn` in the
2219/// name)
2220///
2221/// Typically should not be referred to directly, and instead by a type-alias,
2222/// where the inner-context is set.
2223///
2224/// Notably it has the (optional) module id (inaccessible to the modules
2225/// directly, but used internally) and an inner context `C` injected by the
2226/// outer-layer.
2227///
2228/// `C` is generic, as in different layers / scopes (server vs client, etc.) a
2229/// different (module-typed, type erased, server/client, etc.) contexts might be
2230/// needed, while the database migration logic is kind of generic over that.
2231pub struct DbMigrationFnContext<'tx, C> {
2232    dbtx: DatabaseTransaction<'tx>,
2233    module_instance_id: Option<ModuleInstanceId>,
2234    ctx: C,
2235    __please_use_constructor: (),
2236}
2237
2238impl<'tx, C> DbMigrationFnContext<'tx, C> {
2239    pub fn new(
2240        dbtx: DatabaseTransaction<'tx>,
2241        module_instance_id: Option<ModuleInstanceId>,
2242        ctx: C,
2243    ) -> Self {
2244        dbtx.ensure_global().expect("Must pass global dbtx");
2245        Self {
2246            dbtx,
2247            module_instance_id,
2248            ctx,
2249            // this is a constructor
2250            __please_use_constructor: (),
2251        }
2252    }
2253
2254    pub fn map<R>(self, f: impl FnOnce(C) -> R) -> DbMigrationFnContext<'tx, R> {
2255        DbMigrationFnContext::new(self.dbtx, self.module_instance_id, f(self.ctx))
2256    }
2257
2258    // TODO: this method is currently visible to the module itself, and it shouldn't
2259    #[doc(hidden)]
2260    pub fn split_dbtx_ctx<'s>(&'s mut self) -> (&'s mut DatabaseTransaction<'tx>, &'s C) {
2261        let Self { dbtx, ctx, .. } = self;
2262
2263        (dbtx, ctx)
2264    }
2265
2266    pub fn dbtx(&'_ mut self) -> DatabaseTransaction<'_> {
2267        if let Some(module_instance_id) = self.module_instance_id {
2268            self.dbtx.to_ref_with_prefix_module_id(module_instance_id).0
2269        } else {
2270            self.dbtx.to_ref_nc()
2271        }
2272    }
2273
2274    // TODO: this method is currently visible to the module itself, and it shouldn't
2275    #[doc(hidden)]
2276    pub fn module_instance_id(&self) -> Option<ModuleInstanceId> {
2277        self.module_instance_id
2278    }
2279}
2280
2281/// [`DbMigrationFn`] with no extra context (ATM gateway)
2282pub type GeneralDbMigrationFn = DbMigrationFn<()>;
2283pub type GeneralDbMigrationFnContext<'tx> = DbMigrationFnContext<'tx, ()>;
2284
2285/// [`DbMigrationFn`] used by core client
2286///
2287/// NOTE: client _module_ migrations are handled using separate structs due to
2288/// state machine migrations
2289pub type ClientCoreDbMigrationFn = DbMigrationFn<()>;
2290pub type ClientCoreDbMigrationFnContext<'tx> = DbMigrationFnContext<'tx, ()>;
2291
2292/// `CoreMigrationFn` that modules can implement to "migrate" the database
2293/// to the next database version.
2294///
2295/// It is parametrized over `C` (contents), which is extra data/type/interface
2296/// custom for different part of the codebase, e.g.:
2297///
2298/// * server core
2299/// * server modules
2300/// * client core
2301/// * gateway core
2302pub type DbMigrationFn<C> = Box<
2303    maybe_add_send_sync!(
2304        dyn for<'tx> Fn(
2305            DbMigrationFnContext<'tx, C>,
2306        ) -> Pin<
2307            Box<maybe_add_send!(dyn futures::Future<Output = Result<(), DbMigrationError>> + 'tx)>,
2308        >
2309    ),
2310>;
2311
2312/// Failure while applying database migrations.
2313#[derive(Debug, Error)]
2314#[non_exhaustive]
2315pub enum DbMigrationError {
2316    /// The database itself failed.
2317    #[error("Database error")]
2318    Database(#[from] DatabaseError),
2319    /// A stored value is not a valid consensus encoding.
2320    #[error("Failed to consensus-decode a database entry")]
2321    Decode(#[from] DecodeError),
2322    /// The database was written by newer code than the one applying migrations.
2323    #[error(
2324        "On disk database version {on_disk} for module {kind} is higher than the code \
2325         database version {target}"
2326    )]
2327    VersionTooHigh {
2328        kind: String,
2329        on_disk: DatabaseVersion,
2330        target: DatabaseVersion,
2331    },
2332    /// A migration failed for a reason specific to it.
2333    #[error("Migration failed")]
2334    Other(#[source] Box<dyn Error + Send + Sync>),
2335}
2336
2337impl DbMigrationError {
2338    /// Wraps a migration-specific error; accepts anything convertible into a
2339    /// boxed error, an `anyhow::Error` included.
2340    pub fn other<E>(error: E) -> Self
2341    where
2342        E: Into<Box<dyn Error + Send + Sync>>,
2343    {
2344        Self::Other(error.into())
2345    }
2346}
2347
2348/// Verifies that all database migrations are defined contiguously and returns
2349/// the "current" database version, which is one greater than the last key in
2350/// the map.
2351pub fn get_current_database_version<F>(
2352    migrations: &BTreeMap<DatabaseVersion, F>,
2353) -> DatabaseVersion {
2354    let versions = migrations.keys().copied().collect::<Vec<_>>();
2355
2356    // Verify that all database migrations are defined contiguously. If there is a
2357    // gap, this indicates a programming error and we should panic.
2358    if !versions
2359        .windows(2)
2360        .all(|window| window[0].increment() == window[1])
2361    {
2362        panic!("Database Migrations are not defined contiguously");
2363    }
2364
2365    versions
2366        .last()
2367        .map_or(DatabaseVersion(0), DatabaseVersion::increment)
2368}
2369
2370pub async fn apply_migrations<C>(
2371    db: &Database,
2372    ctx: C,
2373    kind: String,
2374    migrations: BTreeMap<DatabaseVersion, DbMigrationFn<C>>,
2375    module_instance_id: Option<ModuleInstanceId>,
2376    // When used in client side context, we can/should ignore keys that external app
2377    // is allowed to use, and but since this function is shared, we make it optional argument
2378    external_prefixes_above: Option<u8>,
2379) -> Result<(), DbMigrationError>
2380where
2381    C: Clone,
2382{
2383    let mut dbtx = db.begin_transaction().await;
2384    apply_migrations_dbtx(
2385        &mut dbtx.to_ref_nc(),
2386        ctx,
2387        kind,
2388        migrations,
2389        module_instance_id,
2390        external_prefixes_above,
2391    )
2392    .await?;
2393
2394    Ok(dbtx.commit_tx_result().await?)
2395}
2396/// `apply_migrations` iterates from the on disk database version for the
2397/// module.
2398///
2399/// `apply_migrations` iterates from the on disk database version for the module
2400/// up to `target_db_version` and executes all of the migrations that exist in
2401/// the migrations map. Each migration in migrations map updates the
2402/// database to have the correct on-disk structures that the code is expecting.
2403/// The entire migration process is atomic (i.e migration from 0->1 and 1->2
2404/// happen atomically). This function is called before the module is initialized
2405/// and as long as the correct migrations are supplied in the migrations map,
2406/// the module will be able to read and write from the database successfully.
2407pub async fn apply_migrations_dbtx<C>(
2408    global_dbtx: &mut DatabaseTransaction<'_>,
2409    ctx: C,
2410    kind: String,
2411    migrations: BTreeMap<DatabaseVersion, DbMigrationFn<C>>,
2412    module_instance_id: Option<ModuleInstanceId>,
2413    // When used in client side context, we can/should ignore keys that external app
2414    // is allowed to use, and but since this function is shared, we make it optional argument
2415    external_prefixes_above: Option<u8>,
2416) -> Result<(), DbMigrationError>
2417where
2418    C: Clone,
2419{
2420    // Newly created databases will not have any data since they have just been
2421    // instantiated.
2422    let is_new_db = global_dbtx
2423        .raw_find_by_prefix(&[])
2424        .await?
2425        .filter(|(key, _v)| {
2426            std::future::ready(
2427                external_prefixes_above.is_none_or(|external_prefixes_above| {
2428                    !key.is_empty() && key[0] < external_prefixes_above
2429                }),
2430            )
2431        })
2432        .next()
2433        .await
2434        .is_none();
2435
2436    let target_db_version = get_current_database_version(&migrations);
2437
2438    // First write the database version to disk if it does not exist.
2439    create_database_version_dbtx(
2440        global_dbtx,
2441        target_db_version,
2442        module_instance_id,
2443        kind.clone(),
2444        is_new_db,
2445    )
2446    .await;
2447
2448    let module_instance_id_key = module_instance_id_or_global(module_instance_id);
2449
2450    let disk_version = global_dbtx
2451        .get_value(&DatabaseVersionKey(module_instance_id_key))
2452        .await;
2453
2454    let db_version = if let Some(disk_version) = disk_version {
2455        let mut current_db_version = disk_version;
2456
2457        if current_db_version > target_db_version {
2458            return Err(DbMigrationError::VersionTooHigh {
2459                kind,
2460                on_disk: current_db_version,
2461                target: target_db_version,
2462            });
2463        }
2464
2465        while current_db_version < target_db_version {
2466            if let Some(migration) = migrations.get(&current_db_version) {
2467                info!(target: LOG_DB, ?kind, ?current_db_version, ?target_db_version, "Migrating module...");
2468                migration(DbMigrationFnContext::new(
2469                    global_dbtx.to_ref_nc(),
2470                    module_instance_id,
2471                    ctx.clone(),
2472                ))
2473                .await?;
2474            } else {
2475                warn!(target: LOG_DB, ?current_db_version, "Missing server db migration");
2476            }
2477
2478            current_db_version = current_db_version.increment();
2479
2480            global_dbtx
2481                .insert_entry(
2482                    &DatabaseVersionKey(module_instance_id_key),
2483                    &current_db_version,
2484                )
2485                .await;
2486        }
2487
2488        current_db_version
2489    } else {
2490        target_db_version
2491    };
2492
2493    debug!(target: LOG_DB, ?kind, ?db_version, "DB Version");
2494    Ok(())
2495}
2496
2497pub async fn create_database_version(
2498    db: &Database,
2499    target_db_version: DatabaseVersion,
2500    module_instance_id: Option<ModuleInstanceId>,
2501    kind: String,
2502    is_new_db: bool,
2503) -> Result<(), DbMigrationError> {
2504    let mut dbtx = db.begin_transaction().await;
2505
2506    create_database_version_dbtx(
2507        &mut dbtx.to_ref_nc(),
2508        target_db_version,
2509        module_instance_id,
2510        kind,
2511        is_new_db,
2512    )
2513    .await;
2514
2515    dbtx.commit_tx_result().await?;
2516    Ok(())
2517}
2518
2519/// Creates the `DatabaseVersion` inside the database if it does not exist. If
2520/// necessary, this function will migrate the legacy database version to the
2521/// expected `DatabaseVersionKey`.
2522pub async fn create_database_version_dbtx(
2523    global_dbtx: &mut DatabaseTransaction<'_>,
2524    target_db_version: DatabaseVersion,
2525    module_instance_id: Option<ModuleInstanceId>,
2526    kind: String,
2527    is_new_db: bool,
2528) {
2529    let key_module_instance_id = module_instance_id_or_global(module_instance_id);
2530
2531    // First check if the module has a `DatabaseVersion` written to
2532    // `DatabaseVersionKey`. If `DatabaseVersion` already exists, there is
2533    // nothing to do.
2534    if global_dbtx
2535        .get_value(&DatabaseVersionKey(key_module_instance_id))
2536        .await
2537        .is_none()
2538    {
2539        // If it exists, read and remove the legacy `DatabaseVersion`, which used to be
2540        // in the module's isolated namespace (but not for fedimint-server or
2541        // fedimint-client).
2542        //
2543        // Otherwise, if the previous database contains data and no legacy database
2544        // version, use `DatabaseVersion(0)` so that all database migrations are
2545        // run. Otherwise, this database can assumed to be new and can use
2546        // `target_db_version` to skip the database migrations.
2547        let current_version_in_module = if let Some(module_instance_id) = module_instance_id {
2548            remove_current_db_version_if_exists(
2549                &mut global_dbtx
2550                    .to_ref_with_prefix_module_id(module_instance_id)
2551                    .0
2552                    .into_nc(),
2553                is_new_db,
2554                target_db_version,
2555            )
2556            .await
2557        } else {
2558            remove_current_db_version_if_exists(
2559                &mut global_dbtx.to_ref().into_nc(),
2560                is_new_db,
2561                target_db_version,
2562            )
2563            .await
2564        };
2565
2566        // Write the previous `DatabaseVersion` to the new `DatabaseVersionKey`
2567        debug!(target: LOG_DB, ?kind, ?current_version_in_module, ?target_db_version, ?is_new_db, "Creating DatabaseVersionKey...");
2568        global_dbtx
2569            .insert_new_entry(
2570                &DatabaseVersionKey(key_module_instance_id),
2571                &current_version_in_module,
2572            )
2573            .await;
2574    }
2575}
2576
2577/// Removes `DatabaseVersion` from `DatabaseVersionKeyV0` if it exists and
2578/// returns the current database version. If the current version does not
2579/// exist, use `target_db_version` if the database is new. Otherwise, return
2580/// `DatabaseVersion(0)` to ensure all migrations are run.
2581async fn remove_current_db_version_if_exists(
2582    version_dbtx: &mut DatabaseTransaction<'_>,
2583    is_new_db: bool,
2584    target_db_version: DatabaseVersion,
2585) -> DatabaseVersion {
2586    // Remove the previous `DatabaseVersion` in the isolated database. If it doesn't
2587    // exist, just use the 0 for the version so that all of the migrations are
2588    // executed.
2589    let current_version_in_module = version_dbtx.remove_entry(&DatabaseVersionKeyV0).await;
2590    match current_version_in_module {
2591        Some(database_version) => database_version,
2592        None if is_new_db => target_db_version,
2593        None => DatabaseVersion(0),
2594    }
2595}
2596
2597/// Helper function to retrieve the `module_instance_id` for modules, otherwise
2598/// return 0xff for the global namespace.
2599fn module_instance_id_or_global(module_instance_id: Option<ModuleInstanceId>) -> ModuleInstanceId {
2600    // Use 0xff for fedimint-server and the `module_instance_id` for each module
2601    module_instance_id.unwrap_or_else(|| MODULE_GLOBAL_PREFIX.into())
2602}
2603#[allow(unused_imports)]
2604mod test_utils {
2605    use std::collections::BTreeMap;
2606    use std::time::Duration;
2607
2608    use fedimint_core::db::DbMigrationFnContext;
2609    use futures::future::ready;
2610    use futures::{Future, FutureExt, StreamExt};
2611    use rand::Rng;
2612    use tokio::join;
2613
2614    use super::{
2615        Database, DatabaseTransaction, DatabaseVersion, DatabaseVersionKey, DatabaseVersionKeyV0,
2616        DbMigrationError, DbMigrationFn, apply_migrations, apply_migrations_dbtx,
2617        create_database_version_dbtx,
2618    };
2619    use crate::core::ModuleKind;
2620    use crate::db::mem_impl::MemDatabase;
2621    use crate::db::{
2622        IDatabaseTransactionOps, IDatabaseTransactionOpsCoreTyped, MODULE_GLOBAL_PREFIX,
2623    };
2624    use crate::encoding::{Decodable, Encodable};
2625    use crate::module::registry::ModuleDecoderRegistry;
2626
2627    pub async fn future_returns_shortly<F: Future>(fut: F) -> Option<F::Output> {
2628        crate::runtime::timeout(Duration::from_millis(10), fut)
2629            .await
2630            .ok()
2631    }
2632
2633    #[repr(u8)]
2634    #[derive(Clone)]
2635    pub enum TestDbKeyPrefix {
2636        Test = 0x42,
2637        AltTest = 0x43,
2638        PercentTestKey = 0x25,
2639    }
2640
2641    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Encodable, Decodable)]
2642    pub(super) struct TestKey(pub u64);
2643
2644    #[derive(Debug, Encodable, Decodable)]
2645    struct DbPrefixTestPrefix;
2646
2647    impl_db_record!(
2648        key = TestKey,
2649        value = TestVal,
2650        db_prefix = TestDbKeyPrefix::Test,
2651        notify_on_modify = true,
2652    );
2653    impl_db_lookup!(key = TestKey, query_prefix = DbPrefixTestPrefix);
2654
2655    #[derive(Debug, Encodable, Decodable)]
2656    struct TestKeyV0(u64, u64);
2657
2658    #[derive(Debug, Encodable, Decodable)]
2659    struct DbPrefixTestPrefixV0;
2660
2661    impl_db_record!(
2662        key = TestKeyV0,
2663        value = TestVal,
2664        db_prefix = TestDbKeyPrefix::Test,
2665    );
2666    impl_db_lookup!(key = TestKeyV0, query_prefix = DbPrefixTestPrefixV0);
2667
2668    #[derive(Debug, Eq, PartialEq, PartialOrd, Ord, Encodable, Decodable)]
2669    struct AltTestKey(u64);
2670
2671    #[derive(Debug, Encodable, Decodable)]
2672    struct AltDbPrefixTestPrefix;
2673
2674    impl_db_record!(
2675        key = AltTestKey,
2676        value = TestVal,
2677        db_prefix = TestDbKeyPrefix::AltTest,
2678    );
2679    impl_db_lookup!(key = AltTestKey, query_prefix = AltDbPrefixTestPrefix);
2680
2681    #[derive(Debug, Encodable, Decodable)]
2682    struct PercentTestKey(u64);
2683
2684    #[derive(Debug, Encodable, Decodable)]
2685    struct PercentPrefixTestPrefix;
2686
2687    impl_db_record!(
2688        key = PercentTestKey,
2689        value = TestVal,
2690        db_prefix = TestDbKeyPrefix::PercentTestKey,
2691    );
2692
2693    impl_db_lookup!(key = PercentTestKey, query_prefix = PercentPrefixTestPrefix);
2694    #[derive(Debug, Encodable, Decodable, Eq, PartialEq, PartialOrd, Ord)]
2695    pub(super) struct TestVal(pub u64);
2696
2697    const TEST_MODULE_PREFIX: u16 = 1;
2698    const ALT_MODULE_PREFIX: u16 = 2;
2699
2700    pub async fn verify_insert_elements(db: Database) {
2701        let mut dbtx = db.begin_transaction().await;
2702        assert!(dbtx.insert_entry(&TestKey(1), &TestVal(2)).await.is_none());
2703        assert!(dbtx.insert_entry(&TestKey(2), &TestVal(3)).await.is_none());
2704        dbtx.commit_tx().await;
2705
2706        // Test values were persisted
2707        let mut dbtx = db.begin_transaction().await;
2708        assert_eq!(dbtx.get_value(&TestKey(1)).await, Some(TestVal(2)));
2709        assert_eq!(dbtx.get_value(&TestKey(2)).await, Some(TestVal(3)));
2710        dbtx.commit_tx().await;
2711
2712        // Test overwrites work as expected
2713        let mut dbtx = db.begin_transaction().await;
2714        assert_eq!(
2715            dbtx.insert_entry(&TestKey(1), &TestVal(4)).await,
2716            Some(TestVal(2))
2717        );
2718        assert_eq!(
2719            dbtx.insert_entry(&TestKey(2), &TestVal(5)).await,
2720            Some(TestVal(3))
2721        );
2722        dbtx.commit_tx().await;
2723
2724        let mut dbtx = db.begin_transaction().await;
2725        assert_eq!(dbtx.get_value(&TestKey(1)).await, Some(TestVal(4)));
2726        assert_eq!(dbtx.get_value(&TestKey(2)).await, Some(TestVal(5)));
2727        dbtx.commit_tx().await;
2728    }
2729
2730    pub async fn verify_remove_nonexisting(db: Database) {
2731        let mut dbtx = db.begin_transaction().await;
2732        assert_eq!(dbtx.get_value(&TestKey(1)).await, None);
2733        let removed = dbtx.remove_entry(&TestKey(1)).await;
2734        assert!(removed.is_none());
2735
2736        // Commit to suppress the warning message
2737        dbtx.commit_tx().await;
2738    }
2739
2740    pub async fn verify_remove_existing(db: Database) {
2741        let mut dbtx = db.begin_transaction().await;
2742
2743        assert!(dbtx.insert_entry(&TestKey(1), &TestVal(2)).await.is_none());
2744
2745        assert_eq!(dbtx.get_value(&TestKey(1)).await, Some(TestVal(2)));
2746
2747        let removed = dbtx.remove_entry(&TestKey(1)).await;
2748        assert_eq!(removed, Some(TestVal(2)));
2749        assert_eq!(dbtx.get_value(&TestKey(1)).await, None);
2750
2751        // Commit to suppress the warning message
2752        dbtx.commit_tx().await;
2753    }
2754
2755    pub async fn verify_read_own_writes(db: Database) {
2756        let mut dbtx = db.begin_transaction().await;
2757
2758        assert!(dbtx.insert_entry(&TestKey(1), &TestVal(2)).await.is_none());
2759
2760        assert_eq!(dbtx.get_value(&TestKey(1)).await, Some(TestVal(2)));
2761
2762        // Commit to suppress the warning message
2763        dbtx.commit_tx().await;
2764    }
2765
2766    pub async fn verify_prevent_dirty_reads(db: Database) {
2767        let mut dbtx = db.begin_transaction().await;
2768
2769        assert!(dbtx.insert_entry(&TestKey(1), &TestVal(2)).await.is_none());
2770
2771        // dbtx2 should not be able to see uncommitted changes
2772        let mut dbtx2 = db.begin_transaction().await;
2773        assert_eq!(dbtx2.get_value(&TestKey(1)).await, None);
2774
2775        // Commit to suppress the warning message
2776        dbtx.commit_tx().await;
2777    }
2778
2779    pub async fn verify_find_by_range(db: Database) {
2780        let mut dbtx = db.begin_transaction().await;
2781        dbtx.insert_entry(&TestKey(55), &TestVal(9999)).await;
2782        dbtx.insert_entry(&TestKey(54), &TestVal(8888)).await;
2783        dbtx.insert_entry(&TestKey(56), &TestVal(7777)).await;
2784
2785        dbtx.insert_entry(&AltTestKey(55), &TestVal(7777)).await;
2786        dbtx.insert_entry(&AltTestKey(54), &TestVal(6666)).await;
2787
2788        {
2789            let mut module_dbtx = dbtx.to_ref_with_prefix_module_id(2).0;
2790            module_dbtx
2791                .insert_entry(&TestKey(300), &TestVal(3000))
2792                .await;
2793        }
2794
2795        dbtx.commit_tx().await;
2796
2797        // Verify finding by prefix returns the correct set of key pairs
2798        let mut dbtx = db.begin_transaction_nc().await;
2799
2800        let returned_keys = dbtx
2801            .find_by_range(TestKey(55)..TestKey(56))
2802            .await
2803            .collect::<Vec<_>>()
2804            .await;
2805
2806        let expected = vec![(TestKey(55), TestVal(9999))];
2807
2808        assert_eq!(returned_keys, expected);
2809
2810        let returned_keys = dbtx
2811            .find_by_range(TestKey(54)..TestKey(56))
2812            .await
2813            .collect::<Vec<_>>()
2814            .await;
2815
2816        let expected = vec![(TestKey(54), TestVal(8888)), (TestKey(55), TestVal(9999))];
2817        assert_eq!(returned_keys, expected);
2818
2819        let returned_keys = dbtx
2820            .find_by_range(TestKey(54)..TestKey(57))
2821            .await
2822            .collect::<Vec<_>>()
2823            .await;
2824
2825        let expected = vec![
2826            (TestKey(54), TestVal(8888)),
2827            (TestKey(55), TestVal(9999)),
2828            (TestKey(56), TestVal(7777)),
2829        ];
2830        assert_eq!(returned_keys, expected);
2831
2832        let mut module_dbtx = dbtx.with_prefix_module_id(2).0;
2833        let test_range = module_dbtx
2834            .find_by_range(TestKey(300)..TestKey(301))
2835            .await
2836            .collect::<Vec<_>>()
2837            .await;
2838        assert!(test_range.len() == 1);
2839    }
2840
2841    pub async fn verify_find_by_prefix(db: Database) {
2842        let mut dbtx = db.begin_transaction().await;
2843        dbtx.insert_entry(&TestKey(55), &TestVal(9999)).await;
2844        dbtx.insert_entry(&TestKey(54), &TestVal(8888)).await;
2845
2846        dbtx.insert_entry(&AltTestKey(55), &TestVal(7777)).await;
2847        dbtx.insert_entry(&AltTestKey(54), &TestVal(6666)).await;
2848        dbtx.commit_tx().await;
2849
2850        // Verify finding by prefix returns the correct set of key pairs
2851        let mut dbtx = db.begin_transaction().await;
2852
2853        let returned_keys = dbtx
2854            .find_by_prefix(&DbPrefixTestPrefix)
2855            .await
2856            .collect::<Vec<_>>()
2857            .await;
2858
2859        let expected = vec![(TestKey(54), TestVal(8888)), (TestKey(55), TestVal(9999))];
2860        assert_eq!(returned_keys, expected);
2861
2862        let reversed = dbtx
2863            .find_by_prefix_sorted_descending(&DbPrefixTestPrefix)
2864            .await
2865            .collect::<Vec<_>>()
2866            .await;
2867        let mut reversed_expected = expected;
2868        reversed_expected.reverse();
2869        assert_eq!(reversed, reversed_expected);
2870
2871        let returned_keys = dbtx
2872            .find_by_prefix(&AltDbPrefixTestPrefix)
2873            .await
2874            .collect::<Vec<_>>()
2875            .await;
2876
2877        let expected = vec![
2878            (AltTestKey(54), TestVal(6666)),
2879            (AltTestKey(55), TestVal(7777)),
2880        ];
2881        assert_eq!(returned_keys, expected);
2882
2883        let reversed = dbtx
2884            .find_by_prefix_sorted_descending(&AltDbPrefixTestPrefix)
2885            .await
2886            .collect::<Vec<_>>()
2887            .await;
2888        let mut reversed_expected = expected;
2889        reversed_expected.reverse();
2890        assert_eq!(reversed, reversed_expected);
2891    }
2892
2893    pub async fn verify_commit(db: Database) {
2894        let mut dbtx = db.begin_transaction().await;
2895
2896        assert!(dbtx.insert_entry(&TestKey(1), &TestVal(2)).await.is_none());
2897        dbtx.commit_tx().await;
2898
2899        // Verify dbtx2 can see committed transactions
2900        let mut dbtx2 = db.begin_transaction().await;
2901        assert_eq!(dbtx2.get_value(&TestKey(1)).await, Some(TestVal(2)));
2902    }
2903
2904    pub async fn verify_prevent_nonrepeatable_reads(db: Database) {
2905        let mut dbtx = db.begin_transaction().await;
2906        assert_eq!(dbtx.get_value(&TestKey(100)).await, None);
2907
2908        let mut dbtx2 = db.begin_transaction().await;
2909
2910        dbtx2.insert_entry(&TestKey(100), &TestVal(101)).await;
2911
2912        assert_eq!(dbtx.get_value(&TestKey(100)).await, None);
2913
2914        dbtx2.commit_tx().await;
2915
2916        // dbtx should still read None because it is operating over a snapshot
2917        // of the data when the transaction started
2918        assert_eq!(dbtx.get_value(&TestKey(100)).await, None);
2919
2920        let expected_keys = 0;
2921        let returned_keys = dbtx
2922            .find_by_prefix(&DbPrefixTestPrefix)
2923            .await
2924            .fold(0, |returned_keys, (key, value)| async move {
2925                if key == TestKey(100) {
2926                    assert!(value.eq(&TestVal(101)));
2927                }
2928                returned_keys + 1
2929            })
2930            .await;
2931
2932        assert_eq!(returned_keys, expected_keys);
2933    }
2934
2935    pub async fn verify_snapshot_isolation(db: Database) {
2936        async fn random_yield() {
2937            let times = if rand::thread_rng().gen_bool(0.5) {
2938                0
2939            } else {
2940                10
2941            };
2942            for _ in 0..times {
2943                tokio::task::yield_now().await;
2944            }
2945        }
2946
2947        // This scenario is taken straight out of https://github.com/fedimint/fedimint/issues/5195 bug
2948        for i in 0..1000 {
2949            let base_key = i * 2;
2950            let tx_accepted_key = base_key;
2951            let spent_input_key = base_key + 1;
2952
2953            join!(
2954                async {
2955                    random_yield().await;
2956                    let mut dbtx = db.begin_transaction().await;
2957
2958                    random_yield().await;
2959                    let a = dbtx.get_value(&TestKey(tx_accepted_key)).await;
2960                    random_yield().await;
2961                    // we have 4 operations that can give you the db key,
2962                    // try all of them
2963                    let s = match i % 5 {
2964                        0 => dbtx.get_value(&TestKey(spent_input_key)).await,
2965                        1 => dbtx.remove_entry(&TestKey(spent_input_key)).await,
2966                        2 => {
2967                            dbtx.insert_entry(&TestKey(spent_input_key), &TestVal(200))
2968                                .await
2969                        }
2970                        3 => {
2971                            dbtx.find_by_prefix(&DbPrefixTestPrefix)
2972                                .await
2973                                .filter(|(k, _v)| ready(k == &TestKey(spent_input_key)))
2974                                .map(|(_k, v)| v)
2975                                .next()
2976                                .await
2977                        }
2978                        4 => {
2979                            dbtx.find_by_prefix_sorted_descending(&DbPrefixTestPrefix)
2980                                .await
2981                                .filter(|(k, _v)| ready(k == &TestKey(spent_input_key)))
2982                                .map(|(_k, v)| v)
2983                                .next()
2984                                .await
2985                        }
2986                        _ => {
2987                            panic!("woot?");
2988                        }
2989                    };
2990
2991                    match (a, s) {
2992                        (None, None) | (Some(_), Some(_)) => {}
2993                        (None, Some(_)) => panic!("none some?! {i}"),
2994                        (Some(_), None) => panic!("some none?! {i}"),
2995                    }
2996                },
2997                async {
2998                    random_yield().await;
2999
3000                    let mut dbtx = db.begin_transaction().await;
3001                    random_yield().await;
3002                    assert_eq!(dbtx.get_value(&TestKey(tx_accepted_key)).await, None);
3003
3004                    random_yield().await;
3005                    assert_eq!(
3006                        dbtx.insert_entry(&TestKey(spent_input_key), &TestVal(100))
3007                            .await,
3008                        None
3009                    );
3010
3011                    random_yield().await;
3012                    assert_eq!(
3013                        dbtx.insert_entry(&TestKey(tx_accepted_key), &TestVal(100))
3014                            .await,
3015                        None
3016                    );
3017                    random_yield().await;
3018                    dbtx.commit_tx().await;
3019                }
3020            );
3021        }
3022    }
3023
3024    pub async fn verify_phantom_entry(db: Database) {
3025        let mut dbtx = db.begin_transaction().await;
3026
3027        dbtx.insert_entry(&TestKey(100), &TestVal(101)).await;
3028
3029        dbtx.insert_entry(&TestKey(101), &TestVal(102)).await;
3030
3031        dbtx.commit_tx().await;
3032
3033        let mut dbtx = db.begin_transaction().await;
3034        let expected_keys = 2;
3035        let returned_keys = dbtx
3036            .find_by_prefix(&DbPrefixTestPrefix)
3037            .await
3038            .fold(0, |returned_keys, (key, value)| async move {
3039                match key {
3040                    TestKey(100) => {
3041                        assert!(value.eq(&TestVal(101)));
3042                    }
3043                    TestKey(101) => {
3044                        assert!(value.eq(&TestVal(102)));
3045                    }
3046                    _ => {}
3047                }
3048                returned_keys + 1
3049            })
3050            .await;
3051
3052        assert_eq!(returned_keys, expected_keys);
3053
3054        let mut dbtx2 = db.begin_transaction().await;
3055
3056        dbtx2.insert_entry(&TestKey(102), &TestVal(103)).await;
3057
3058        dbtx2.commit_tx().await;
3059
3060        let returned_keys = dbtx
3061            .find_by_prefix(&DbPrefixTestPrefix)
3062            .await
3063            .fold(0, |returned_keys, (key, value)| async move {
3064                match key {
3065                    TestKey(100) => {
3066                        assert!(value.eq(&TestVal(101)));
3067                    }
3068                    TestKey(101) => {
3069                        assert!(value.eq(&TestVal(102)));
3070                    }
3071                    _ => {}
3072                }
3073                returned_keys + 1
3074            })
3075            .await;
3076
3077        assert_eq!(returned_keys, expected_keys);
3078    }
3079
3080    pub async fn expect_write_conflict(db: Database) {
3081        let mut dbtx = db.begin_transaction().await;
3082        dbtx.insert_entry(&TestKey(100), &TestVal(101)).await;
3083        dbtx.commit_tx().await;
3084
3085        let mut dbtx2 = db.begin_transaction().await;
3086        let mut dbtx3 = db.begin_transaction().await;
3087
3088        dbtx2.insert_entry(&TestKey(100), &TestVal(102)).await;
3089
3090        // Depending on if the database implementation supports optimistic or
3091        // pessimistic transactions, this test should generate an error here
3092        // (pessimistic) or at commit time (optimistic)
3093        dbtx3.insert_entry(&TestKey(100), &TestVal(103)).await;
3094
3095        dbtx2.commit_tx().await;
3096        dbtx3.commit_tx_result().await.expect_err("Expecting an error to be returned because this transaction is in a write-write conflict with dbtx");
3097    }
3098
3099    pub async fn verify_string_prefix(db: Database) {
3100        let mut dbtx = db.begin_transaction().await;
3101        dbtx.insert_entry(&PercentTestKey(100), &TestVal(101)).await;
3102
3103        assert_eq!(
3104            dbtx.get_value(&PercentTestKey(100)).await,
3105            Some(TestVal(101))
3106        );
3107
3108        dbtx.insert_entry(&PercentTestKey(101), &TestVal(100)).await;
3109
3110        dbtx.insert_entry(&PercentTestKey(101), &TestVal(100)).await;
3111
3112        dbtx.insert_entry(&PercentTestKey(101), &TestVal(100)).await;
3113
3114        // If the wildcard character ('%') is not handled properly, this will make
3115        // find_by_prefix return 5 results instead of 4
3116        dbtx.insert_entry(&TestKey(101), &TestVal(100)).await;
3117
3118        let expected_keys = 4;
3119        let returned_keys = dbtx
3120            .find_by_prefix(&PercentPrefixTestPrefix)
3121            .await
3122            .fold(0, |returned_keys, (key, value)| async move {
3123                if matches!(key, PercentTestKey(101)) {
3124                    assert!(value.eq(&TestVal(100)));
3125                }
3126                returned_keys + 1
3127            })
3128            .await;
3129
3130        assert_eq!(returned_keys, expected_keys);
3131    }
3132
3133    pub async fn verify_remove_by_prefix(db: Database) {
3134        let mut dbtx = db.begin_transaction().await;
3135
3136        dbtx.insert_entry(&TestKey(100), &TestVal(101)).await;
3137
3138        dbtx.insert_entry(&TestKey(101), &TestVal(102)).await;
3139
3140        dbtx.commit_tx().await;
3141
3142        let mut remove_dbtx = db.begin_transaction().await;
3143        remove_dbtx.remove_by_prefix(&DbPrefixTestPrefix).await;
3144        remove_dbtx.commit_tx().await;
3145
3146        let mut dbtx = db.begin_transaction().await;
3147        let expected_keys = 0;
3148        let returned_keys = dbtx
3149            .find_by_prefix(&DbPrefixTestPrefix)
3150            .await
3151            .fold(0, |returned_keys, (key, value)| async move {
3152                match key {
3153                    TestKey(100) => {
3154                        assert!(value.eq(&TestVal(101)));
3155                    }
3156                    TestKey(101) => {
3157                        assert!(value.eq(&TestVal(102)));
3158                    }
3159                    _ => {}
3160                }
3161                returned_keys + 1
3162            })
3163            .await;
3164
3165        assert_eq!(returned_keys, expected_keys);
3166    }
3167
3168    pub async fn verify_module_db(db: Database, module_db: Database) {
3169        let mut dbtx = db.begin_transaction().await;
3170
3171        dbtx.insert_entry(&TestKey(100), &TestVal(101)).await;
3172
3173        dbtx.insert_entry(&TestKey(101), &TestVal(102)).await;
3174
3175        dbtx.commit_tx().await;
3176
3177        // verify module_dbtx can only read key/value pairs from its own module
3178        let mut module_dbtx = module_db.begin_transaction().await;
3179        assert_eq!(module_dbtx.get_value(&TestKey(100)).await, None);
3180
3181        assert_eq!(module_dbtx.get_value(&TestKey(101)).await, None);
3182
3183        // verify module_dbtx can read key/value pairs that it wrote
3184        let mut dbtx = db.begin_transaction().await;
3185        assert_eq!(dbtx.get_value(&TestKey(100)).await, Some(TestVal(101)));
3186
3187        assert_eq!(dbtx.get_value(&TestKey(101)).await, Some(TestVal(102)));
3188
3189        let mut module_dbtx = module_db.begin_transaction().await;
3190
3191        module_dbtx.insert_entry(&TestKey(100), &TestVal(103)).await;
3192
3193        module_dbtx.insert_entry(&TestKey(101), &TestVal(104)).await;
3194
3195        module_dbtx.commit_tx().await;
3196
3197        let expected_keys = 2;
3198        let mut dbtx = db.begin_transaction().await;
3199        let returned_keys = dbtx
3200            .find_by_prefix(&DbPrefixTestPrefix)
3201            .await
3202            .fold(0, |returned_keys, (key, value)| async move {
3203                match key {
3204                    TestKey(100) => {
3205                        assert!(value.eq(&TestVal(101)));
3206                    }
3207                    TestKey(101) => {
3208                        assert!(value.eq(&TestVal(102)));
3209                    }
3210                    _ => {}
3211                }
3212                returned_keys + 1
3213            })
3214            .await;
3215
3216        assert_eq!(returned_keys, expected_keys);
3217
3218        let removed = dbtx.remove_entry(&TestKey(100)).await;
3219        assert_eq!(removed, Some(TestVal(101)));
3220        assert_eq!(dbtx.get_value(&TestKey(100)).await, None);
3221
3222        let mut module_dbtx = module_db.begin_transaction().await;
3223        assert_eq!(
3224            module_dbtx.get_value(&TestKey(100)).await,
3225            Some(TestVal(103))
3226        );
3227    }
3228
3229    pub async fn verify_module_prefix(db: Database) {
3230        let mut test_dbtx = db.begin_transaction().await;
3231        {
3232            let mut test_module_dbtx = test_dbtx.to_ref_with_prefix_module_id(TEST_MODULE_PREFIX).0;
3233
3234            test_module_dbtx
3235                .insert_entry(&TestKey(100), &TestVal(101))
3236                .await;
3237
3238            test_module_dbtx
3239                .insert_entry(&TestKey(101), &TestVal(102))
3240                .await;
3241        }
3242
3243        test_dbtx.commit_tx().await;
3244
3245        let mut alt_dbtx = db.begin_transaction().await;
3246        {
3247            let mut alt_module_dbtx = alt_dbtx.to_ref_with_prefix_module_id(ALT_MODULE_PREFIX).0;
3248
3249            alt_module_dbtx
3250                .insert_entry(&TestKey(100), &TestVal(103))
3251                .await;
3252
3253            alt_module_dbtx
3254                .insert_entry(&TestKey(101), &TestVal(104))
3255                .await;
3256        }
3257
3258        alt_dbtx.commit_tx().await;
3259
3260        // verify test_module_dbtx can only see key/value pairs from its own module
3261        let mut test_dbtx = db.begin_transaction().await;
3262        let mut test_module_dbtx = test_dbtx.to_ref_with_prefix_module_id(TEST_MODULE_PREFIX).0;
3263        assert_eq!(
3264            test_module_dbtx.get_value(&TestKey(100)).await,
3265            Some(TestVal(101))
3266        );
3267
3268        assert_eq!(
3269            test_module_dbtx.get_value(&TestKey(101)).await,
3270            Some(TestVal(102))
3271        );
3272
3273        let expected_keys = 2;
3274        let returned_keys = test_module_dbtx
3275            .find_by_prefix(&DbPrefixTestPrefix)
3276            .await
3277            .fold(0, |returned_keys, (key, value)| async move {
3278                match key {
3279                    TestKey(100) => {
3280                        assert!(value.eq(&TestVal(101)));
3281                    }
3282                    TestKey(101) => {
3283                        assert!(value.eq(&TestVal(102)));
3284                    }
3285                    _ => {}
3286                }
3287                returned_keys + 1
3288            })
3289            .await;
3290
3291        assert_eq!(returned_keys, expected_keys);
3292
3293        let removed = test_module_dbtx.remove_entry(&TestKey(100)).await;
3294        assert_eq!(removed, Some(TestVal(101)));
3295        assert_eq!(test_module_dbtx.get_value(&TestKey(100)).await, None);
3296
3297        // test_dbtx on its own wont find the key because it does not use a module
3298        // prefix
3299        let mut test_dbtx = db.begin_transaction().await;
3300        assert_eq!(test_dbtx.get_value(&TestKey(101)).await, None);
3301
3302        test_dbtx.commit_tx().await;
3303    }
3304
3305    #[cfg(test)]
3306    #[tokio::test]
3307    pub async fn verify_test_migration() {
3308        // Insert a bunch of old dummy data that needs to be migrated to a new version
3309        let db = Database::new(MemDatabase::new(), ModuleDecoderRegistry::default());
3310        let expected_test_keys_size: usize = 100;
3311        let mut dbtx = db.begin_transaction().await;
3312        for i in 0..expected_test_keys_size {
3313            dbtx.insert_new_entry(&TestKeyV0(i as u64, (i + 1) as u64), &TestVal(i as u64))
3314                .await;
3315        }
3316
3317        // Will also be migrated to `DatabaseVersionKey`
3318        dbtx.insert_new_entry(&DatabaseVersionKeyV0, &DatabaseVersion(0))
3319            .await;
3320        dbtx.commit_tx().await;
3321
3322        let mut migrations: BTreeMap<DatabaseVersion, DbMigrationFn<()>> = BTreeMap::new();
3323
3324        migrations.insert(
3325            DatabaseVersion(0),
3326            Box::new(|ctx| migrate_test_db_version_0(ctx).boxed()),
3327        );
3328
3329        apply_migrations(&db, (), "TestModule".to_string(), migrations, None, None)
3330            .await
3331            .expect("Error applying migrations for TestModule");
3332
3333        // Verify that the migrations completed successfully
3334        let mut dbtx = db.begin_transaction().await;
3335
3336        // Verify that the old `DatabaseVersion` under `DatabaseVersionKeyV0` migrated
3337        // to `DatabaseVersionKey`
3338        assert!(
3339            dbtx.get_value(&DatabaseVersionKey(MODULE_GLOBAL_PREFIX.into()))
3340                .await
3341                .is_some()
3342        );
3343
3344        // Verify Dummy module migration
3345        let test_keys = dbtx
3346            .find_by_prefix(&DbPrefixTestPrefix)
3347            .await
3348            .collect::<Vec<_>>()
3349            .await;
3350        let test_keys_size = test_keys.len();
3351        assert_eq!(test_keys_size, expected_test_keys_size);
3352        for (key, val) in test_keys {
3353            assert_eq!(key.0, val.0 + 1);
3354        }
3355    }
3356
3357    #[cfg(test)]
3358    #[tokio::test]
3359    async fn apply_migrations_rejects_newer_on_disk_version() {
3360        let db = Database::new(MemDatabase::new(), ModuleDecoderRegistry::default());
3361        let mut dbtx = db.begin_transaction().await;
3362
3363        // Pretend the code that wrote this database was at version 5.
3364        create_database_version_dbtx(
3365            &mut dbtx.to_ref_nc(),
3366            DatabaseVersion(5),
3367            None,
3368            "test".to_owned(),
3369            true,
3370        )
3371        .await;
3372
3373        // This code knows no migrations at all, i.e. it is at version 0.
3374        let err = apply_migrations_dbtx(
3375            &mut dbtx.to_ref_nc(),
3376            (),
3377            "test".to_owned(),
3378            BTreeMap::new(),
3379            None,
3380            None,
3381        )
3382        .await
3383        .expect_err("on-disk version 5 must not be accepted by code at version 0");
3384
3385        assert!(matches!(
3386            err,
3387            DbMigrationError::VersionTooHigh {
3388                on_disk: DatabaseVersion(5),
3389                target: DatabaseVersion(0),
3390                ..
3391            }
3392        ));
3393    }
3394
3395    #[cfg(test)]
3396    #[test]
3397    fn db_migration_error_other_accepts_anyhow() {
3398        let err = DbMigrationError::other(anyhow::anyhow!("legacy"));
3399
3400        assert!(matches!(
3401            &err,
3402            DbMigrationError::Other(source) if source.to_string() == "legacy"
3403        ));
3404        assert!(std::error::Error::source(&err).is_some());
3405    }
3406
3407    #[allow(dead_code)]
3408    async fn migrate_test_db_version_0(
3409        mut ctx: DbMigrationFnContext<'_, ()>,
3410    ) -> Result<(), DbMigrationError> {
3411        let mut dbtx = ctx.dbtx();
3412        let example_keys_v0 = dbtx
3413            .find_by_prefix(&DbPrefixTestPrefixV0)
3414            .await
3415            .collect::<Vec<_>>()
3416            .await;
3417        dbtx.remove_by_prefix(&DbPrefixTestPrefixV0).await;
3418        for (key, val) in example_keys_v0 {
3419            let key_v2 = TestKey(key.1);
3420            dbtx.insert_new_entry(&key_v2, &val).await;
3421        }
3422        Ok(())
3423    }
3424
3425    #[cfg(test)]
3426    #[tokio::test]
3427    async fn test_autocommit() {
3428        use std::marker::PhantomData;
3429        use std::ops::Range;
3430        use std::path::Path;
3431
3432        use async_trait::async_trait;
3433
3434        use crate::ModuleDecoderRegistry;
3435        use crate::db::{
3436            AutocommitError, BaseDatabaseTransaction, DatabaseError, DatabaseResult,
3437            IDatabaseTransaction, IDatabaseTransactionOps, IDatabaseTransactionOpsCore,
3438            IRawDatabase, IRawDatabaseTransaction,
3439        };
3440
3441        #[derive(Debug)]
3442        struct FakeDatabase;
3443
3444        #[async_trait]
3445        impl IRawDatabase for FakeDatabase {
3446            type Transaction<'a> = FakeTransaction<'a>;
3447            async fn begin_transaction(&self) -> FakeTransaction {
3448                FakeTransaction(PhantomData)
3449            }
3450
3451            fn checkpoint(&self, _backup_path: &Path) -> DatabaseResult<()> {
3452                Ok(())
3453            }
3454        }
3455
3456        #[derive(Debug)]
3457        struct FakeTransaction<'a>(PhantomData<&'a ()>);
3458
3459        #[async_trait]
3460        impl IDatabaseTransactionOpsCore for FakeTransaction<'_> {
3461            async fn raw_insert_bytes(
3462                &mut self,
3463                _key: &[u8],
3464                _value: &[u8],
3465            ) -> DatabaseResult<Option<Vec<u8>>> {
3466                unimplemented!()
3467            }
3468
3469            async fn raw_get_bytes(&mut self, _key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
3470                unimplemented!()
3471            }
3472
3473            async fn raw_remove_entry(&mut self, _key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
3474                unimplemented!()
3475            }
3476
3477            async fn raw_find_by_range(
3478                &mut self,
3479                _key_range: Range<&[u8]>,
3480            ) -> DatabaseResult<crate::db::PrefixStream<'_>> {
3481                unimplemented!()
3482            }
3483
3484            async fn raw_find_by_prefix(
3485                &mut self,
3486                _key_prefix: &[u8],
3487            ) -> DatabaseResult<crate::db::PrefixStream<'_>> {
3488                unimplemented!()
3489            }
3490
3491            async fn raw_remove_by_prefix(&mut self, _key_prefix: &[u8]) -> DatabaseResult<()> {
3492                unimplemented!()
3493            }
3494
3495            async fn raw_find_by_prefix_sorted_descending(
3496                &mut self,
3497                _key_prefix: &[u8],
3498            ) -> DatabaseResult<crate::db::PrefixStream<'_>> {
3499                unimplemented!()
3500            }
3501        }
3502
3503        impl IDatabaseTransactionOps for FakeTransaction<'_> {}
3504
3505        #[async_trait]
3506        impl IRawDatabaseTransaction for FakeTransaction<'_> {
3507            async fn commit_tx(self) -> DatabaseResult<()> {
3508                use crate::db::DatabaseError;
3509
3510                Err(DatabaseError::backend(std::io::Error::other(
3511                    "Can't commit!",
3512                )))
3513            }
3514        }
3515
3516        let db = Database::new(FakeDatabase, ModuleDecoderRegistry::default());
3517        let err = db
3518            .autocommit::<_, _, ()>(|_dbtx, _| Box::pin(async { Ok(()) }), Some(5))
3519            .await
3520            .unwrap_err();
3521
3522        match err {
3523            AutocommitError::CommitFailed {
3524                attempts: failed_attempts,
3525                ..
3526            } => {
3527                assert_eq!(failed_attempts, 5);
3528            }
3529            AutocommitError::ClosureError { .. } => panic!("Closure did not return error"),
3530        }
3531    }
3532}
3533
3534pub async fn find_by_prefix_sorted_descending<'r, 'inner, KP>(
3535    tx: &'r mut (dyn IDatabaseTransaction + 'inner),
3536    decoders: ModuleDecoderRegistry,
3537    key_prefix: &KP,
3538) -> impl Stream<
3539    Item = (
3540        KP::Record,
3541        <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
3542    ),
3543>
3544+ 'r
3545+ use<'r, KP>
3546where
3547    'inner: 'r,
3548    KP: DatabaseLookup,
3549    KP::Record: DatabaseKey,
3550{
3551    debug!(target: LOG_DB, "find by prefix sorted descending");
3552    let prefix_bytes = key_prefix.to_bytes();
3553    tx.raw_find_by_prefix_sorted_descending(&prefix_bytes)
3554        .await
3555        .expect("Error doing prefix search in database")
3556        .map(move |(key_bytes, value_bytes)| {
3557            let key = decode_key_expect(&key_bytes, &decoders);
3558            let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
3559            (key, value)
3560        })
3561}
3562
3563pub async fn verify_module_db_integrity_dbtx(
3564    dbtx: &mut DatabaseTransaction<'_>,
3565    module_id: ModuleInstanceId,
3566    module_kind: ModuleKind,
3567    prefixes: &BTreeSet<u8>,
3568) {
3569    let module_db_prefix = module_instance_id_to_byte_prefix(module_id);
3570    if module_id < 250 {
3571        assert_eq!(module_db_prefix.len(), 2);
3572    }
3573    let mut records = dbtx
3574        .raw_find_by_prefix(&module_db_prefix)
3575        .await
3576        .expect("DB fail");
3577    while let Some((k, v)) = records.next().await {
3578        assert!(
3579            prefixes.contains(&k[module_db_prefix.len()]),
3580            "Unexpected module {module_kind} {module_id} db record found: {}: {}",
3581            k.as_hex(),
3582            v.as_hex()
3583        );
3584    }
3585}
3586
3587#[cfg(test)]
3588mod tests;