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//!             Ok::<(), std::convert::Infallible>(())
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::borrow::Cow;
107use std::collections::{BTreeMap, BTreeSet};
108use std::error::Error;
109use std::fmt::{self, Debug};
110use std::marker::{self, PhantomData};
111use std::ops::{self, DerefMut, Range};
112use std::path::Path;
113use std::pin::Pin;
114use std::sync::Arc;
115use std::time::Duration;
116
117use bitcoin::hex::DisplayHex as _;
118use fedimint_core::util::BoxFuture;
119use fedimint_logging::LOG_DB;
120use fedimint_util_error::FmtCompact as _;
121use futures::{Stream, StreamExt};
122use macro_rules_attribute::apply;
123use rand::Rng;
124use serde::Serialize;
125use strum_macros::EnumIter;
126use thiserror::Error;
127use tracing::{debug, info, instrument, trace, warn};
128
129use crate::core::{ModuleInstanceId, ModuleKind};
130use crate::encoding::{Decodable, DecodeError, Encodable};
131use crate::fmt_utils::AbbreviateHexBytes;
132use crate::task::{MaybeSend, MaybeSync};
133use crate::{async_trait_maybe_send, maybe_add_send, maybe_add_send_sync, timing};
134
135pub mod mem_impl;
136pub mod notifications;
137
138pub use test_utils::*;
139
140use self::notifications::{Notifications, NotifyQueue};
141use crate::module::registry::{ModuleDecoderRegistry, ModuleRegistry};
142
143pub const MODULE_GLOBAL_PREFIX: u8 = 0xff;
144
145/// Result type for database operations
146pub type DatabaseResult<T> = std::result::Result<T, DatabaseError>;
147
148pub trait DatabaseKeyPrefix: Debug {
149    fn to_bytes(&self) -> Vec<u8>;
150}
151
152/// A key + value pair in the database with a unique prefix
153/// Extends `DatabaseKeyPrefix` to prepend the key's prefix.
154pub trait DatabaseRecord: DatabaseKeyPrefix {
155    const DB_PREFIX: u8;
156    const NOTIFY_ON_MODIFY: bool = false;
157    type Key: DatabaseKey + Debug;
158    type Value: DatabaseValue + Debug;
159}
160
161/// A key that can be used to query one or more `DatabaseRecord`
162/// Extends `DatabaseKeyPrefix` to prepend the key's prefix.
163pub trait DatabaseLookup: DatabaseKeyPrefix {
164    type Record: DatabaseRecord;
165}
166
167// Every `DatabaseRecord` is automatically a `DatabaseLookup`
168impl<Record> DatabaseLookup for Record
169where
170    Record: DatabaseRecord + Debug + Decodable + Encodable,
171{
172    type Record = Record;
173}
174
175/// `DatabaseKey` that represents the lookup structure for retrieving key/value
176/// pairs from the database.
177pub trait DatabaseKey: Sized {
178    /// Send a notification to tasks waiting to be notified if the value of
179    /// `DatabaseKey` is modified
180    ///
181    /// For instance, this can be used to be notified when a key in the
182    /// database is created. It is also possible to run a closure with the
183    /// value of the `DatabaseKey` as parameter to verify some changes to
184    /// that value.
185    const NOTIFY_ON_MODIFY: bool = false;
186    fn from_bytes(
187        data: &[u8],
188        modules: &ModuleDecoderRegistry,
189    ) -> std::result::Result<Self, DecodingError>;
190}
191
192/// Marker trait for `DatabaseKey`s where `NOTIFY` is true
193pub trait DatabaseKeyWithNotify {}
194
195/// `DatabaseValue` that represents the value structure of database records.
196pub trait DatabaseValue: Sized + Debug {
197    fn from_bytes(
198        data: &[u8],
199        modules: &ModuleDecoderRegistry,
200    ) -> std::result::Result<Self, DecodingError>;
201    fn to_bytes(&self) -> Vec<u8>;
202}
203
204pub type PrefixStream<'a> = Pin<Box<maybe_add_send!(dyn Stream<Item = (Vec<u8>, Vec<u8>)> + 'a)>>;
205
206/// Just ignore this type, it's only there to make compiler happy
207///
208/// See <https://users.rust-lang.org/t/argument-requires-that-is-borrowed-for-static/66503/2?u=yandros> for details.
209pub type PhantomBound<'big, 'small> = PhantomData<&'small &'big ()>;
210
211/// Error returned when the autocommit function fails
212#[derive(Debug, Error)]
213pub enum AutocommitError<E> {
214    /// Committing the transaction failed too many times, giving up
215    #[error("Commit Failed: {last_error}")]
216    CommitFailed {
217        /// Number of attempts
218        attempts: usize,
219        /// Last error on commit
220        last_error: DatabaseError,
221    },
222    /// Error returned by the closure provided to `autocommit`. If returned no
223    /// commit was attempted in that round
224    #[error("Closure error: {error}")]
225    ClosureError {
226        /// The attempt on which the closure returned an error
227        ///
228        /// Values other than 0 typically indicate a logic error since the
229        /// closure given to `autocommit` should not have side effects
230        /// and thus keep succeeding if it succeeded once.
231        attempts: usize,
232        /// Error returned by the closure
233        error: E,
234    },
235}
236
237pub trait AutocommitResultExt<T, E> {
238    /// Unwraps the "commit failed" error variant. Use this in cases where
239    /// autocommit is instructed to run indefinitely and commit will thus never
240    /// fail.
241    fn unwrap_autocommit(self) -> std::result::Result<T, E>;
242}
243
244impl<T, E> AutocommitResultExt<T, E> for std::result::Result<T, AutocommitError<E>> {
245    fn unwrap_autocommit(self) -> std::result::Result<T, E> {
246        match self {
247            Ok(value) => Ok(value),
248            Err(AutocommitError::CommitFailed { .. }) => {
249                panic!("`unwrap_autocommit` called on a autocommit result with finite retries");
250            }
251            Err(AutocommitError::ClosureError { error, .. }) => Err(error),
252        }
253    }
254}
255
256/// Raw database implementation
257///
258/// This and [`IRawDatabaseTransaction`] are meant to be implemented
259/// by crates like `fedimint-rocksdb` to provide a concrete implementation
260/// of a database to be used by Fedimint.
261///
262/// This is in contrast of [`IDatabase`] which includes extra
263/// functionality that Fedimint needs (and adds) on top of it.
264#[apply(async_trait_maybe_send!)]
265pub trait IRawDatabase: Debug + MaybeSend + MaybeSync + 'static {
266    /// A raw database transaction type
267    type Transaction<'a>: IRawDatabaseTransaction + Debug;
268
269    /// Start a database transaction
270    async fn begin_transaction<'a>(&'a self) -> Self::Transaction<'a>;
271
272    // Checkpoint the database to a backup directory
273    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()>;
274}
275
276#[apply(async_trait_maybe_send!)]
277impl<T> IRawDatabase for Box<T>
278where
279    T: IRawDatabase,
280{
281    type Transaction<'a> = <T as IRawDatabase>::Transaction<'a>;
282
283    async fn begin_transaction<'a>(&'a self) -> Self::Transaction<'a> {
284        (**self).begin_transaction().await
285    }
286
287    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
288        (**self).checkpoint(backup_path)
289    }
290}
291
292/// An extension trait with convenience operations on [`IRawDatabase`]
293pub trait IRawDatabaseExt: IRawDatabase + Sized {
294    /// Convert to type implementing [`IRawDatabase`] into [`Database`].
295    ///
296    /// When type inference is not an issue, [`Into::into`] can be used instead.
297    fn into_database(self) -> Database {
298        Database::new(self, ModuleRegistry::default())
299    }
300}
301
302impl<T> IRawDatabaseExt for T where T: IRawDatabase {}
303
304impl<T> From<T> for Database
305where
306    T: IRawDatabase,
307{
308    fn from(raw: T) -> Self {
309        Self::new(raw, ModuleRegistry::default())
310    }
311}
312
313/// A database that on top of a raw database operation, implements
314/// key notification system.
315#[apply(async_trait_maybe_send!)]
316pub trait IDatabase: Debug + MaybeSend + MaybeSync + 'static {
317    /// Start a database transaction
318    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a>;
319    /// Register for `key` updates, returning a future that waits for an update.
320    ///
321    /// Await registration before taking the database snapshot used to check a
322    /// condition. Once registration completes, updates must be remembered even
323    /// before the returned future is polled. The future may wake spuriously;
324    /// register again before checking a fresh snapshot after each wake.
325    ///
326    /// The returned future borrows the database, not the key. Dropping it
327    /// cancels the wait.
328    async fn register<'a>(&'a self, key: &[u8]) -> BoxFuture<'a, ()>;
329    /// Notify about `key` update (creation, modification, deletion)
330    async fn notify(&self, key: &[u8]);
331
332    /// The prefix len of this database refers to the global (as opposed to
333    /// module-isolated) key space
334    fn is_global(&self) -> bool;
335
336    /// Checkpoints the database to a backup directory
337    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()>;
338}
339
340#[apply(async_trait_maybe_send!)]
341impl<T> IDatabase for Arc<T>
342where
343    T: IDatabase + ?Sized,
344{
345    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a> {
346        (**self).begin_transaction().await
347    }
348    async fn register<'a>(&'a self, key: &[u8]) -> BoxFuture<'a, ()> {
349        (**self).register(key).await
350    }
351    async fn notify(&self, key: &[u8]) {
352        (**self).notify(key).await;
353    }
354
355    fn is_global(&self) -> bool {
356        (**self).is_global()
357    }
358
359    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
360        (**self).checkpoint(backup_path)
361    }
362}
363
364/// Base functionality around [`IRawDatabase`] to make it a [`IDatabase`]
365///
366/// Mostly notification system, but also run-time single-commit handling.
367struct BaseDatabase<RawDatabase> {
368    notifications: Arc<Notifications>,
369    raw: RawDatabase,
370}
371
372impl<RawDatabase> fmt::Debug for BaseDatabase<RawDatabase> {
373    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
374        f.write_str("BaseDatabase")
375    }
376}
377
378#[apply(async_trait_maybe_send!)]
379impl<RawDatabase: IRawDatabase + MaybeSend + 'static> IDatabase for BaseDatabase<RawDatabase> {
380    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a> {
381        Box::new(BaseDatabaseTransaction::new(
382            self.raw.begin_transaction().await,
383            self.notifications.clone(),
384        ))
385    }
386    async fn register<'a>(&'a self, key: &[u8]) -> BoxFuture<'a, ()> {
387        Box::pin(self.notifications.register(key))
388    }
389    async fn notify(&self, key: &[u8]) {
390        self.notifications.notify(key);
391    }
392
393    fn is_global(&self) -> bool {
394        true
395    }
396
397    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
398        self.raw.checkpoint(backup_path)
399    }
400}
401
402/// A public-facing newtype over `IDatabase`
403///
404/// Notably carries set of module decoders (`ModuleDecoderRegistry`)
405/// and implements common utility function for auto-commits, db isolation,
406/// and other.
407#[derive(Clone, Debug)]
408pub struct Database {
409    inner: Arc<dyn IDatabase + 'static>,
410    module_decoders: ModuleDecoderRegistry,
411}
412
413impl Database {
414    pub fn strong_count(&self) -> usize {
415        Arc::strong_count(&self.inner)
416    }
417
418    pub fn into_inner(self) -> Arc<dyn IDatabase + 'static> {
419        self.inner
420    }
421}
422
423impl Database {
424    /// Creates a new Fedimint database from any object implementing
425    /// [`IDatabase`].
426    ///
427    /// See also [`Database::new_from_arc`].
428    pub fn new(raw: impl IRawDatabase + 'static, module_decoders: ModuleDecoderRegistry) -> Self {
429        let inner = BaseDatabase {
430            raw,
431            notifications: Arc::new(Notifications::new()),
432        };
433        Self::new_from_arc(
434            Arc::new(inner) as Arc<dyn IDatabase + 'static>,
435            module_decoders,
436        )
437    }
438
439    /// Create [`Database`] from an already typed-erased `IDatabase`.
440    pub fn new_from_arc(
441        inner: Arc<dyn IDatabase + 'static>,
442        module_decoders: ModuleDecoderRegistry,
443    ) -> Self {
444        Self {
445            inner,
446            module_decoders,
447        }
448    }
449
450    /// Create [`Database`] isolated to a partition with a given `prefix`
451    pub fn with_prefix(&self, prefix: Vec<u8>) -> Self {
452        Self {
453            inner: Arc::new(PrefixDatabase {
454                inner: self.inner.clone(),
455                global_dbtx_access_token: None,
456                prefix,
457            }),
458            module_decoders: self.module_decoders.clone(),
459        }
460    }
461
462    /// Create [`Database`] isolated to a partition with a prefix for a given
463    /// `module_instance_id`, allowing the module to access `global_dbtx` with
464    /// the right `access_token`
465    pub fn with_prefix_module_id(
466        &self,
467        module_instance_id: ModuleInstanceId,
468    ) -> (Self, GlobalDBTxAccessToken) {
469        let prefix = module_instance_id_to_byte_prefix(module_instance_id);
470        let global_dbtx_access_token = GlobalDBTxAccessToken::from_prefix(&prefix);
471        (
472            Self {
473                inner: Arc::new(PrefixDatabase {
474                    inner: self.inner.clone(),
475                    global_dbtx_access_token: Some(global_dbtx_access_token),
476                    prefix,
477                }),
478                module_decoders: self.module_decoders.clone(),
479            },
480            global_dbtx_access_token,
481        )
482    }
483
484    pub fn with_decoders(&self, module_decoders: ModuleDecoderRegistry) -> Self {
485        Self {
486            inner: self.inner.clone(),
487            module_decoders,
488        }
489    }
490
491    /// Is this `Database` a global, unpartitioned `Database`
492    pub fn is_global(&self) -> bool {
493        self.inner.is_global()
494    }
495
496    /// `Err` if [`Self::is_global`] is not true
497    pub fn ensure_global(&self) -> DatabaseResult<()> {
498        if !self.is_global() {
499            return Err(DatabaseError::NotGlobal);
500        }
501
502        Ok(())
503    }
504
505    /// `Err` if [`Self::is_global`] is true
506    pub fn ensure_isolated(&self) -> DatabaseResult<()> {
507        if self.is_global() {
508            return Err(DatabaseError::NotIsolated);
509        }
510
511        Ok(())
512    }
513
514    /// Begin a new committable database transaction
515    pub async fn begin_transaction<'s, 'tx>(&'s self) -> DatabaseTransaction<'tx, Committable>
516    where
517        's: 'tx,
518    {
519        DatabaseTransaction::<Committable>::new(
520            self.inner.begin_transaction().await,
521            self.module_decoders.clone(),
522        )
523    }
524
525    /// Begin a new non-committable database transaction
526    pub async fn begin_transaction_nc<'s, 'tx>(&'s self) -> DatabaseTransaction<'tx, NonCommittable>
527    where
528        's: 'tx,
529    {
530        self.begin_transaction().await.into_nc()
531    }
532
533    pub fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
534        self.inner.checkpoint(backup_path)
535    }
536
537    /// Runs a closure with a reference to a database transaction and tries to
538    /// commit the transaction if the closure returns `Ok` and rolls it back
539    /// otherwise. If committing fails the closure is run for up to
540    /// `max_attempts` times. If `max_attempts` is `None` it will run
541    /// `usize::MAX` times which is close enough to infinite times.
542    ///
543    /// The closure `tx_fn` provided should not have side effects outside of the
544    /// database transaction provided, or if it does these should be
545    /// idempotent, since the closure might be run multiple times.
546    ///
547    /// # Lifetime Parameters
548    ///
549    /// The higher rank trait bound (HRTB) `'a` that is applied to the the
550    /// mutable reference to the database transaction ensures that the
551    /// reference lives as least as long as the returned future of the
552    /// closure.
553    ///
554    /// Further, the reference to self (`'s`) must outlive the
555    /// `DatabaseTransaction<'dt>`. In other words, the
556    /// `DatabaseTransaction` must live as least as long as `self` and that is
557    /// true as the `DatabaseTransaction` is only dropped at the end of the
558    /// `loop{}`.
559    ///
560    /// # Panics
561    ///
562    /// This function panics when the given number of maximum attempts is zero.
563    /// `max_attempts` must be greater or equal to one.
564    pub async fn autocommit<'s, 'dbtx, F, T, E>(
565        &'s self,
566        tx_fn: F,
567        max_attempts: Option<usize>,
568    ) -> std::result::Result<T, AutocommitError<E>>
569    where
570        's: 'dbtx,
571        for<'r, 'o> F: Fn(
572            &'r mut DatabaseTransaction<'o>,
573            PhantomBound<'dbtx, 'o>,
574        ) -> BoxFuture<'r, std::result::Result<T, E>>,
575    {
576        assert_ne!(max_attempts, Some(0));
577        let mut curr_attempts: usize = 0;
578
579        loop {
580            // The `checked_add()` function is used to catch the `usize` overflow.
581            // With `usize=32bit` and an assumed time of 1ms per iteration, this would crash
582            // after ~50 days. But if that's the case, something else must be wrong.
583            // With `usize=64bit` it would take much longer, obviously.
584            curr_attempts = curr_attempts
585                .checked_add(1)
586                .expect("db autocommit attempt counter overflowed");
587
588            let mut dbtx = self.begin_transaction().await;
589
590            let tx_fn_res = tx_fn(&mut dbtx.to_ref_nc(), PhantomData).await;
591            let val = match tx_fn_res {
592                Ok(val) => val,
593                Err(err) => {
594                    dbtx.ignore_uncommitted();
595                    return Err(AutocommitError::ClosureError {
596                        attempts: curr_attempts,
597                        error: err,
598                    });
599                }
600            };
601
602            let _timing /* logs on drop */ = timing::TimeReporter::new("autocommit - commit_tx");
603
604            match dbtx.commit_tx_result().await {
605                Ok(()) => {
606                    return Ok(val);
607                }
608                Err(err) => {
609                    if max_attempts.is_some_and(|max_att| max_att <= curr_attempts) {
610                        warn!(
611                            target: LOG_DB,
612                            curr_attempts,
613                            err = %err.fmt_compact(),
614                            "Database commit failed in an autocommit block - terminating"
615                        );
616                        return Err(AutocommitError::CommitFailed {
617                            attempts: curr_attempts,
618                            last_error: err,
619                        });
620                    }
621
622                    let delay = (2u64.pow(curr_attempts.min(7) as u32) * 10).min(1000);
623                    let delay = rand::thread_rng().gen_range(delay..(2 * delay));
624                    warn!(
625                        target: LOG_DB,
626                        curr_attempts,
627                        err = %err.fmt_compact(),
628                        delay_ms = %delay,
629                        "Database commit failed in an autocommit block - retrying"
630                    );
631                    crate::runtime::sleep(Duration::from_millis(delay)).await;
632                }
633            }
634        }
635    }
636
637    /// Waits for key to be notified.
638    ///
639    /// Calls the `checker` when value of the key may have changed.
640    /// Returns the value when `checker` returns a `Some(T)`.
641    pub async fn wait_key_check<'a, K, T>(
642        &'a self,
643        key: &K,
644        checker: impl Fn(Option<K::Value>) -> Option<T>,
645    ) -> (T, DatabaseTransaction<'a, Committable>)
646    where
647        K: DatabaseKey + DatabaseRecord + DatabaseKeyWithNotify,
648    {
649        let key_bytes = key.to_bytes();
650        loop {
651            // register for notification
652            let notify = self.inner.register(&key_bytes).await;
653
654            // check for value in db
655            let mut tx = self.inner.begin_transaction().await;
656
657            let maybe_value_bytes = tx
658                .raw_get_bytes(&key_bytes)
659                .await
660                .expect("Unrecoverable error when reading from database")
661                .map(|value_bytes| {
662                    decode_value_expect(&value_bytes, &self.module_decoders, &key_bytes)
663                });
664
665            if let Some(value) = checker(maybe_value_bytes) {
666                return (
667                    value,
668                    DatabaseTransaction::new(tx, self.module_decoders.clone()),
669                );
670            }
671
672            // key not found, try again
673            notify.await;
674            // if miss a notification between await and next register, it is
675            // fine. because we are going check the database
676        }
677    }
678
679    /// Waits for key to be present in database.
680    pub async fn wait_key_exists<K>(&self, key: &K) -> K::Value
681    where
682        K: DatabaseKey + DatabaseRecord + DatabaseKeyWithNotify,
683    {
684        self.wait_key_check(key, std::convert::identity).await.0
685    }
686}
687
688fn module_instance_id_to_byte_prefix(module_instance_id: u16) -> Vec<u8> {
689    let mut bytes = vec![MODULE_GLOBAL_PREFIX];
690    bytes.append(&mut module_instance_id.consensus_encode_to_vec());
691    bytes
692}
693
694/// A database that wraps an `inner` one and adds a prefix to all operations,
695/// effectively creating an isolated partition.
696#[derive(Clone, Debug)]
697struct PrefixDatabase<Inner>
698where
699    Inner: Debug,
700{
701    prefix: Vec<u8>,
702    global_dbtx_access_token: Option<GlobalDBTxAccessToken>,
703    inner: Inner,
704}
705
706impl<Inner> PrefixDatabase<Inner>
707where
708    Inner: Debug,
709{
710    // TODO: we should optimize these concatenations, maybe by having an internal
711    // `key: &[&[u8]]` that we flatten once, when passing to lowest layer, or
712    // something
713    fn get_full_key(&self, key: &[u8]) -> Vec<u8> {
714        let mut full_key = self.prefix.clone();
715        full_key.extend_from_slice(key);
716        full_key
717    }
718}
719
720#[apply(async_trait_maybe_send!)]
721impl<Inner> IDatabase for PrefixDatabase<Inner>
722where
723    Inner: Debug + MaybeSend + MaybeSync + 'static + IDatabase,
724{
725    async fn begin_transaction<'a>(&'a self) -> Box<dyn IDatabaseTransaction + 'a> {
726        Box::new(PrefixDatabaseTransaction {
727            inner: self.inner.begin_transaction().await,
728            global_dbtx_access_token: self.global_dbtx_access_token,
729            prefix: self.prefix.clone(),
730        })
731    }
732    async fn register<'a>(&'a self, key: &[u8]) -> BoxFuture<'a, ()> {
733        self.inner.register(&self.get_full_key(key)).await
734    }
735
736    async fn notify(&self, key: &[u8]) {
737        self.inner.notify(&self.get_full_key(key)).await;
738    }
739
740    fn is_global(&self) -> bool {
741        if self.global_dbtx_access_token.is_some() {
742            false
743        } else {
744            self.inner.is_global()
745        }
746    }
747
748    fn checkpoint(&self, backup_path: &Path) -> DatabaseResult<()> {
749        self.inner.checkpoint(backup_path)
750    }
751}
752
753/// A database transactions that wraps an `inner` one and adds a prefix to all
754/// operations, effectively creating an isolated partition.
755///
756/// Produced by [`PrefixDatabase`].
757#[derive(Debug)]
758struct PrefixDatabaseTransaction<Inner> {
759    inner: Inner,
760    global_dbtx_access_token: Option<GlobalDBTxAccessToken>,
761    prefix: Vec<u8>,
762}
763
764impl<Inner> PrefixDatabaseTransaction<Inner> {
765    // TODO: we should optimize these concatenations, maybe by having an internal
766    // `key: &[&[u8]]` that we flatten once, when passing to lowest layer, or
767    // something
768    fn get_full_key(&self, key: &[u8]) -> Vec<u8> {
769        let mut full_key = self.prefix.clone();
770        full_key.extend_from_slice(key);
771        full_key
772    }
773
774    fn get_full_range(&self, range: Range<&[u8]>) -> Range<Vec<u8>> {
775        Range {
776            start: self.get_full_key(range.start),
777            end: self.get_full_key(range.end),
778        }
779    }
780
781    fn adapt_prefix_stream(stream: PrefixStream<'_>, prefix_len: usize) -> PrefixStream<'_> {
782        Box::pin(stream.map(move |(mut k, v)| {
783            k.drain(..prefix_len);
784            (k, v)
785        }))
786    }
787}
788
789#[apply(async_trait_maybe_send!)]
790impl<Inner> IDatabaseTransaction for PrefixDatabaseTransaction<Inner>
791where
792    Inner: IDatabaseTransaction,
793{
794    async fn commit_tx(&mut self) -> DatabaseResult<()> {
795        self.inner.commit_tx().await
796    }
797
798    fn is_global(&self) -> bool {
799        if self.global_dbtx_access_token.is_some() {
800            false
801        } else {
802            self.inner.is_global()
803        }
804    }
805
806    fn global_dbtx(
807        &mut self,
808        access_token: GlobalDBTxAccessToken,
809    ) -> &mut dyn IDatabaseTransaction {
810        if let Some(self_global_dbtx_access_token) = self.global_dbtx_access_token {
811            assert_eq!(
812                access_token, self_global_dbtx_access_token,
813                "Invalid access key used to access global_dbtx"
814            );
815            &mut self.inner
816        } else {
817            self.inner.global_dbtx(access_token)
818        }
819    }
820}
821
822#[apply(async_trait_maybe_send!)]
823impl<Inner> IDatabaseTransactionOpsCore for PrefixDatabaseTransaction<Inner>
824where
825    Inner: IDatabaseTransactionOpsCore,
826{
827    async fn raw_insert_bytes_no_previous(
828        &mut self,
829        key: &[u8],
830        value: &[u8],
831    ) -> DatabaseResult<()> {
832        let key = self.get_full_key(key);
833        self.inner.raw_insert_bytes_no_previous(&key, value).await
834    }
835
836    async fn raw_remove_entry_no_previous(&mut self, key: &[u8]) -> DatabaseResult<()> {
837        let key = self.get_full_key(key);
838        self.inner.raw_remove_entry_no_previous(&key).await
839    }
840
841    async fn raw_insert_bytes(
842        &mut self,
843        key: &[u8],
844        value: &[u8],
845    ) -> DatabaseResult<Option<Vec<u8>>> {
846        let key = self.get_full_key(key);
847        self.inner.raw_insert_bytes(&key, value).await
848    }
849
850    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
851        let key = self.get_full_key(key);
852        self.inner.raw_get_bytes(&key).await
853    }
854
855    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
856        let key = self.get_full_key(key);
857        self.inner.raw_remove_entry(&key).await
858    }
859
860    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
861        let key = self.get_full_key(key_prefix);
862        let stream = self.inner.raw_find_by_prefix(&key).await?;
863        Ok(Self::adapt_prefix_stream(stream, self.prefix.len()))
864    }
865
866    async fn raw_find_by_prefix_sorted_descending(
867        &mut self,
868        key_prefix: &[u8],
869    ) -> DatabaseResult<PrefixStream<'_>> {
870        let key = self.get_full_key(key_prefix);
871        let stream = self
872            .inner
873            .raw_find_by_prefix_sorted_descending(&key)
874            .await?;
875        Ok(Self::adapt_prefix_stream(stream, self.prefix.len()))
876    }
877
878    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
879        let range = self.get_full_range(range);
880        let stream = self
881            .inner
882            .raw_find_by_range(Range {
883                start: &range.start,
884                end: &range.end,
885            })
886            .await?;
887        Ok(Self::adapt_prefix_stream(stream, self.prefix.len()))
888    }
889
890    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
891        let key = self.get_full_key(key_prefix);
892        self.inner.raw_remove_by_prefix(&key).await
893    }
894}
895
896impl<Inner> IDatabaseTransactionOps for PrefixDatabaseTransaction<Inner> where
897    Inner: IDatabaseTransactionOps
898{
899}
900
901/// Core raw a operations database transactions supports
902///
903/// Used to enforce the same signature on all types supporting it
904#[apply(async_trait_maybe_send!)]
905pub trait IDatabaseTransactionOpsCore: MaybeSend {
906    /// Insert entry
907    async fn raw_insert_bytes(
908        &mut self,
909        key: &[u8],
910        value: &[u8],
911    ) -> DatabaseResult<Option<Vec<u8>>>;
912
913    /// Get key value
914    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>>;
915
916    /// Insert without returning the previous bytes. Backends may avoid reading
917    /// them, but must retain the same write-conflict tracking.
918    async fn raw_insert_bytes_no_previous(
919        &mut self,
920        key: &[u8],
921        value: &[u8],
922    ) -> DatabaseResult<()> {
923        self.raw_insert_bytes(key, value).await.map(|_| ())
924    }
925
926    /// Remove entry by `key`
927    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>>;
928
929    /// Remove without returning the previous bytes. Backends may avoid reading
930    /// them, but must retain the same write-conflict tracking.
931    async fn raw_remove_entry_no_previous(&mut self, key: &[u8]) -> DatabaseResult<()> {
932        self.raw_remove_entry(key).await.map(|_| ())
933    }
934
935    /// Returns an stream of key-value pairs with keys that start with
936    /// `key_prefix`, sorted by key.
937    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>>;
938
939    /// Same as [`Self::raw_find_by_prefix`] but the order is descending by key.
940    async fn raw_find_by_prefix_sorted_descending(
941        &mut self,
942        key_prefix: &[u8],
943    ) -> DatabaseResult<PrefixStream<'_>>;
944
945    /// Returns an stream of key-value pairs with keys within a `range`, sorted
946    /// by key. [`Range`] is an (half-open) range bounded inclusively below and
947    /// exclusively above.
948    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>>;
949
950    /// Delete keys matching prefix
951    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()>;
952}
953
954#[apply(async_trait_maybe_send!)]
955impl<T> IDatabaseTransactionOpsCore for Box<T>
956where
957    T: IDatabaseTransactionOpsCore + ?Sized,
958{
959    async fn raw_insert_bytes_no_previous(
960        &mut self,
961        key: &[u8],
962        value: &[u8],
963    ) -> DatabaseResult<()> {
964        (**self).raw_insert_bytes_no_previous(key, value).await
965    }
966
967    async fn raw_remove_entry_no_previous(&mut self, key: &[u8]) -> DatabaseResult<()> {
968        (**self).raw_remove_entry_no_previous(key).await
969    }
970
971    async fn raw_insert_bytes(
972        &mut self,
973        key: &[u8],
974        value: &[u8],
975    ) -> DatabaseResult<Option<Vec<u8>>> {
976        (**self).raw_insert_bytes(key, value).await
977    }
978
979    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
980        (**self).raw_get_bytes(key).await
981    }
982
983    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
984        (**self).raw_remove_entry(key).await
985    }
986
987    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
988        (**self).raw_find_by_prefix(key_prefix).await
989    }
990
991    async fn raw_find_by_prefix_sorted_descending(
992        &mut self,
993        key_prefix: &[u8],
994    ) -> DatabaseResult<PrefixStream<'_>> {
995        (**self)
996            .raw_find_by_prefix_sorted_descending(key_prefix)
997            .await
998    }
999
1000    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
1001        (**self).raw_find_by_range(range).await
1002    }
1003
1004    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
1005        (**self).raw_remove_by_prefix(key_prefix).await
1006    }
1007}
1008
1009#[apply(async_trait_maybe_send!)]
1010impl<T> IDatabaseTransactionOpsCore for &mut T
1011where
1012    T: IDatabaseTransactionOpsCore + ?Sized,
1013{
1014    async fn raw_insert_bytes_no_previous(
1015        &mut self,
1016        key: &[u8],
1017        value: &[u8],
1018    ) -> DatabaseResult<()> {
1019        (**self).raw_insert_bytes_no_previous(key, value).await
1020    }
1021
1022    async fn raw_remove_entry_no_previous(&mut self, key: &[u8]) -> DatabaseResult<()> {
1023        (**self).raw_remove_entry_no_previous(key).await
1024    }
1025
1026    async fn raw_insert_bytes(
1027        &mut self,
1028        key: &[u8],
1029        value: &[u8],
1030    ) -> DatabaseResult<Option<Vec<u8>>> {
1031        (**self).raw_insert_bytes(key, value).await
1032    }
1033
1034    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1035        (**self).raw_get_bytes(key).await
1036    }
1037
1038    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1039        (**self).raw_remove_entry(key).await
1040    }
1041
1042    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
1043        (**self).raw_find_by_prefix(key_prefix).await
1044    }
1045
1046    async fn raw_find_by_prefix_sorted_descending(
1047        &mut self,
1048        key_prefix: &[u8],
1049    ) -> DatabaseResult<PrefixStream<'_>> {
1050        (**self)
1051            .raw_find_by_prefix_sorted_descending(key_prefix)
1052            .await
1053    }
1054
1055    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
1056        (**self).raw_find_by_range(range).await
1057    }
1058
1059    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
1060        (**self).raw_remove_by_prefix(key_prefix).await
1061    }
1062}
1063
1064/// Additional operations (only some) database transactions expose, on top of
1065/// [`IDatabaseTransactionOpsCore`]
1066///
1067/// In certain contexts exposing these operations would be a problem, so they
1068/// are moved to a separate trait.
1069pub trait IDatabaseTransactionOps: IDatabaseTransactionOpsCore + MaybeSend {}
1070
1071impl<T> IDatabaseTransactionOps for Box<T> where T: IDatabaseTransactionOps + ?Sized {}
1072
1073impl<T> IDatabaseTransactionOps for &mut T where T: IDatabaseTransactionOps + ?Sized {}
1074
1075/// Like [`IDatabaseTransactionOpsCore`], but typed
1076///
1077/// Implemented via blanket impl for everything that implements
1078/// [`IDatabaseTransactionOpsCore`] that has decoders (implements
1079/// [`WithDecoders`]).
1080#[apply(async_trait_maybe_send!)]
1081pub trait IDatabaseTransactionOpsCoreTyped<'a> {
1082    async fn get_value<K>(&mut self, key: &K) -> Option<K::Value>
1083    where
1084        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync;
1085
1086    async fn insert_entry<K>(&mut self, key: &K, value: &K::Value) -> Option<K::Value>
1087    where
1088        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1089        K::Value: MaybeSend + MaybeSync;
1090
1091    /// Overwrite without returning or decoding the previous value.
1092    ///
1093    /// Unlike [`Self::insert_entry`], this does not validate existing bytes.
1094    /// Do not use it when the previous value or its validity is a guard.
1095    async fn insert_entry_no_previous<K>(&mut self, key: &K, value: &K::Value)
1096    where
1097        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1098        K::Value: MaybeSend + MaybeSync;
1099
1100    async fn insert_new_entry<K>(&mut self, key: &K, value: &K::Value)
1101    where
1102        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1103        K::Value: MaybeSend + MaybeSync;
1104
1105    async fn find_by_range<K>(
1106        &mut self,
1107        key_range: Range<K>,
1108    ) -> Pin<Box<maybe_add_send!(dyn Stream<Item = (K, K::Value)> + '_)>>
1109    where
1110        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1111        K::Value: MaybeSend + MaybeSync;
1112
1113    async fn find_by_prefix<KP>(
1114        &mut self,
1115        key_prefix: &KP,
1116    ) -> Pin<
1117        Box<
1118            maybe_add_send!(
1119                dyn Stream<
1120                        Item = (
1121                            KP::Record,
1122                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1123                        ),
1124                    > + '_
1125            ),
1126        >,
1127    >
1128    where
1129        KP: DatabaseLookup + MaybeSend + MaybeSync,
1130        KP::Record: DatabaseKey;
1131
1132    async fn find_by_prefix_sorted_descending<KP>(
1133        &mut self,
1134        key_prefix: &KP,
1135    ) -> Pin<
1136        Box<
1137            maybe_add_send!(
1138                dyn Stream<
1139                        Item = (
1140                            KP::Record,
1141                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1142                        ),
1143                    > + '_
1144            ),
1145        >,
1146    >
1147    where
1148        KP: DatabaseLookup + MaybeSend + MaybeSync,
1149        KP::Record: DatabaseKey;
1150
1151    async fn remove_entry<K>(&mut self, key: &K) -> Option<K::Value>
1152    where
1153        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync;
1154
1155    /// Remove without returning or decoding the previous value.
1156    ///
1157    /// Unlike [`Self::remove_entry`], this does not validate existing bytes.
1158    /// Do not use it when existence or validity of the previous value is a
1159    /// guard.
1160    async fn remove_entry_no_previous<K>(&mut self, key: &K)
1161    where
1162        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync;
1163
1164    async fn remove_by_prefix<KP>(&mut self, key_prefix: &KP)
1165    where
1166        KP: DatabaseLookup + MaybeSend + MaybeSync;
1167}
1168
1169// blanket implementation of typed ops for anything that implements raw ops and
1170// has decoders
1171#[apply(async_trait_maybe_send!)]
1172impl<T> IDatabaseTransactionOpsCoreTyped<'_> for T
1173where
1174    T: IDatabaseTransactionOpsCore + WithDecoders,
1175{
1176    async fn get_value<K>(&mut self, key: &K) -> Option<K::Value>
1177    where
1178        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1179    {
1180        let key_bytes = key.to_bytes();
1181        let raw = self
1182            .raw_get_bytes(&key_bytes)
1183            .await
1184            .expect("Unrecoverable error occurred while reading and entry from the database");
1185        raw.map(|value_bytes| {
1186            decode_value_expect::<K::Value>(&value_bytes, self.decoders(), &key_bytes)
1187        })
1188    }
1189
1190    async fn insert_entry<K>(&mut self, key: &K, value: &K::Value) -> Option<K::Value>
1191    where
1192        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1193        K::Value: MaybeSend + MaybeSync,
1194    {
1195        let key_bytes = key.to_bytes();
1196        self.raw_insert_bytes(&key_bytes, &value.to_bytes())
1197            .await
1198            .expect("Unrecoverable error occurred while inserting entry into the database")
1199            .map(|value_bytes| {
1200                decode_value_expect::<K::Value>(&value_bytes, self.decoders(), &key_bytes)
1201            })
1202    }
1203
1204    async fn insert_entry_no_previous<K>(&mut self, key: &K, value: &K::Value)
1205    where
1206        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1207        K::Value: MaybeSend + MaybeSync,
1208    {
1209        self.raw_insert_bytes_no_previous(&key.to_bytes(), &value.to_bytes())
1210            .await
1211            .expect("Unrecoverable error occurred while inserting entry into the database");
1212    }
1213
1214    async fn insert_new_entry<K>(&mut self, key: &K, value: &K::Value)
1215    where
1216        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1217        K::Value: MaybeSend + MaybeSync,
1218    {
1219        if let Some(prev) = self.insert_entry(key, value).await {
1220            panic!(
1221                "Database overwriting element when expecting insertion of new entry. Key: {key:?} Prev Value: {prev:?}"
1222            );
1223        }
1224    }
1225
1226    async fn find_by_range<K>(
1227        &mut self,
1228        key_range: Range<K>,
1229    ) -> Pin<Box<maybe_add_send!(dyn Stream<Item = (K, K::Value)> + '_)>>
1230    where
1231        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1232        K::Value: MaybeSend + MaybeSync,
1233    {
1234        let decoders = self.decoders().clone();
1235        Box::pin(
1236            self.raw_find_by_range(Range {
1237                start: &key_range.start.to_bytes(),
1238                end: &key_range.end.to_bytes(),
1239            })
1240            .await
1241            .expect("Unrecoverable error occurred while listing entries from the database")
1242            .map(move |(key_bytes, value_bytes)| {
1243                let key = decode_key_expect(&key_bytes, &decoders);
1244                let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
1245                (key, value)
1246            }),
1247        )
1248    }
1249
1250    async fn find_by_prefix<KP>(
1251        &mut self,
1252        key_prefix: &KP,
1253    ) -> Pin<
1254        Box<
1255            maybe_add_send!(
1256                dyn Stream<
1257                        Item = (
1258                            KP::Record,
1259                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1260                        ),
1261                    > + '_
1262            ),
1263        >,
1264    >
1265    where
1266        KP: DatabaseLookup + MaybeSend + MaybeSync,
1267        KP::Record: DatabaseKey,
1268    {
1269        let decoders = self.decoders().clone();
1270        Box::pin(
1271            self.raw_find_by_prefix(&key_prefix.to_bytes())
1272                .await
1273                .expect("Unrecoverable error occurred while listing entries from the database")
1274                .map(move |(key_bytes, value_bytes)| {
1275                    let key = decode_key_expect(&key_bytes, &decoders);
1276                    let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
1277                    (key, value)
1278                }),
1279        )
1280    }
1281
1282    async fn find_by_prefix_sorted_descending<KP>(
1283        &mut self,
1284        key_prefix: &KP,
1285    ) -> Pin<
1286        Box<
1287            maybe_add_send!(
1288                dyn Stream<
1289                        Item = (
1290                            KP::Record,
1291                            <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
1292                        ),
1293                    > + '_
1294            ),
1295        >,
1296    >
1297    where
1298        KP: DatabaseLookup + MaybeSend + MaybeSync,
1299        KP::Record: DatabaseKey,
1300    {
1301        let decoders = self.decoders().clone();
1302        Box::pin(
1303            self.raw_find_by_prefix_sorted_descending(&key_prefix.to_bytes())
1304                .await
1305                .expect("Unrecoverable error occurred while listing entries from the database")
1306                .map(move |(key_bytes, value_bytes)| {
1307                    let key = decode_key_expect(&key_bytes, &decoders);
1308                    let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
1309                    (key, value)
1310                }),
1311        )
1312    }
1313    async fn remove_entry<K>(&mut self, key: &K) -> Option<K::Value>
1314    where
1315        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1316    {
1317        let key_bytes = key.to_bytes();
1318        self.raw_remove_entry(&key_bytes)
1319            .await
1320            .expect("Unrecoverable error occurred while inserting removing entry from the database")
1321            .map(|value_bytes| {
1322                decode_value_expect::<K::Value>(&value_bytes, self.decoders(), &key_bytes)
1323            })
1324    }
1325    async fn remove_entry_no_previous<K>(&mut self, key: &K)
1326    where
1327        K: DatabaseKey + DatabaseRecord + MaybeSend + MaybeSync,
1328    {
1329        self.raw_remove_entry_no_previous(&key.to_bytes())
1330            .await
1331            .expect("Unrecoverable error occurred while removing entry from the database");
1332    }
1333
1334    async fn remove_by_prefix<KP>(&mut self, key_prefix: &KP)
1335    where
1336        KP: DatabaseLookup + MaybeSend + MaybeSync,
1337    {
1338        self.raw_remove_by_prefix(&key_prefix.to_bytes())
1339            .await
1340            .expect("Unrecoverable error when removing entries from the database");
1341    }
1342}
1343
1344/// A database type that has decoders, which allows it to implement
1345/// [`IDatabaseTransactionOpsCoreTyped`]
1346pub trait WithDecoders {
1347    fn decoders(&self) -> &ModuleDecoderRegistry;
1348}
1349
1350/// Raw database transaction (e.g. rocksdb implementation)
1351#[apply(async_trait_maybe_send!)]
1352pub trait IRawDatabaseTransaction: MaybeSend + IDatabaseTransactionOps {
1353    async fn commit_tx(self) -> DatabaseResult<()>;
1354}
1355
1356/// Fedimint database transaction
1357///
1358/// See [`IDatabase`] for more info.
1359#[apply(async_trait_maybe_send!)]
1360pub trait IDatabaseTransaction: MaybeSend + IDatabaseTransactionOps + fmt::Debug {
1361    /// Commit the transaction
1362    async fn commit_tx(&mut self) -> DatabaseResult<()>;
1363
1364    /// Is global database
1365    fn is_global(&self) -> bool;
1366
1367    /// Get the global database tx from a module-prefixed database transaction
1368    ///
1369    /// Meant to be called only by core internals, and module developers should
1370    /// not call it directly.
1371    #[doc(hidden)]
1372    fn global_dbtx(&mut self, access_token: GlobalDBTxAccessToken)
1373    -> &mut dyn IDatabaseTransaction;
1374}
1375
1376#[apply(async_trait_maybe_send!)]
1377impl<T> IDatabaseTransaction for Box<T>
1378where
1379    T: IDatabaseTransaction + ?Sized,
1380{
1381    async fn commit_tx(&mut self) -> DatabaseResult<()> {
1382        (**self).commit_tx().await
1383    }
1384
1385    fn is_global(&self) -> bool {
1386        (**self).is_global()
1387    }
1388
1389    fn global_dbtx(
1390        &mut self,
1391        access_token: GlobalDBTxAccessToken,
1392    ) -> &mut dyn IDatabaseTransaction {
1393        (**self).global_dbtx(access_token)
1394    }
1395}
1396
1397#[apply(async_trait_maybe_send!)]
1398impl<'a, T> IDatabaseTransaction for &'a mut T
1399where
1400    T: IDatabaseTransaction + ?Sized,
1401{
1402    async fn commit_tx(&mut self) -> DatabaseResult<()> {
1403        (**self).commit_tx().await
1404    }
1405
1406    fn is_global(&self) -> bool {
1407        (**self).is_global()
1408    }
1409
1410    fn global_dbtx(&mut self, access_key: GlobalDBTxAccessToken) -> &mut dyn IDatabaseTransaction {
1411        (**self).global_dbtx(access_key)
1412    }
1413}
1414
1415/// Struct that implements `IRawDatabaseTransaction` and can be wrapped
1416/// easier in other structs since it does not consumed `self` by move.
1417struct BaseDatabaseTransaction<Tx> {
1418    // TODO: merge options
1419    raw: Option<Tx>,
1420    notify_queue: Option<NotifyQueue>,
1421    notifications: Arc<Notifications>,
1422}
1423
1424impl<Tx> fmt::Debug for BaseDatabaseTransaction<Tx>
1425where
1426    Tx: fmt::Debug,
1427{
1428    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1429        f.write_fmt(format_args!(
1430            "BaseDatabaseTransaction{{ raw={:?} }}",
1431            self.raw
1432        ))
1433    }
1434}
1435impl<Tx> BaseDatabaseTransaction<Tx>
1436where
1437    Tx: IRawDatabaseTransaction,
1438{
1439    fn new(dbtx: Tx, notifications: Arc<Notifications>) -> Self {
1440        Self {
1441            raw: Some(dbtx),
1442            notifications,
1443            notify_queue: Some(NotifyQueue::new()),
1444        }
1445    }
1446
1447    fn add_notification_key(&mut self, key: &[u8]) -> DatabaseResult<()> {
1448        self.notify_queue
1449            .as_mut()
1450            .ok_or(DatabaseError::TransactionConsumed)?
1451            .add(key);
1452        Ok(())
1453    }
1454}
1455
1456#[apply(async_trait_maybe_send!)]
1457impl<Tx: IRawDatabaseTransaction> IDatabaseTransactionOpsCore for BaseDatabaseTransaction<Tx> {
1458    async fn raw_insert_bytes_no_previous(
1459        &mut self,
1460        key: &[u8],
1461        value: &[u8],
1462    ) -> DatabaseResult<()> {
1463        self.add_notification_key(key)?;
1464        self.raw
1465            .as_mut()
1466            .ok_or(DatabaseError::TransactionConsumed)?
1467            .raw_insert_bytes_no_previous(key, value)
1468            .await
1469    }
1470
1471    async fn raw_remove_entry_no_previous(&mut self, key: &[u8]) -> DatabaseResult<()> {
1472        self.add_notification_key(key)?;
1473        self.raw
1474            .as_mut()
1475            .ok_or(DatabaseError::TransactionConsumed)?
1476            .raw_remove_entry_no_previous(key)
1477            .await
1478    }
1479
1480    async fn raw_insert_bytes(
1481        &mut self,
1482        key: &[u8],
1483        value: &[u8],
1484    ) -> DatabaseResult<Option<Vec<u8>>> {
1485        self.add_notification_key(key)?;
1486        self.raw
1487            .as_mut()
1488            .ok_or(DatabaseError::TransactionConsumed)?
1489            .raw_insert_bytes(key, value)
1490            .await
1491    }
1492
1493    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1494        self.raw
1495            .as_mut()
1496            .ok_or(DatabaseError::TransactionConsumed)?
1497            .raw_get_bytes(key)
1498            .await
1499    }
1500
1501    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
1502        self.add_notification_key(key)?;
1503        self.raw
1504            .as_mut()
1505            .ok_or(DatabaseError::TransactionConsumed)?
1506            .raw_remove_entry(key)
1507            .await
1508    }
1509
1510    async fn raw_find_by_range(
1511        &mut self,
1512        key_range: Range<&[u8]>,
1513    ) -> DatabaseResult<PrefixStream<'_>> {
1514        self.raw
1515            .as_mut()
1516            .ok_or(DatabaseError::TransactionConsumed)?
1517            .raw_find_by_range(key_range)
1518            .await
1519    }
1520
1521    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
1522        self.raw
1523            .as_mut()
1524            .ok_or(DatabaseError::TransactionConsumed)?
1525            .raw_find_by_prefix(key_prefix)
1526            .await
1527    }
1528
1529    async fn raw_find_by_prefix_sorted_descending(
1530        &mut self,
1531        key_prefix: &[u8],
1532    ) -> DatabaseResult<PrefixStream<'_>> {
1533        self.raw
1534            .as_mut()
1535            .ok_or(DatabaseError::TransactionConsumed)?
1536            .raw_find_by_prefix_sorted_descending(key_prefix)
1537            .await
1538    }
1539
1540    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
1541        self.raw
1542            .as_mut()
1543            .ok_or(DatabaseError::TransactionConsumed)?
1544            .raw_remove_by_prefix(key_prefix)
1545            .await
1546    }
1547}
1548
1549impl<Tx: IRawDatabaseTransaction> IDatabaseTransactionOps for BaseDatabaseTransaction<Tx> {}
1550
1551#[apply(async_trait_maybe_send!)]
1552impl<Tx: IRawDatabaseTransaction + fmt::Debug> IDatabaseTransaction
1553    for BaseDatabaseTransaction<Tx>
1554{
1555    async fn commit_tx(&mut self) -> DatabaseResult<()> {
1556        self.raw
1557            .take()
1558            .ok_or(DatabaseError::TransactionConsumed)?
1559            .commit_tx()
1560            .await?;
1561        self.notifications.submit_queue(
1562            &self
1563                .notify_queue
1564                .take()
1565                .expect("commit must be called only once"),
1566        );
1567        Ok(())
1568    }
1569
1570    fn is_global(&self) -> bool {
1571        true
1572    }
1573
1574    fn global_dbtx(
1575        &mut self,
1576        _access_token: GlobalDBTxAccessToken,
1577    ) -> &mut dyn IDatabaseTransaction {
1578        panic!("Illegal to call global_dbtx on BaseDatabaseTransaction");
1579    }
1580}
1581
1582/// A helper for tracking and logging on `Drop` any instances of uncommitted
1583/// writes
1584#[derive(Clone)]
1585struct CommitTracker {
1586    /// Is the dbtx committed
1587    is_committed: bool,
1588    /// Does the dbtx have any writes
1589    has_writes: bool,
1590    /// Don't warn-log uncommitted writes
1591    ignore_uncommitted: bool,
1592}
1593
1594impl Drop for CommitTracker {
1595    fn drop(&mut self) {
1596        if self.has_writes && !self.is_committed {
1597            if self.ignore_uncommitted {
1598                trace!(
1599                    target: LOG_DB,
1600                    "DatabaseTransaction has writes and has not called commit, but that's expected."
1601                );
1602            } else {
1603                warn!(
1604                    target: LOG_DB,
1605                    location = ?backtrace::Backtrace::new(),
1606                    "DatabaseTransaction has writes and has not called commit."
1607                );
1608            }
1609        }
1610    }
1611}
1612
1613enum MaybeRef<'a, T> {
1614    Owned(T),
1615    Borrowed(&'a mut T),
1616}
1617
1618impl<T> ops::Deref for MaybeRef<'_, T> {
1619    type Target = T;
1620
1621    fn deref(&self) -> &Self::Target {
1622        match self {
1623            MaybeRef::Owned(o) => o,
1624            MaybeRef::Borrowed(r) => r,
1625        }
1626    }
1627}
1628
1629impl<T> ops::DerefMut for MaybeRef<'_, T> {
1630    fn deref_mut(&mut self) -> &mut Self::Target {
1631        match self {
1632            MaybeRef::Owned(o) => o,
1633            MaybeRef::Borrowed(r) => r,
1634        }
1635    }
1636}
1637
1638/// Session type for [`DatabaseTransaction`] that is allowed to commit
1639///
1640/// Opposite of [`NonCommittable`].
1641pub struct Committable;
1642
1643/// Session type for a [`DatabaseTransaction`] that is not allowed to commit
1644///
1645/// Opposite of [`Committable`].
1646pub struct NonCommittable;
1647
1648/// A high level database transaction handle
1649///
1650/// `Cap` is a session type
1651pub struct DatabaseTransaction<'tx, Cap = NonCommittable> {
1652    tx: Box<dyn IDatabaseTransaction + 'tx>,
1653    // Owning transactions keep their decoder context; references borrow it for
1654    // the same lifetime as the underlying transaction instead of cloning the
1655    // registry's outer map. Decoder metadata is immutable through this handle.
1656    decoders: Cow<'tx, ModuleDecoderRegistry>,
1657    commit_tracker: MaybeRef<'tx, CommitTracker>,
1658    on_commit_hooks: MaybeRef<'tx, Vec<Box<maybe_add_send!(dyn FnOnce())>>>,
1659    capability: marker::PhantomData<Cap>,
1660}
1661
1662impl<Cap> fmt::Debug for DatabaseTransaction<'_, Cap> {
1663    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1664        f.write_fmt(format_args!(
1665            "DatabaseTransaction {{ tx: {:?}, decoders={:?} }}",
1666            self.tx, self.decoders
1667        ))
1668    }
1669}
1670
1671impl<Cap> WithDecoders for DatabaseTransaction<'_, Cap> {
1672    fn decoders(&self) -> &ModuleDecoderRegistry {
1673        &self.decoders
1674    }
1675}
1676
1677#[instrument(target = LOG_DB, level = "trace", skip_all, fields(value_type = std::any::type_name::<V>()), err)]
1678fn decode_value<V: DatabaseValue>(
1679    value_bytes: &[u8],
1680    decoders: &ModuleDecoderRegistry,
1681) -> std::result::Result<V, DecodingError> {
1682    trace!(
1683        bytes = %AbbreviateHexBytes(value_bytes),
1684        "decoding value",
1685    );
1686    V::from_bytes(value_bytes, decoders)
1687}
1688
1689#[track_caller]
1690fn decode_value_expect<V: DatabaseValue>(
1691    value_bytes: &[u8],
1692    decoders: &ModuleDecoderRegistry,
1693    key_bytes: &[u8],
1694) -> V {
1695    decode_value(value_bytes, decoders).unwrap_or_else(|err| {
1696        panic!(
1697            "Unrecoverable decoding DatabaseValue as {}; err={}, key_bytes={}, val_bytes={}",
1698            any::type_name::<V>(),
1699            err.fmt_compact(),
1700            AbbreviateHexBytes(key_bytes),
1701            AbbreviateHexBytes(value_bytes),
1702        )
1703    })
1704}
1705
1706#[track_caller]
1707fn decode_key_expect<K: DatabaseKey>(key_bytes: &[u8], decoders: &ModuleDecoderRegistry) -> K {
1708    trace!(
1709        bytes = %AbbreviateHexBytes(key_bytes),
1710        "decoding key",
1711    );
1712    K::from_bytes(key_bytes, decoders).unwrap_or_else(|err| {
1713        panic!(
1714            "Unrecoverable decoding DatabaseKey as {}; err={}; bytes={}",
1715            any::type_name::<K>(),
1716            err.fmt_compact(),
1717            AbbreviateHexBytes(key_bytes)
1718        )
1719    })
1720}
1721
1722impl<'tx, Cap> DatabaseTransaction<'tx, Cap> {
1723    /// Convert into a non-committable version
1724    pub fn into_nc(self) -> DatabaseTransaction<'tx, NonCommittable> {
1725        DatabaseTransaction {
1726            tx: self.tx,
1727            decoders: self.decoders,
1728            commit_tracker: self.commit_tracker,
1729            on_commit_hooks: self.on_commit_hooks,
1730            capability: PhantomData::<NonCommittable>,
1731        }
1732    }
1733
1734    /// Get a reference to a non-committeable version
1735    pub fn to_ref_nc<'s, 'a>(&'s mut self) -> DatabaseTransaction<'a, NonCommittable>
1736    where
1737        's: 'a,
1738    {
1739        self.to_ref().into_nc()
1740    }
1741
1742    /// Get [`DatabaseTransaction`] isolated to a `prefix`
1743    pub fn with_prefix<'a: 'tx>(self, prefix: Vec<u8>) -> DatabaseTransaction<'a, Cap>
1744    where
1745        'tx: 'a,
1746    {
1747        DatabaseTransaction {
1748            tx: Box::new(PrefixDatabaseTransaction {
1749                inner: self.tx,
1750                global_dbtx_access_token: None,
1751                prefix,
1752            }),
1753            decoders: self.decoders,
1754            commit_tracker: self.commit_tracker,
1755            on_commit_hooks: self.on_commit_hooks,
1756            capability: self.capability,
1757        }
1758    }
1759
1760    /// Get [`DatabaseTransaction`] isolated to a prefix of a given
1761    /// `module_instance_id`, allowing the module to access global_dbtx
1762    /// with the right access token.
1763    pub fn with_prefix_module_id<'a: 'tx>(
1764        self,
1765        module_instance_id: ModuleInstanceId,
1766    ) -> (DatabaseTransaction<'a, Cap>, GlobalDBTxAccessToken)
1767    where
1768        'tx: 'a,
1769    {
1770        let prefix = module_instance_id_to_byte_prefix(module_instance_id);
1771        let global_dbtx_access_token = GlobalDBTxAccessToken::from_prefix(&prefix);
1772        (
1773            DatabaseTransaction {
1774                tx: Box::new(PrefixDatabaseTransaction {
1775                    inner: self.tx,
1776                    global_dbtx_access_token: Some(global_dbtx_access_token),
1777                    prefix,
1778                }),
1779                decoders: self.decoders,
1780                commit_tracker: self.commit_tracker,
1781                on_commit_hooks: self.on_commit_hooks,
1782                capability: self.capability,
1783            },
1784            global_dbtx_access_token,
1785        )
1786    }
1787
1788    /// Get [`DatabaseTransaction`] to `self`
1789    pub fn to_ref<'s, 'a>(&'s mut self) -> DatabaseTransaction<'a, Cap>
1790    where
1791        's: 'a,
1792    {
1793        let decoders = Cow::Borrowed(&*self.decoders);
1794
1795        DatabaseTransaction {
1796            tx: Box::new(&mut self.tx),
1797            decoders,
1798            commit_tracker: match self.commit_tracker {
1799                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1800                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1801            },
1802            on_commit_hooks: match self.on_commit_hooks {
1803                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1804                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1805            },
1806            capability: self.capability,
1807        }
1808    }
1809
1810    /// Get [`DatabaseTransaction`] isolated to a `prefix` of `self`
1811    pub fn to_ref_with_prefix<'a>(&'a mut self, prefix: Vec<u8>) -> DatabaseTransaction<'a, Cap>
1812    where
1813        'tx: 'a,
1814    {
1815        DatabaseTransaction {
1816            tx: Box::new(PrefixDatabaseTransaction {
1817                inner: &mut self.tx,
1818                global_dbtx_access_token: None,
1819                prefix,
1820            }),
1821            decoders: Cow::Borrowed(&*self.decoders),
1822            commit_tracker: match self.commit_tracker {
1823                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1824                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1825            },
1826            on_commit_hooks: match self.on_commit_hooks {
1827                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1828                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1829            },
1830            capability: self.capability,
1831        }
1832    }
1833
1834    pub fn to_ref_with_prefix_module_id<'a>(
1835        &'a mut self,
1836        module_instance_id: ModuleInstanceId,
1837    ) -> (DatabaseTransaction<'a, Cap>, GlobalDBTxAccessToken)
1838    where
1839        'tx: 'a,
1840    {
1841        let prefix = module_instance_id_to_byte_prefix(module_instance_id);
1842        let global_dbtx_access_token = GlobalDBTxAccessToken::from_prefix(&prefix);
1843        (
1844            DatabaseTransaction {
1845                tx: Box::new(PrefixDatabaseTransaction {
1846                    inner: &mut self.tx,
1847                    global_dbtx_access_token: Some(global_dbtx_access_token),
1848                    prefix,
1849                }),
1850                decoders: Cow::Borrowed(&*self.decoders),
1851                commit_tracker: match self.commit_tracker {
1852                    MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1853                    MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1854                },
1855                on_commit_hooks: match self.on_commit_hooks {
1856                    MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1857                    MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1858                },
1859                capability: self.capability,
1860            },
1861            global_dbtx_access_token,
1862        )
1863    }
1864
1865    /// Is this `Database` a global, unpartitioned `Database`
1866    pub fn is_global(&self) -> bool {
1867        self.tx.is_global()
1868    }
1869
1870    /// `Err` if [`Self::is_global`] is not true
1871    pub fn ensure_global(&self) -> DatabaseResult<()> {
1872        if !self.is_global() {
1873            return Err(DatabaseError::NotGlobal);
1874        }
1875
1876        Ok(())
1877    }
1878
1879    /// `Err` if [`Self::is_global`] is true
1880    pub fn ensure_isolated(&self) -> DatabaseResult<()> {
1881        if self.is_global() {
1882            return Err(DatabaseError::NotIsolated);
1883        }
1884
1885        Ok(())
1886    }
1887
1888    /// Cancel the tx to avoid debugging warnings about uncommitted writes
1889    pub fn ignore_uncommitted(&mut self) -> &mut Self {
1890        self.commit_tracker.ignore_uncommitted = true;
1891        self
1892    }
1893
1894    /// Create warnings about uncommitted writes
1895    pub fn warn_uncommitted(&mut self) -> &mut Self {
1896        self.commit_tracker.ignore_uncommitted = false;
1897        self
1898    }
1899
1900    /// Register a hook that will be run after commit succeeds.
1901    #[instrument(target = LOG_DB, level = "trace", skip_all)]
1902    pub fn on_commit(&mut self, f: maybe_add_send!(impl FnOnce() + 'static)) {
1903        self.on_commit_hooks.push(Box::new(f));
1904    }
1905
1906    pub fn global_dbtx<'a>(
1907        &'a mut self,
1908        access_token: GlobalDBTxAccessToken,
1909    ) -> DatabaseTransaction<'a, Cap>
1910    where
1911        'tx: 'a,
1912    {
1913        let decoders = Cow::Borrowed(&*self.decoders);
1914
1915        DatabaseTransaction {
1916            tx: Box::new(self.tx.global_dbtx(access_token)),
1917            decoders,
1918            commit_tracker: match self.commit_tracker {
1919                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1920                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1921            },
1922            on_commit_hooks: match self.on_commit_hooks {
1923                MaybeRef::Owned(ref mut o) => MaybeRef::Borrowed(o),
1924                MaybeRef::Borrowed(ref mut b) => MaybeRef::Borrowed(b),
1925            },
1926            capability: self.capability,
1927        }
1928    }
1929}
1930
1931/// Code used to access `global_dbtx`
1932#[derive(Copy, Clone, Debug, PartialEq, Eq)]
1933pub struct GlobalDBTxAccessToken(u32);
1934
1935impl GlobalDBTxAccessToken {
1936    /// Calculate an access code for accessing global_dbtx from a prefixed
1937    /// database tx
1938    ///
1939    /// Since we need to do it at runtime, we want the user modules not to be
1940    /// able to call `global_dbtx` too easily. But at the same time we don't
1941    /// need to be paranoid.
1942    ///
1943    /// This must be deterministic during whole instance of the software running
1944    /// (because it's being rederived independently in multiple codepahs) , but
1945    /// it could be somewhat randomized between different runs and releases.
1946    fn from_prefix(prefix: &[u8]) -> Self {
1947        Self(prefix.iter().fold(0, |acc, b| acc + u32::from(*b)) + 513)
1948    }
1949}
1950
1951impl<'tx> DatabaseTransaction<'tx, Committable> {
1952    pub fn new(dbtx: Box<dyn IDatabaseTransaction + 'tx>, decoders: ModuleDecoderRegistry) -> Self {
1953        Self {
1954            tx: dbtx,
1955            decoders: Cow::Owned(decoders),
1956            commit_tracker: MaybeRef::Owned(CommitTracker {
1957                is_committed: false,
1958                has_writes: false,
1959                ignore_uncommitted: false,
1960            }),
1961            on_commit_hooks: MaybeRef::Owned(vec![]),
1962            capability: PhantomData,
1963        }
1964    }
1965
1966    pub async fn commit_tx_result(mut self) -> DatabaseResult<()> {
1967        self.commit_tracker.is_committed = true;
1968        let commit_result = self.tx.commit_tx().await;
1969
1970        // Run commit hooks in case commit was successful
1971        if commit_result.is_ok() {
1972            for hook in self.on_commit_hooks.deref_mut().drain(..) {
1973                hook();
1974            }
1975        }
1976
1977        commit_result
1978    }
1979
1980    pub async fn commit_tx(mut self) {
1981        self.commit_tracker.is_committed = true;
1982        self.commit_tx_result()
1983            .await
1984            .expect("Unrecoverable error occurred while committing to the database.");
1985    }
1986}
1987
1988#[apply(async_trait_maybe_send!)]
1989impl<Cap> IDatabaseTransactionOpsCore for DatabaseTransaction<'_, Cap>
1990where
1991    Cap: Send,
1992{
1993    async fn raw_insert_bytes_no_previous(
1994        &mut self,
1995        key: &[u8],
1996        value: &[u8],
1997    ) -> DatabaseResult<()> {
1998        self.commit_tracker.has_writes = true;
1999        self.tx.raw_insert_bytes_no_previous(key, value).await
2000    }
2001
2002    async fn raw_remove_entry_no_previous(&mut self, key: &[u8]) -> DatabaseResult<()> {
2003        self.tx.raw_remove_entry_no_previous(key).await
2004    }
2005
2006    async fn raw_insert_bytes(
2007        &mut self,
2008        key: &[u8],
2009        value: &[u8],
2010    ) -> DatabaseResult<Option<Vec<u8>>> {
2011        self.commit_tracker.has_writes = true;
2012        self.tx.raw_insert_bytes(key, value).await
2013    }
2014
2015    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
2016        self.tx.raw_get_bytes(key).await
2017    }
2018
2019    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
2020        self.tx.raw_remove_entry(key).await
2021    }
2022
2023    async fn raw_find_by_range(
2024        &mut self,
2025        key_range: Range<&[u8]>,
2026    ) -> DatabaseResult<PrefixStream<'_>> {
2027        self.tx.raw_find_by_range(key_range).await
2028    }
2029
2030    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
2031        self.tx.raw_find_by_prefix(key_prefix).await
2032    }
2033
2034    async fn raw_find_by_prefix_sorted_descending(
2035        &mut self,
2036        key_prefix: &[u8],
2037    ) -> DatabaseResult<PrefixStream<'_>> {
2038        self.tx
2039            .raw_find_by_prefix_sorted_descending(key_prefix)
2040            .await
2041    }
2042
2043    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
2044        self.commit_tracker.has_writes = true;
2045        self.tx.raw_remove_by_prefix(key_prefix).await
2046    }
2047}
2048impl IDatabaseTransactionOps for DatabaseTransaction<'_, Committable> {}
2049
2050impl<T> DatabaseKeyPrefix for T
2051where
2052    T: DatabaseLookup + crate::encoding::Encodable + Debug,
2053{
2054    fn to_bytes(&self) -> Vec<u8> {
2055        let mut data = vec![<Self as DatabaseLookup>::Record::DB_PREFIX];
2056        data.append(&mut self.consensus_encode_to_vec());
2057        data
2058    }
2059}
2060
2061impl<T> DatabaseKey for T
2062where
2063    // Note: key can only be `T` that can be decoded without modules (even if
2064    // module type is `()`)
2065    T: DatabaseRecord + crate::encoding::Decodable + Sized,
2066{
2067    const NOTIFY_ON_MODIFY: bool = <T as DatabaseRecord>::NOTIFY_ON_MODIFY;
2068    fn from_bytes(
2069        data: &[u8],
2070        modules: &ModuleDecoderRegistry,
2071    ) -> std::result::Result<Self, DecodingError> {
2072        if data.is_empty() {
2073            // TODO: build better coding errors, pretty useless right now
2074            return Err(DecodingError::wrong_length(1, 0));
2075        }
2076
2077        if data[0] != Self::DB_PREFIX {
2078            return Err(DecodingError::wrong_prefix(Self::DB_PREFIX, data[0]));
2079        }
2080
2081        <Self as crate::encoding::Decodable>::consensus_decode_whole(&data[1..], modules)
2082            .map_err(DecodingError::from)
2083    }
2084}
2085
2086impl<T> DatabaseValue for T
2087where
2088    T: Debug + Encodable + Decodable,
2089{
2090    fn from_bytes(
2091        data: &[u8],
2092        modules: &ModuleDecoderRegistry,
2093    ) -> std::result::Result<Self, DecodingError> {
2094        T::consensus_decode_whole(data, modules).map_err(DecodingError::from)
2095    }
2096
2097    fn to_bytes(&self) -> Vec<u8> {
2098        self.consensus_encode_to_vec()
2099    }
2100}
2101
2102/// This is a helper macro that generates the implementations of
2103/// `DatabaseRecord` necessary for reading/writing to the
2104/// database and fetching by prefix.
2105///
2106/// - `key`: This is the type of struct that will be used as the key into the
2107///   database
2108/// - `value`: This is the type of struct that will be used as the value into
2109///   the database
2110/// - `db_prefix`: Required enum expression that is represented as a `u8` and is
2111///   prepended to this key
2112/// - `query_prefix`: Optional type of struct that can be passed zero or more
2113///   times. Every query prefix can be used to query the database via
2114///   `find_by_prefix`
2115///
2116/// # Examples
2117///
2118/// ```
2119/// use fedimint_core::encoding::{Decodable, Encodable};
2120/// use fedimint_core::impl_db_record;
2121///
2122/// #[derive(Debug, Encodable, Decodable)]
2123/// struct MyKey;
2124///
2125/// #[derive(Debug, Encodable, Decodable)]
2126/// struct MyValue;
2127///
2128/// #[repr(u8)]
2129/// #[derive(Clone, Debug)]
2130/// pub enum DbKeyPrefix {
2131///     MyKey = 0x50,
2132/// }
2133///
2134/// impl_db_record!(key = MyKey, value = MyValue, db_prefix = DbKeyPrefix::MyKey);
2135/// ```
2136///
2137/// Use the required parameters and specify one `query_prefix`
2138///
2139/// ```
2140/// use fedimint_core::encoding::{Decodable, Encodable};
2141/// use fedimint_core::{impl_db_lookup, impl_db_record};
2142///
2143/// #[derive(Debug, Encodable, Decodable)]
2144/// struct MyKey;
2145///
2146/// #[derive(Debug, Encodable, Decodable)]
2147/// struct MyValue;
2148///
2149/// #[repr(u8)]
2150/// #[derive(Clone, Debug)]
2151/// pub enum DbKeyPrefix {
2152///     MyKey = 0x50,
2153/// }
2154///
2155/// #[derive(Debug, Encodable, Decodable)]
2156/// struct MyKeyPrefix;
2157///
2158/// impl_db_record!(key = MyKey, value = MyValue, db_prefix = DbKeyPrefix::MyKey,);
2159///
2160/// impl_db_lookup!(key = MyKey, query_prefix = MyKeyPrefix);
2161/// ```
2162#[macro_export]
2163macro_rules! impl_db_record {
2164    (key = $key:ty, value = $val:ty, db_prefix = $db_prefix:expr_2021 $(, notify_on_modify = $notify:tt)? $(,)?) => {
2165        impl $crate::db::DatabaseRecord for $key {
2166            const DB_PREFIX: u8 = $db_prefix as u8;
2167            $(const NOTIFY_ON_MODIFY: bool = $notify;)?
2168            type Key = Self;
2169            type Value = $val;
2170        }
2171        $(
2172            impl_db_record! {
2173                @impl_notify_marker key = $key, notify_on_modify = $notify
2174            }
2175        )?
2176    };
2177    // if notify is set to true
2178    (@impl_notify_marker key = $key:ty, notify_on_modify = true) => {
2179        impl $crate::db::DatabaseKeyWithNotify for $key {}
2180    };
2181    // if notify is set to false
2182    (@impl_notify_marker key = $key:ty, notify_on_modify = false) => {};
2183}
2184
2185#[macro_export]
2186macro_rules! impl_db_lookup{
2187    (key = $key:ty $(, query_prefix = $query_prefix:ty)* $(,)?) => {
2188        $(
2189            impl $crate::db::DatabaseLookup for $query_prefix {
2190                type Record = $key;
2191            }
2192        )*
2193    };
2194}
2195
2196/// Deprecated: Use `DatabaseVersionKey(ModuleInstanceId)` instead.
2197#[derive(Debug, Encodable, Decodable, Serialize)]
2198pub struct DatabaseVersionKeyV0;
2199
2200#[derive(Debug, Encodable, Decodable, Serialize)]
2201pub struct DatabaseVersionKey(pub ModuleInstanceId);
2202
2203#[derive(Debug, Encodable, Decodable, Serialize, Clone, PartialOrd, Ord, PartialEq, Eq, Copy)]
2204pub struct DatabaseVersion(pub u64);
2205
2206impl_db_record!(
2207    key = DatabaseVersionKeyV0,
2208    value = DatabaseVersion,
2209    db_prefix = DbKeyPrefix::DatabaseVersion
2210);
2211
2212impl_db_record!(
2213    key = DatabaseVersionKey,
2214    value = DatabaseVersion,
2215    db_prefix = DbKeyPrefix::DatabaseVersion
2216);
2217
2218impl std::fmt::Display for DatabaseVersion {
2219    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2220        write!(f, "{}", self.0)
2221    }
2222}
2223
2224impl DatabaseVersion {
2225    pub fn increment(&self) -> Self {
2226        Self(self.0 + 1)
2227    }
2228}
2229
2230impl std::fmt::Display for DbKeyPrefix {
2231    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2232        write!(f, "{self:?}")
2233    }
2234}
2235
2236#[repr(u8)]
2237#[derive(Clone, EnumIter, Debug)]
2238pub enum DbKeyPrefix {
2239    DatabaseVersion = 0x50,
2240    ClientBackup = 0x51,
2241}
2242
2243#[derive(Debug, Error)]
2244#[non_exhaustive]
2245pub enum DecodingError {
2246    #[error("Key had a wrong prefix, expected {expected} but got {found}")]
2247    WrongPrefix { expected: u8, found: u8 },
2248    #[error("Key had a wrong length, expected {expected} but got {found}")]
2249    WrongLength { expected: usize, found: usize },
2250    #[error("Other decoding error")]
2251    Other(#[source] Box<dyn Error + Send + Sync>),
2252    /// The bytes are not a valid consensus encoding of the record type.
2253    #[error("Invalid consensus encoding")]
2254    Decode(#[from] DecodeError),
2255}
2256
2257impl DecodingError {
2258    pub fn other<E: Error + Send + Sync + 'static>(error: E) -> Self {
2259        Self::Other(Box::new(error))
2260    }
2261
2262    pub fn wrong_prefix(expected: u8, found: u8) -> Self {
2263        Self::WrongPrefix { expected, found }
2264    }
2265
2266    pub fn wrong_length(expected: usize, found: usize) -> Self {
2267        Self::WrongLength { expected, found }
2268    }
2269}
2270
2271/// Error type for database operations
2272#[derive(Debug, Error)]
2273#[non_exhaustive]
2274pub enum DatabaseError {
2275    /// Write-write conflict during optimistic transaction commit.
2276    /// This occurs when two transactions attempt to modify the same key.
2277    #[error("Write-write conflict detected")]
2278    WriteConflict,
2279
2280    /// The backend could not tell whether this transaction conflicted with
2281    /// another one, because its snapshot is older than the write history the
2282    /// backend still keeps around for conflict detection.
2283    ///
2284    /// This is *not* a conflict: it says nothing about whether any other
2285    /// transaction touched the same keys, only that the transaction was open
2286    /// for too many intervening writes to check. Nothing was written, and
2287    /// rerunning the whole operation against a fresh transaction is the only
2288    /// supported recovery — the failed transaction cannot be committed again.
2289    ///
2290    /// The usual cause is a transaction held open across something slow.
2291    #[error("Transaction snapshot is older than the retained write history: {0}")]
2292    SnapshotTooOld(Box<dyn Error + Send + Sync>),
2293
2294    /// The transaction has already been consumed (committed or dropped).
2295    /// Operations cannot be performed on a consumed transaction.
2296    #[error("Transaction already consumed")]
2297    TransactionConsumed,
2298
2299    /// Error from the underlying database backend (e.g., RocksDB I/O errors).
2300    #[error("Database backend error: {0}")]
2301    DatabaseBackend(#[from] Box<dyn Error + Send + Sync>),
2302
2303    /// The operation requires a global (non module-isolated) database instance.
2304    #[error("Database instance is not global")]
2305    NotGlobal,
2306
2307    /// The operation requires a module-isolated database instance.
2308    #[error("Database instance is not isolated")]
2309    NotIsolated,
2310}
2311
2312impl DatabaseError {
2313    /// Create a DatabaseBackend error from any error type
2314    pub fn backend<E: Error + Send + Sync + 'static>(error: E) -> Self {
2315        Self::DatabaseBackend(Box::new(error))
2316    }
2317
2318    /// Create a `SnapshotTooOld` error, preserving the backend's own error
2319    pub fn snapshot_too_old<E: Error + Send + Sync + 'static>(error: E) -> Self {
2320        Self::SnapshotTooOld(Box::new(error))
2321    }
2322}
2323
2324#[macro_export]
2325macro_rules! push_db_pair_items {
2326    ($dbtx:ident, $prefix_type:expr_2021, $key_type:ty, $value_type:ty, $map:ident, $key_literal:literal) => {
2327        let db_items =
2328            $crate::db::IDatabaseTransactionOpsCoreTyped::find_by_prefix($dbtx, &$prefix_type)
2329                .await
2330                .map(|(key, val)| {
2331                    (
2332                        $crate::encoding::Encodable::consensus_encode_to_hex(&key),
2333                        val,
2334                    )
2335                })
2336                .collect::<BTreeMap<String, $value_type>>()
2337                .await;
2338
2339        $map.insert($key_literal.to_string(), Box::new(db_items));
2340    };
2341}
2342
2343#[macro_export]
2344macro_rules! push_db_key_items {
2345    ($dbtx:ident, $prefix_type:expr_2021, $key_type:ty, $map:ident, $key_literal:literal) => {
2346        let db_items =
2347            $crate::db::IDatabaseTransactionOpsCoreTyped::find_by_prefix($dbtx, &$prefix_type)
2348                .await
2349                .map(|(key, _)| key)
2350                .collect::<Vec<$key_type>>()
2351                .await;
2352
2353        $map.insert($key_literal.to_string(), Box::new(db_items));
2354    };
2355}
2356
2357/// Context passed to the db migration _functions_ (pay attention to `Fn` in the
2358/// name)
2359///
2360/// Typically should not be referred to directly, and instead by a type-alias,
2361/// where the inner-context is set.
2362///
2363/// Notably it has the (optional) module id (inaccessible to the modules
2364/// directly, but used internally) and an inner context `C` injected by the
2365/// outer-layer.
2366///
2367/// `C` is generic, as in different layers / scopes (server vs client, etc.) a
2368/// different (module-typed, type erased, server/client, etc.) contexts might be
2369/// needed, while the database migration logic is kind of generic over that.
2370pub struct DbMigrationFnContext<'tx, C> {
2371    dbtx: DatabaseTransaction<'tx>,
2372    module_instance_id: Option<ModuleInstanceId>,
2373    ctx: C,
2374    __please_use_constructor: (),
2375}
2376
2377impl<'tx, C> DbMigrationFnContext<'tx, C> {
2378    pub fn new(
2379        dbtx: DatabaseTransaction<'tx>,
2380        module_instance_id: Option<ModuleInstanceId>,
2381        ctx: C,
2382    ) -> Self {
2383        dbtx.ensure_global().expect("Must pass global dbtx");
2384        Self {
2385            dbtx,
2386            module_instance_id,
2387            ctx,
2388            // this is a constructor
2389            __please_use_constructor: (),
2390        }
2391    }
2392
2393    pub fn map<R>(self, f: impl FnOnce(C) -> R) -> DbMigrationFnContext<'tx, R> {
2394        DbMigrationFnContext::new(self.dbtx, self.module_instance_id, f(self.ctx))
2395    }
2396
2397    // TODO: this method is currently visible to the module itself, and it shouldn't
2398    #[doc(hidden)]
2399    pub fn split_dbtx_ctx<'s>(&'s mut self) -> (&'s mut DatabaseTransaction<'tx>, &'s C) {
2400        let Self { dbtx, ctx, .. } = self;
2401
2402        (dbtx, ctx)
2403    }
2404
2405    pub fn dbtx(&'_ mut self) -> DatabaseTransaction<'_> {
2406        if let Some(module_instance_id) = self.module_instance_id {
2407            self.dbtx.to_ref_with_prefix_module_id(module_instance_id).0
2408        } else {
2409            self.dbtx.to_ref_nc()
2410        }
2411    }
2412
2413    // TODO: this method is currently visible to the module itself, and it shouldn't
2414    #[doc(hidden)]
2415    pub fn module_instance_id(&self) -> Option<ModuleInstanceId> {
2416        self.module_instance_id
2417    }
2418}
2419
2420/// [`DbMigrationFn`] with no extra context (ATM gateway)
2421pub type GeneralDbMigrationFn = DbMigrationFn<()>;
2422pub type GeneralDbMigrationFnContext<'tx> = DbMigrationFnContext<'tx, ()>;
2423
2424/// [`DbMigrationFn`] used by core client
2425///
2426/// NOTE: client _module_ migrations are handled using separate structs due to
2427/// state machine migrations
2428pub type ClientCoreDbMigrationFn = DbMigrationFn<()>;
2429pub type ClientCoreDbMigrationFnContext<'tx> = DbMigrationFnContext<'tx, ()>;
2430
2431/// `CoreMigrationFn` that modules can implement to "migrate" the database
2432/// to the next database version.
2433///
2434/// It is parametrized over `C` (contents), which is extra data/type/interface
2435/// custom for different part of the codebase, e.g.:
2436///
2437/// * server core
2438/// * server modules
2439/// * client core
2440/// * gateway core
2441pub type DbMigrationFn<C> = Box<
2442    maybe_add_send_sync!(
2443        dyn for<'tx> Fn(
2444            DbMigrationFnContext<'tx, C>,
2445        ) -> Pin<
2446            Box<maybe_add_send!(dyn futures::Future<Output = Result<(), DbMigrationError>> + 'tx)>,
2447        >
2448    ),
2449>;
2450
2451/// Failure while applying database migrations.
2452#[derive(Debug, Error)]
2453#[non_exhaustive]
2454pub enum DbMigrationError {
2455    /// The database itself failed.
2456    #[error("Database error")]
2457    Database(#[from] DatabaseError),
2458    /// A stored value is not a valid consensus encoding.
2459    #[error("Failed to consensus-decode a database entry")]
2460    Decode(#[from] DecodeError),
2461    /// The database was written by newer code than the one applying migrations.
2462    #[error(
2463        "On disk database version {on_disk} for module {kind} is higher than the code \
2464         database version {target}"
2465    )]
2466    VersionTooHigh {
2467        kind: String,
2468        on_disk: DatabaseVersion,
2469        target: DatabaseVersion,
2470    },
2471    /// A migration failed for a reason specific to it.
2472    #[error("Migration failed")]
2473    Other(#[source] Box<dyn Error + Send + Sync>),
2474}
2475
2476impl DbMigrationError {
2477    /// Wraps a migration-specific error; accepts anything convertible into a
2478    /// boxed error, an `anyhow::Error` included.
2479    pub fn other<E>(error: E) -> Self
2480    where
2481        E: Into<Box<dyn Error + Send + Sync>>,
2482    {
2483        Self::Other(error.into())
2484    }
2485}
2486
2487/// Verifies that all database migrations are defined contiguously and returns
2488/// the "current" database version, which is one greater than the last key in
2489/// the map.
2490pub fn get_current_database_version<F>(
2491    migrations: &BTreeMap<DatabaseVersion, F>,
2492) -> DatabaseVersion {
2493    let versions = migrations.keys().copied().collect::<Vec<_>>();
2494
2495    // Verify that all database migrations are defined contiguously. If there is a
2496    // gap, this indicates a programming error and we should panic.
2497    if !versions
2498        .windows(2)
2499        .all(|window| window[0].increment() == window[1])
2500    {
2501        panic!("Database Migrations are not defined contiguously");
2502    }
2503
2504    versions
2505        .last()
2506        .map_or(DatabaseVersion(0), DatabaseVersion::increment)
2507}
2508
2509pub async fn apply_migrations<C>(
2510    db: &Database,
2511    ctx: C,
2512    kind: String,
2513    migrations: BTreeMap<DatabaseVersion, DbMigrationFn<C>>,
2514    module_instance_id: Option<ModuleInstanceId>,
2515    // When used in client side context, we can/should ignore keys that external app
2516    // is allowed to use, and but since this function is shared, we make it optional argument
2517    external_prefixes_above: Option<u8>,
2518) -> Result<(), DbMigrationError>
2519where
2520    C: Clone,
2521{
2522    let mut dbtx = db.begin_transaction().await;
2523    apply_migrations_dbtx(
2524        &mut dbtx.to_ref_nc(),
2525        ctx,
2526        kind,
2527        migrations,
2528        module_instance_id,
2529        external_prefixes_above,
2530    )
2531    .await?;
2532
2533    Ok(dbtx.commit_tx_result().await?)
2534}
2535/// `apply_migrations` iterates from the on disk database version for the
2536/// module.
2537///
2538/// `apply_migrations` iterates from the on disk database version for the module
2539/// up to `target_db_version` and executes all of the migrations that exist in
2540/// the migrations map. Each migration in migrations map updates the
2541/// database to have the correct on-disk structures that the code is expecting.
2542/// The entire migration process is atomic (i.e migration from 0->1 and 1->2
2543/// happen atomically). This function is called before the module is initialized
2544/// and as long as the correct migrations are supplied in the migrations map,
2545/// the module will be able to read and write from the database successfully.
2546pub async fn apply_migrations_dbtx<C>(
2547    global_dbtx: &mut DatabaseTransaction<'_>,
2548    ctx: C,
2549    kind: String,
2550    migrations: BTreeMap<DatabaseVersion, DbMigrationFn<C>>,
2551    module_instance_id: Option<ModuleInstanceId>,
2552    // When used in client side context, we can/should ignore keys that external app
2553    // is allowed to use, and but since this function is shared, we make it optional argument
2554    external_prefixes_above: Option<u8>,
2555) -> Result<(), DbMigrationError>
2556where
2557    C: Clone,
2558{
2559    // Newly created databases will not have any data since they have just been
2560    // instantiated.
2561    let is_new_db = global_dbtx
2562        .raw_find_by_prefix(&[])
2563        .await?
2564        .filter(|(key, _v)| {
2565            std::future::ready(
2566                external_prefixes_above.is_none_or(|external_prefixes_above| {
2567                    !key.is_empty() && key[0] < external_prefixes_above
2568                }),
2569            )
2570        })
2571        .next()
2572        .await
2573        .is_none();
2574
2575    let target_db_version = get_current_database_version(&migrations);
2576
2577    // First write the database version to disk if it does not exist.
2578    create_database_version_dbtx(
2579        global_dbtx,
2580        target_db_version,
2581        module_instance_id,
2582        kind.clone(),
2583        is_new_db,
2584    )
2585    .await;
2586
2587    let module_instance_id_key = module_instance_id_or_global(module_instance_id);
2588
2589    let disk_version = global_dbtx
2590        .get_value(&DatabaseVersionKey(module_instance_id_key))
2591        .await;
2592
2593    let db_version = if let Some(disk_version) = disk_version {
2594        let mut current_db_version = disk_version;
2595
2596        if current_db_version > target_db_version {
2597            return Err(DbMigrationError::VersionTooHigh {
2598                kind,
2599                on_disk: current_db_version,
2600                target: target_db_version,
2601            });
2602        }
2603
2604        while current_db_version < target_db_version {
2605            if let Some(migration) = migrations.get(&current_db_version) {
2606                info!(target: LOG_DB, ?kind, ?current_db_version, ?target_db_version, "Migrating module...");
2607                migration(DbMigrationFnContext::new(
2608                    global_dbtx.to_ref_nc(),
2609                    module_instance_id,
2610                    ctx.clone(),
2611                ))
2612                .await?;
2613            } else {
2614                warn!(target: LOG_DB, ?current_db_version, "Missing server db migration");
2615            }
2616
2617            current_db_version = current_db_version.increment();
2618
2619            global_dbtx
2620                .insert_entry(
2621                    &DatabaseVersionKey(module_instance_id_key),
2622                    &current_db_version,
2623                )
2624                .await;
2625        }
2626
2627        current_db_version
2628    } else {
2629        target_db_version
2630    };
2631
2632    debug!(target: LOG_DB, ?kind, ?db_version, "DB Version");
2633    Ok(())
2634}
2635
2636pub async fn create_database_version(
2637    db: &Database,
2638    target_db_version: DatabaseVersion,
2639    module_instance_id: Option<ModuleInstanceId>,
2640    kind: String,
2641    is_new_db: bool,
2642) -> Result<(), DbMigrationError> {
2643    let mut dbtx = db.begin_transaction().await;
2644
2645    create_database_version_dbtx(
2646        &mut dbtx.to_ref_nc(),
2647        target_db_version,
2648        module_instance_id,
2649        kind,
2650        is_new_db,
2651    )
2652    .await;
2653
2654    dbtx.commit_tx_result().await?;
2655    Ok(())
2656}
2657
2658/// Creates the `DatabaseVersion` inside the database if it does not exist. If
2659/// necessary, this function will migrate the legacy database version to the
2660/// expected `DatabaseVersionKey`.
2661pub async fn create_database_version_dbtx(
2662    global_dbtx: &mut DatabaseTransaction<'_>,
2663    target_db_version: DatabaseVersion,
2664    module_instance_id: Option<ModuleInstanceId>,
2665    kind: String,
2666    is_new_db: bool,
2667) {
2668    let key_module_instance_id = module_instance_id_or_global(module_instance_id);
2669
2670    // First check if the module has a `DatabaseVersion` written to
2671    // `DatabaseVersionKey`. If `DatabaseVersion` already exists, there is
2672    // nothing to do.
2673    if global_dbtx
2674        .get_value(&DatabaseVersionKey(key_module_instance_id))
2675        .await
2676        .is_none()
2677    {
2678        // If it exists, read and remove the legacy `DatabaseVersion`, which used to be
2679        // in the module's isolated namespace (but not for fedimint-server or
2680        // fedimint-client).
2681        //
2682        // Otherwise, if the previous database contains data and no legacy database
2683        // version, use `DatabaseVersion(0)` so that all database migrations are
2684        // run. Otherwise, this database can assumed to be new and can use
2685        // `target_db_version` to skip the database migrations.
2686        let current_version_in_module = if let Some(module_instance_id) = module_instance_id {
2687            remove_current_db_version_if_exists(
2688                &mut global_dbtx
2689                    .to_ref_with_prefix_module_id(module_instance_id)
2690                    .0
2691                    .into_nc(),
2692                is_new_db,
2693                target_db_version,
2694            )
2695            .await
2696        } else {
2697            remove_current_db_version_if_exists(
2698                &mut global_dbtx.to_ref().into_nc(),
2699                is_new_db,
2700                target_db_version,
2701            )
2702            .await
2703        };
2704
2705        // Write the previous `DatabaseVersion` to the new `DatabaseVersionKey`
2706        debug!(target: LOG_DB, ?kind, ?current_version_in_module, ?target_db_version, ?is_new_db, "Creating DatabaseVersionKey...");
2707        global_dbtx
2708            .insert_new_entry(
2709                &DatabaseVersionKey(key_module_instance_id),
2710                &current_version_in_module,
2711            )
2712            .await;
2713    }
2714}
2715
2716/// Removes `DatabaseVersion` from `DatabaseVersionKeyV0` if it exists and
2717/// returns the current database version. If the current version does not
2718/// exist, use `target_db_version` if the database is new. Otherwise, return
2719/// `DatabaseVersion(0)` to ensure all migrations are run.
2720async fn remove_current_db_version_if_exists(
2721    version_dbtx: &mut DatabaseTransaction<'_>,
2722    is_new_db: bool,
2723    target_db_version: DatabaseVersion,
2724) -> DatabaseVersion {
2725    // Remove the previous `DatabaseVersion` in the isolated database. If it doesn't
2726    // exist, just use the 0 for the version so that all of the migrations are
2727    // executed.
2728    let current_version_in_module = version_dbtx.remove_entry(&DatabaseVersionKeyV0).await;
2729    match current_version_in_module {
2730        Some(database_version) => database_version,
2731        None if is_new_db => target_db_version,
2732        None => DatabaseVersion(0),
2733    }
2734}
2735
2736/// Helper function to retrieve the `module_instance_id` for modules, otherwise
2737/// return 0xff for the global namespace.
2738fn module_instance_id_or_global(module_instance_id: Option<ModuleInstanceId>) -> ModuleInstanceId {
2739    // Use 0xff for fedimint-server and the `module_instance_id` for each module
2740    module_instance_id.unwrap_or_else(|| MODULE_GLOBAL_PREFIX.into())
2741}
2742#[allow(unused_imports)]
2743mod test_utils;
2744
2745pub async fn find_by_prefix_sorted_descending<'r, 'inner, KP>(
2746    tx: &'r mut (dyn IDatabaseTransaction + 'inner),
2747    decoders: ModuleDecoderRegistry,
2748    key_prefix: &KP,
2749) -> impl Stream<
2750    Item = (
2751        KP::Record,
2752        <<KP as DatabaseLookup>::Record as DatabaseRecord>::Value,
2753    ),
2754>
2755+ 'r
2756+ use<'r, KP>
2757where
2758    'inner: 'r,
2759    KP: DatabaseLookup,
2760    KP::Record: DatabaseKey,
2761{
2762    debug!(target: LOG_DB, "find by prefix sorted descending");
2763    let prefix_bytes = key_prefix.to_bytes();
2764    tx.raw_find_by_prefix_sorted_descending(&prefix_bytes)
2765        .await
2766        .expect("Error doing prefix search in database")
2767        .map(move |(key_bytes, value_bytes)| {
2768            let key = decode_key_expect(&key_bytes, &decoders);
2769            let value = decode_value_expect(&value_bytes, &decoders, &key_bytes);
2770            (key, value)
2771        })
2772}
2773
2774pub async fn verify_module_db_integrity_dbtx(
2775    dbtx: &mut DatabaseTransaction<'_>,
2776    module_id: ModuleInstanceId,
2777    module_kind: ModuleKind,
2778    prefixes: &BTreeSet<u8>,
2779) {
2780    let module_db_prefix = module_instance_id_to_byte_prefix(module_id);
2781    if module_id < 250 {
2782        assert_eq!(module_db_prefix.len(), 2);
2783    }
2784    let mut records = dbtx
2785        .raw_find_by_prefix(&module_db_prefix)
2786        .await
2787        .expect("DB fail");
2788    while let Some((k, v)) = records.next().await {
2789        assert!(
2790            prefixes.contains(&k[module_db_prefix.len()]),
2791            "Unexpected module {module_kind} {module_id} db record found: {}: {}",
2792            k.as_hex(),
2793            v.as_hex()
2794        );
2795    }
2796}
2797
2798#[cfg(test)]
2799mod tests;