reth_db/implementation/mdbx/

tx.rs

//! Transaction wrapper for libmdbx-sys.

use super::{cursor::Cursor, utils::*, TableDbis};
use crate::{
    metrics::{DatabaseEnvMetrics, Operation, TransactionMode, TransactionOutcome},
    DatabaseError, Tables,
};
use reth_db_api::{
    table::{Compress, DupSort, Encode, IntoVec, Table, TableImporter},
    transaction::{DbTx, DbTxMut},
};
use reth_libmdbx::{ffi::MDBX_dbi, CommitLatency, Transaction, TransactionKind, WriteFlags, RW};
use reth_primitives_traits::FastInstant as Instant;
use reth_storage_errors::db::{DatabaseWriteError, DatabaseWriteOperation};
use std::{
    backtrace::Backtrace,
    marker::PhantomData,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};
use tracing::{debug, instrument, trace, warn};

/// Duration after which we emit the log about long-lived database transactions.
const LONG_TRANSACTION_DURATION: Duration = Duration::from_secs(60);

/// Wrapper for the libmdbx transaction.
#[derive(Debug)]
pub struct Tx<K: TransactionKind> {
    /// Libmdbx-sys transaction.
    inner: Transaction<K>,

    /// Cached MDBX DBIs for reuse.
    dbis: Arc<TableDbis>,

    /// Handler for metrics with its own [Drop] implementation for cases when the transaction isn't
    /// closed by [`Tx::commit`] or [`Tx::abort`], but we still need to report it in the metrics.
    ///
    /// If [Some], then metrics are reported.
    metrics_handler: Option<MetricsHandler<K>>,
    /// Persistent read-ahead setting for gets and newly created cursors.
    #[cfg(target_os = "linux")]
    prefetch: AtomicBool,
}

impl<K: TransactionKind> Tx<K> {
    /// Creates new `Tx` object with a `RO` or `RW` transaction and optionally enables metrics.
    #[inline]
    #[track_caller]
    pub(crate) fn new(
        inner: Transaction<K>,
        dbis: Arc<TableDbis>,
        env_metrics: Option<Arc<DatabaseEnvMetrics>>,
    ) -> reth_libmdbx::Result<Self> {
        let metrics_handler = env_metrics
            .map(|env_metrics| {
                let handler = MetricsHandler::<K>::new(inner.id()?, env_metrics);
                handler.env_metrics.record_opened_transaction(handler.transaction_mode());
                handler.log_transaction_opened();
                Ok(handler)
            })
            .transpose()?;
        Ok(Self {
            inner,
            dbis,
            metrics_handler,
            #[cfg(target_os = "linux")]
            prefetch: AtomicBool::new(false),
        })
    }

    /// Returns a reference to the inner libmdbx transaction.
    pub const fn inner(&self) -> &Transaction<K> {
        &self.inner
    }

    /// Gets this transaction ID.
    pub fn id(&self) -> reth_libmdbx::Result<u64> {
        self.metrics_handler.as_ref().map_or_else(|| self.inner.id(), |handler| Ok(handler.txn_id))
    }

    /// Gets a table database handle by name if it exists, otherwise, check the
    /// database, opening the DB if it exists.
    pub fn get_dbi_raw(&self, name: &str) -> Result<MDBX_dbi, DatabaseError> {
        self.get_table_dbi(name.parse::<Tables>().ok(), name)
    }

    /// Gets a table database handle by name if it exists, otherwise, check the
    /// database, opening the DB if it exists.
    pub fn get_dbi<T: Table>(&self) -> Result<MDBX_dbi, DatabaseError> {
        self.get_table_dbi(T::TABLE, T::NAME)
    }

    /// Gets the tracked database handle of a table, which is looked up by `table` if it is a
    /// [`Tables`] entry and by `name` otherwise, opening the DB by name if it is not tracked.
    fn get_table_dbi(&self, table: Option<Tables>, name: &str) -> Result<MDBX_dbi, DatabaseError> {
        if let Some(dbi) = self.dbis.get(table, name) {
            Ok(dbi)
        } else {
            self.inner
                .open_db(Some(name))
                .map(|db| db.dbi())
                .map_err(|e| DatabaseError::Open(e.into()))
        }
    }

    /// Create db Cursor
    pub fn new_cursor<T: Table>(&self) -> Result<Cursor<K, T>, DatabaseError> {
        let inner = self
            .inner
            .cursor_with_dbi(self.get_dbi::<T>()?)
            .map_err(|e| DatabaseError::InitCursor(e.into()))?;

        Ok(Cursor::new_with_metrics(
            inner,
            self.metrics_handler
                .as_ref()
                .zip(T::TABLE)
                .map(|(h, table)| h.env_metrics.table_operation_metrics(table)),
            #[cfg(target_os = "linux")]
            self.prefetch.load(Ordering::Relaxed),
        ))
    }

    /// If `self.metrics_handler == Some(_)`, measure the time it takes to execute the closure and
    /// record a metric with the provided transaction outcome.
    ///
    /// Otherwise, just execute the closure.
    fn execute_with_close_transaction_metric<R>(
        mut self,
        outcome: TransactionOutcome,
        f: impl FnOnce(Self) -> (R, Option<CommitLatency>),
    ) -> R {
        let run = |tx| {
            let start = Instant::now();
            let (result, commit_latency) = f(tx);
            let total_duration = start.elapsed();

            if outcome.is_commit() {
                debug!(
                    target: "storage::db::mdbx",
                    ?total_duration,
                    ?commit_latency,
                    is_read_only = K::IS_READ_ONLY,
                    "Commit"
                );
            }

            (result, commit_latency, total_duration)
        };

        if let Some(mut metrics_handler) = self.metrics_handler.take() {
            metrics_handler.close_recorded = true;
            metrics_handler.log_backtrace_on_long_read_transaction();

            let (result, commit_latency, close_duration) = run(self);
            let open_duration = metrics_handler.start.elapsed();
            metrics_handler.env_metrics.record_closed_transaction(
                metrics_handler.transaction_mode(),
                outcome,
                open_duration,
                Some(close_duration),
                commit_latency,
            );

            result
        } else {
            run(self).0
        }
    }

    /// If `self.metrics_handler == Some(_)`, measure the time it takes to execute the closure and
    /// record a metric with the provided operation.
    ///
    /// Otherwise, just execute the closure.
    fn execute_with_operation_metric<T: Table, R>(
        &self,
        operation: Operation,
        value_size: Option<usize>,
        f: impl FnOnce(&Transaction<K>) -> R,
    ) -> R {
        if let Some(metrics_handler) = &self.metrics_handler {
            metrics_handler.log_backtrace_on_long_read_transaction();
            metrics_handler
                .env_metrics
                .record_operation(T::TABLE, operation, value_size, || f(&self.inner))
        } else {
            f(&self.inner)
        }
    }
}

#[derive(Debug)]
struct MetricsHandler<K: TransactionKind> {
    /// Cached internal transaction ID provided by libmdbx.
    txn_id: u64,
    /// The time when transaction has started.
    ///
    /// This is a TSC-backed [`reth_primitives_traits::FastInstant`] rather than
    /// [`std::time::Instant`] because
    /// [`MetricsHandler::log_backtrace_on_long_read_transaction`] reads it on every database
    /// operation.
    start: Instant,
    /// The time after which we emit the log about long-lived database transactions.
    ///
    /// Stored as a deadline so that the check on every database operation only compares two
    /// timestamps instead of computing the open duration.
    long_transaction_deadline: Instant,
    /// If `true`, the metric about transaction closing has already been recorded and we don't need
    /// to do anything on [`Drop::drop`].
    close_recorded: bool,
    /// If `true`, the backtrace of transaction will be recorded and logged.
    /// See [`MetricsHandler::log_backtrace_on_long_read_transaction`].
    record_backtrace: bool,
    /// If `true`, the backtrace of transaction has already been recorded and logged.
    /// See [`MetricsHandler::log_backtrace_on_long_read_transaction`].
    backtrace_recorded: AtomicBool,
    /// Shared database environment metrics.
    env_metrics: Arc<DatabaseEnvMetrics>,
    /// Backtrace of the location where the transaction has been opened. Reported only with debug
    /// assertions outside DST, because capturing every transaction opening is expensive.
    #[cfg(all(debug_assertions, not(feature = "fuzz")))]
    open_backtrace: Backtrace,
    _marker: PhantomData<K>,
}

impl<K: TransactionKind> MetricsHandler<K> {
    fn new(txn_id: u64, env_metrics: Arc<DatabaseEnvMetrics>) -> Self {
        let start = Instant::now();
        Self {
            txn_id,
            start,
            long_transaction_deadline: start + LONG_TRANSACTION_DURATION,
            close_recorded: false,
            record_backtrace: true,
            backtrace_recorded: AtomicBool::new(false),
            #[cfg(all(debug_assertions, not(feature = "fuzz")))]
            open_backtrace: Backtrace::force_capture(),
            env_metrics,
            _marker: PhantomData,
        }
    }

    const fn transaction_mode(&self) -> TransactionMode {
        if K::IS_READ_ONLY {
            TransactionMode::ReadOnly
        } else {
            TransactionMode::ReadWrite
        }
    }

    /// Logs the caller location and ID of the transaction that was opened.
    #[track_caller]
    fn log_transaction_opened(&self) {
        trace!(
            target: "storage::db::mdbx",
            caller = %core::panic::Location::caller(),
            id = %self.txn_id,
            mode = %self.transaction_mode().as_str(),
            "Transaction opened",
        );
    }

    /// Logs the backtrace of current call if the duration that the read transaction has been open
    /// is more than [`LONG_TRANSACTION_DURATION`] and `record_backtrace == true`.
    /// The backtrace is recorded and logged just once, guaranteed by `backtrace_recorded` atomic.
    ///
    /// NOTE: Backtrace is recorded using [`Backtrace::force_capture`], so `RUST_BACKTRACE` env var
    /// is not needed.
    fn log_backtrace_on_long_read_transaction(&self) {
        if self.record_backtrace &&
            !self.backtrace_recorded.load(Ordering::Relaxed) &&
            self.transaction_mode().is_read_only()
        {
            let now = Instant::now();
            if now >= self.long_transaction_deadline {
                let open_duration = now - self.start;
                self.backtrace_recorded.store(true, Ordering::Relaxed);
                #[cfg(all(debug_assertions, not(feature = "fuzz")))]
                let open_backtrace = format_args!("{}", self.open_backtrace);
                #[cfg(not(all(debug_assertions, not(feature = "fuzz"))))]
                let open_backtrace = tracing::field::Empty;
                warn!(
                    target: "storage::db::mdbx",
                    ?open_duration,
                    id=%self.txn_id,
                    backtrace=%Backtrace::force_capture(),
                    open_backtrace,
                    "A database read transaction has been open for too long"
                );
            }
        }
    }
}

impl<K: TransactionKind> Drop for MetricsHandler<K> {
    fn drop(&mut self) {
        if !self.close_recorded {
            self.log_backtrace_on_long_read_transaction();
            self.env_metrics.record_closed_transaction(
                self.transaction_mode(),
                TransactionOutcome::Drop,
                self.start.elapsed(),
                None,
                None,
            );
        }
    }
}

impl TableImporter for Tx<RW> {}

impl<K: TransactionKind> DbTx for Tx<K> {
    type Cursor<T: Table> = Cursor<K, T>;
    type DupCursor<T: DupSort> = Cursor<K, T>;

    fn get<T: Table>(&self, key: T::Key) -> Result<Option<<T as Table>::Value>, DatabaseError> {
        self.get_by_encoded_key::<T>(&key.encode())
    }

    fn get_by_encoded_key<T: Table>(
        &self,
        key: &<T::Key as Encode>::Encoded,
    ) -> Result<Option<T::Value>, DatabaseError> {
        self.execute_with_operation_metric::<T, _>(Operation::Get, None, |tx| {
            let dbi = self.get_dbi::<T>()?;
            #[cfg(target_os = "linux")]
            let value = if self.prefetch.load(Ordering::Relaxed) {
                tx.get::<super::value_prefetch::PrefetchValue<'_>>(dbi, key.as_ref())
                    .map(|value| value.map(|value| value.0))
            } else {
                tx.get(dbi, key.as_ref())
            };
            #[cfg(not(target_os = "linux"))]
            let value = tx.get(dbi, key.as_ref());
            value.map_err(|e| DatabaseError::Read(e.into()))?.map(decode_one::<T>).transpose()
        })
    }

    #[instrument(name = "Tx::commit", level = "debug", target = "providers::db", skip_all)]
    fn commit(self) -> Result<(), DatabaseError> {
        self.execute_with_close_transaction_metric(TransactionOutcome::Commit, |this| {
            match this.inner.commit().map_err(|e| DatabaseError::Commit(e.into())) {
                Ok(latency) => (Ok(()), Some(latency)),
                Err(e) => (Err(e), None),
            }
        })
    }

    fn abort(self) {
        self.execute_with_close_transaction_metric(TransactionOutcome::Abort, |this| {
            (drop(this.inner), None)
        })
    }

    // Iterate over read only values in database.
    fn cursor_read<T: Table>(&self) -> Result<Self::Cursor<T>, DatabaseError> {
        self.new_cursor()
    }

    /// Iterate over read only values in database.
    fn cursor_dup_read<T: DupSort>(&self) -> Result<Self::DupCursor<T>, DatabaseError> {
        self.new_cursor()
    }

    /// Returns number of entries in the table using cheap DB stats invocation.
    fn entries<T: Table>(&self) -> Result<usize, DatabaseError> {
        Ok(self
            .inner
            .db_stat_with_dbi(self.get_dbi::<T>()?)
            .map_err(|e| DatabaseError::Stats(e.into()))?
            .entries())
    }

    /// Disables long-lived read transaction safety guarantees, such as backtrace recording and
    /// timeout.
    fn disable_long_read_transaction_safety(&mut self) {
        if let Some(metrics_handler) = self.metrics_handler.as_mut() {
            metrics_handler.record_backtrace = false;
        }

        self.inner.disable_timeout();
    }

    #[cfg(target_os = "linux")]
    fn prefetch(&self, enabled: bool) -> &Self {
        self.prefetch.store(enabled, Ordering::Relaxed);
        self
    }
}

#[derive(Clone, Copy)]
enum PutKind {
    /// Default kind that inserts a new key-value or overwrites an existed key.
    Upsert,
    /// Append the key-value to the end of the table -- fast path when the new
    /// key is the highest so far, like the latest block number.
    Append,
}

impl PutKind {
    const fn into_operation_and_flags(self) -> (Operation, DatabaseWriteOperation, WriteFlags) {
        match self {
            Self::Upsert => {
                (Operation::PutUpsert, DatabaseWriteOperation::PutUpsert, WriteFlags::UPSERT)
            }
            Self::Append => {
                (Operation::PutAppend, DatabaseWriteOperation::PutAppend, WriteFlags::APPEND)
            }
        }
    }
}

impl Tx<RW> {
    /// The inner implementation mapping to `mdbx_put` that supports different
    /// put kinds like upserting and appending.
    fn put<T: Table>(
        &self,
        kind: PutKind,
        key: T::Key,
        value: T::Value,
    ) -> Result<(), DatabaseError> {
        let key = key.encode();
        let value = value.compress();
        let (operation, write_operation, flags) = kind.into_operation_and_flags();
        self.execute_with_operation_metric::<T, _>(operation, Some(value.as_ref().len()), |tx| {
            tx.put(self.get_dbi::<T>()?, key.as_ref(), value, flags).map_err(|e| {
                DatabaseWriteError {
                    info: e.into(),
                    operation: write_operation,
                    table_name: T::NAME,
                    key: key.into_vec(),
                }
                .into()
            })
        })
    }
}

impl DbTxMut for Tx<RW> {
    type CursorMut<T: Table> = Cursor<RW, T>;
    type DupCursorMut<T: DupSort> = Cursor<RW, T>;

    fn put<T: Table>(&self, key: T::Key, value: T::Value) -> Result<(), DatabaseError> {
        self.put::<T>(PutKind::Upsert, key, value)
    }

    fn append<T: Table>(&self, key: T::Key, value: T::Value) -> Result<(), DatabaseError> {
        self.put::<T>(PutKind::Append, key, value)
    }

    fn delete<T: Table>(
        &self,
        key: T::Key,
        value: Option<T::Value>,
    ) -> Result<bool, DatabaseError> {
        let mut data = None;

        let value = value.map(Compress::compress);
        if let Some(value) = &value {
            data = Some(value.as_ref());
        };

        self.execute_with_operation_metric::<T, _>(Operation::Delete, None, |tx| {
            tx.del(self.get_dbi::<T>()?, key.encode(), data)
                .map_err(|e| DatabaseError::Delete(e.into()))
        })
    }

    fn clear<T: Table>(&self) -> Result<(), DatabaseError> {
        self.inner.clear_db(self.get_dbi::<T>()?).map_err(|e| DatabaseError::Delete(e.into()))?;

        Ok(())
    }

    fn cursor_write<T: Table>(&self) -> Result<Self::CursorMut<T>, DatabaseError> {
        self.new_cursor()
    }

    fn cursor_dup_write<T: DupSort>(&self) -> Result<Self::DupCursorMut<T>, DatabaseError> {
        self.new_cursor()
    }
}

#[cfg(test)]
mod tests {
    use crate::{mdbx::DatabaseArguments, tables, DatabaseEnv, DatabaseEnvKind};
    use reth_db_api::{database::Database, models::ClientVersion, transaction::DbTx};
    use reth_libmdbx::MaxReadTransactionDuration;
    use reth_storage_errors::db::DatabaseError;
    use std::{sync::atomic::Ordering, thread::sleep, time::Duration};
    use tempfile::tempdir;

    #[test]
    fn long_read_transaction_safety_disabled() {
        const MAX_DURATION: Duration = Duration::from_secs(1);

        let dir = tempdir().unwrap();
        let args = DatabaseArguments::new(ClientVersion::default())
            .with_max_read_transaction_duration(Some(MaxReadTransactionDuration::Set(
                MAX_DURATION,
            )));
        let db = DatabaseEnv::open(dir.path(), DatabaseEnvKind::RW, args)
            .unwrap()
            .with_metrics()
            .with_disable_long_read_transaction_safety();

        let mut tx = db.tx().unwrap();
        let metrics_handler = tx.metrics_handler.as_mut().unwrap();
        metrics_handler.long_transaction_deadline = metrics_handler.start + MAX_DURATION;
        // Give the `TxnManager` some time to time out the transaction.
        sleep(MAX_DURATION + Duration::from_millis(100));

        // Transaction has not timed out.
        assert!(matches!(
            tx.get::<tables::Transactions>(0).unwrap_err(),
            DatabaseError::Open(err) if err == reth_libmdbx::Error::NotFound.into()));
        // Backtrace is not recorded.
        assert!(!tx.metrics_handler.unwrap().backtrace_recorded.load(Ordering::Relaxed));
    }

    #[test]
    fn long_read_transaction_safety_enabled() {
        const MAX_DURATION: Duration = Duration::from_secs(1);

        let dir = tempdir().unwrap();
        let args = DatabaseArguments::new(ClientVersion::default())
            .with_max_read_transaction_duration(Some(MaxReadTransactionDuration::Set(
                MAX_DURATION,
            )));
        let db = DatabaseEnv::open(dir.path(), DatabaseEnvKind::RW, args).unwrap().with_metrics();

        let mut tx = db.tx().unwrap();
        let metrics_handler = tx.metrics_handler.as_mut().unwrap();
        metrics_handler.long_transaction_deadline = metrics_handler.start + MAX_DURATION;
        // Give the `TxnManager` some time to time out the transaction.
        sleep(MAX_DURATION + Duration::from_millis(100));

        // Transaction has timed out.
        assert!(matches!(
            tx.get::<tables::Transactions>(0).unwrap_err(),
            DatabaseError::Open(err) if err == reth_libmdbx::Error::ReadTransactionTimeout.into()));
        // Backtrace is recorded.
        assert!(tx.metrics_handler.unwrap().backtrace_recorded.load(Ordering::Relaxed));
    }
}