typed_store/

database.rs

// Copyright (c) 2026 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

use std::{
    borrow::Borrow,
    marker::PhantomData,
    ops::{Bound, Deref, RangeBounds},
    path::{Path, PathBuf},
    sync::{Arc, Weak},
    time::Duration,
};

use fastcrypto::hash::{Digest, HashFunction};
use iota_common::debug_fatal;
use iota_macros::{fail_point, nondeterministic};
use prometheus_filtered::{Histogram, HistogramTimer};
use rocksdb::{
    DBPinnableSlice, DBWithThreadMode, Error, LiveFile, MultiThreaded, ReadOptions,
    SnapshotWithThreadMode, WriteBatch,
};
use serde::{Serialize, de::DeserializeOwned};
use tokio::{sync::oneshot, time::Instant};
use tracing::{debug, error, instrument, warn};
use typed_store_error::TypedStoreError;

use crate::{
    DbIterator,
    memstore::{InMemoryBatch, InMemoryDB},
    metrics::{DBMetrics, RocksDBPerfContext, SamplingInterval},
    rocks::{
        RocksDB,
        errors::{typed_store_err_from_bcs_err, typed_store_err_from_rocks_err},
        options::ReadWriteOptions,
        rocks_cf, rocks_util,
        safe_iter::{SafeIter, SafeRevIter},
    },
    traits::{Map, TableSummary},
    util::{
        be_fix_int_ser, iterator_bounds_with_range, prefix_iterator_bounds,
        prefix_iterator_bounds_with_range,
    },
};

#[derive(Clone)]
pub(crate) enum ColumnFamily {
    Rocks(String),
    #[allow(dead_code)]
    InMemory(String),
}

impl std::fmt::Debug for ColumnFamily {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ColumnFamily::Rocks(name) => write!(f, "RocksDB cf: {name}"),
            ColumnFamily::InMemory(name) => write!(f, "InMemory cf: {name}"),
        }
    }
}

impl ColumnFamily {
    pub(crate) fn name(&self) -> &str {
        match self {
            ColumnFamily::Rocks(name) => name,
            ColumnFamily::InMemory(name) => name,
        }
    }
}

pub(crate) enum Storage {
    Rocks(RocksDB),
    #[allow(dead_code)]
    InMemory(InMemoryDB),
}

/// A point-in-time RocksDB snapshot of a database.
///
/// Reads through it ignore every write made after it was taken. It pins a
/// sequence number rather than copying data, and does not survive a restart.
///
/// While it is alive, RocksDB keeps on disk every older version it may still
/// read, so take it as late as possible and drop it as soon as the read is
/// done.
///
/// A `CompactionFilter` ignores snapshots and may drop rows one is still
/// reading, so never install one on a column family read through a snapshot.
pub struct DbSnapshot<'db> {
    snapshot: SnapshotWithThreadMode<'db, DBWithThreadMode<MultiThreaded>>,
}

impl std::fmt::Debug for Storage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Storage::Rocks(db) => write!(f, "RocksDB Storage {db:?}"),
            Storage::InMemory(db) => write!(f, "InMemoryDB Storage {db:?}"),
        }
    }
}

pub(crate) enum GetResult<'a> {
    Rocks(DBPinnableSlice<'a>),
    InMemory(Vec<u8>),
}

impl Deref for GetResult<'_> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        match self {
            GetResult::Rocks(d) => d.deref(),
            GetResult::InMemory(d) => d.deref(),
        }
    }
}

pub enum StorageWriteBatch {
    Rocks(rocksdb::WriteBatch),
    InMemory(InMemoryBatch),
}

#[derive(Debug, Default)]
pub struct MetricConf {
    pub db_name: String,
    pub read_sample_interval: SamplingInterval,
    pub write_sample_interval: SamplingInterval,
    pub iter_sample_interval: SamplingInterval,
}

impl MetricConf {
    pub fn new(db_name: &str) -> Self {
        if db_name.is_empty() {
            error!("A meaningful db name should be used for metrics reporting.")
        }
        Self {
            db_name: db_name.to_string(),
            read_sample_interval: SamplingInterval::default(),
            write_sample_interval: SamplingInterval::default(),
            iter_sample_interval: SamplingInterval::default(),
        }
    }

    pub fn with_sampling(self, read_interval: SamplingInterval) -> Self {
        Self {
            db_name: self.db_name,
            read_sample_interval: read_interval,
            write_sample_interval: SamplingInterval::default(),
            iter_sample_interval: SamplingInterval::default(),
        }
    }
}

const CF_METRICS_REPORT_PERIOD_SECS: u64 = 30;

#[derive(Debug)]
pub struct Database {
    pub(crate) storage: Storage,
    pub(crate) metric_conf: MetricConf,
    /// Declared after `storage`: fields are dropped in declaration order, so
    /// this goes only once RocksDB has closed and released its `LOCK` file.
    closed: Arc<()>,
}

/// What [`wait_for_database_close`] waits on; see [`Database::close_handle`].
pub struct DatabaseClose(Weak<()>);

impl Drop for Database {
    fn drop(&mut self) {
        DBMetrics::get().decrement_num_active_dbs(&self.metric_conf.db_name);
    }
}

impl Database {
    /// A handle for [`wait_for_database_close`]. Take it while the database
    /// is still open, before its last handle is dropped.
    pub fn close_handle(&self) -> DatabaseClose {
        DatabaseClose(Arc::downgrade(&self.closed))
    }

    pub(crate) fn new(storage: Storage, metric_conf: MetricConf) -> Self {
        DBMetrics::get().increment_num_active_dbs(&metric_conf.db_name);
        Self {
            storage,
            metric_conf,
            closed: Arc::new(()),
        }
    }

    pub(crate) fn get<K: AsRef<[u8]>>(
        &self,
        cf: &ColumnFamily,
        key: K,
        readopts: &ReadOptions,
    ) -> Result<Option<GetResult<'_>>, TypedStoreError> {
        match (&self.storage, cf) {
            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => Ok(db
                .underlying
                .get_pinned_cf_opt(&rocks_cf(db, cf.name())?, key, readopts)
                .map_err(typed_store_err_from_rocks_err)?
                .map(GetResult::Rocks)),
            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => {
                Ok(db.get(cf_name, key)?.map(GetResult::InMemory))
            }

            _ => Err(TypedStoreError::RocksDB(
                "typed store invariant violation".to_string(),
            )),
        }
    }

    pub(crate) fn multi_get<I, K>(
        &self,
        cf: &ColumnFamily,
        keys: I,
        readopts: &ReadOptions,
    ) -> Vec<Result<Option<GetResult<'_>>, TypedStoreError>>
    where
        I: IntoIterator<Item = K>,
        K: AsRef<[u8]>,
    {
        match (&self.storage, cf) {
            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => {
                let keys_vec: Vec<K> = keys.into_iter().collect();
                let handle = match rocks_cf(db, cf.name()) {
                    Ok(handle) => handle,
                    // Fail every key, so the result keeps one slot per key.
                    Err(e) => return keys_vec.iter().map(|_| Err(e.clone())).collect(),
                };
                let res = db.underlying.batched_multi_get_cf_opt(
                    &handle,
                    keys_vec.iter(),
                    // sorted_input
                    false,
                    readopts,
                );
                res.into_iter()
                    .map(|r| {
                        r.map_err(typed_store_err_from_rocks_err)
                            .map(|item| item.map(GetResult::Rocks))
                    })
                    .collect()
            }
            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => db
                .multi_get(cf_name, keys)
                .into_iter()
                .map(|r| r.map(|value| value.map(GetResult::InMemory)))
                .collect(),
            _ => unreachable!("typed store invariant violation"),
        }
    }

    pub fn cf_handle(&self, name: &str) -> Option<()> {
        match &self.storage {
            Storage::Rocks(db) => db.underlying.cf_handle(name).map(|_| ()),
            Storage::InMemory(db) => db.has_cf(name).then_some(()),
        }
    }

    /// Creates a new column family at runtime. Fails if a column family with
    /// this name already exists.
    ///
    /// `options` hold until the database is next opened, after which a
    /// column family the caller does not declare is opened with the
    /// crate's options instead. Declare it to keep them.
    #[instrument(level = "debug", skip(self, options), err)]
    pub fn create_cf(&self, name: &str, options: &rocksdb::Options) -> Result<(), TypedStoreError> {
        match &self.storage {
            Storage::Rocks(db) => nondeterministic!(db.underlying.create_cf(name, options))
                .map_err(typed_store_err_from_rocks_err),
            Storage::InMemory(db) => db.create_cf(name),
        }
    }

    /// Maps still holding the column family keep working but report
    /// [`TypedStoreError::UnregisteredColumn`] on every operation.
    #[instrument(level = "debug", skip(self), err)]
    pub fn drop_cf(&self, name: &str) -> Result<(), TypedStoreError> {
        if name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
            // rocksdb refuses this, but not before its wrapper has taken the
            // handle out of its own map, losing it for the rest of the process.
            return Err(TypedStoreError::RocksDB(format!(
                "the {name} column family cannot be dropped"
            )));
        }
        match &self.storage {
            Storage::Rocks(db) => {
                // RocksDB deletes a dropped column family's files once its
                // last handle is released. Holding one across the drop moves
                // that past the drop, which blocks every column family lookup
                // of this database while it runs.
                let handle = db.underlying.cf_handle(name);
                let dropped = nondeterministic!(db.underlying.drop_cf(name))
                    .map_err(typed_store_err_from_rocks_err);
                drop(handle);
                if dropped.is_ok() {
                    DBMetrics::get().remove_column_family(name);
                }
                dropped
            }
            Storage::InMemory(db) => {
                db.drop_cf(name);
                Ok(())
            }
        }
    }

    pub(crate) fn delete_cf<K: AsRef<[u8]>>(
        &self,
        cf: &ColumnFamily,
        key: K,
    ) -> Result<(), TypedStoreError> {
        fail_point!("delete-cf-before");
        let ret = match (&self.storage, cf) {
            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => {
                rocks_cf(db, cf.name()).and_then(|handle| {
                    db.underlying
                        .delete_cf(&handle, key)
                        .map_err(typed_store_err_from_rocks_err)
                })
            }
            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => {
                db.delete(cf_name, key.as_ref())
            }
            _ => Err(TypedStoreError::RocksDB(
                "typed store invariant violation".to_string(),
            )),
        };
        fail_point!("delete-cf-after");
        #[allow(clippy::let_and_return)]
        ret
    }

    pub fn path_for_pruning(&self) -> &Path {
        match &self.storage {
            Storage::Rocks(rocks) => rocks.underlying.path(),
            _ => unimplemented!("method is only supported for rocksdb backend"),
        }
    }

    pub(crate) fn put_cf(
        &self,
        cf: &ColumnFamily,
        key: Vec<u8>,
        value: Vec<u8>,
    ) -> Result<(), TypedStoreError> {
        fail_point!("put-cf-before");
        let ret = match (&self.storage, cf) {
            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => {
                rocks_cf(db, cf.name()).and_then(|handle| {
                    db.underlying
                        .put_cf(&handle, key, value)
                        .map_err(typed_store_err_from_rocks_err)
                })
            }
            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => db.put(cf_name, key, value),
            _ => Err(TypedStoreError::RocksDB(
                "typed store invariant violation".to_string(),
            )),
        };
        fail_point!("put-cf-after");
        #[allow(clippy::let_and_return)]
        ret
    }

    pub(crate) fn key_may_exist_cf<K: AsRef<[u8]>>(
        &self,
        cf: &ColumnFamily,
        key: K,
        readopts: &ReadOptions,
    ) -> bool {
        match &self.storage {
            // [`rocksdb::DBWithThreadMode::key_may_exist_cf`] can have false positives,
            // but no false negatives. We use it to short-circuit the absent case
            Storage::Rocks(rocks) => match rocks_cf(rocks, cf.name()) {
                Ok(handle) => rocks
                    .underlying
                    .key_may_exist_cf_opt(&handle, key, readopts),
                // Not a definitive absence: let the real read report it.
                Err(_) => true,
            },
            _ => true,
        }
    }

    /// Flush the memtable of a single column family to SST files on disk.
    pub(crate) fn flush_cf(&self, cf: &ColumnFamily) -> Result<(), TypedStoreError> {
        match &self.storage {
            // A flush blocks on RocksDB background threads; under the simulator it
            // must run off the test thread, like the opens in `rocks/mod.rs`, or
            // real flush durations leak into the simulated schedule. No-op outside
            // msim.
            Storage::Rocks(rocks) => nondeterministic!({
                let cf_name = cf.name();
                // Deliberately not `rocks_cf`: flushing a column family
                // dropped in the meantime is a no-op, not an error.
                if let Some(handle) = rocks.underlying.cf_handle(cf_name) {
                    rocks.underlying.flush_cf(&handle).map_err(|e| {
                        TypedStoreError::RocksDB(format!(
                            "Failed to flush column family {cf_name}: {e}"
                        ))
                    })?;
                }
                Ok(())
            }),
            // InMemory databases don't need flushing.
            Storage::InMemory(_) => Ok(()),
        }
    }

    /// Flush the memtables of every column family of the database to SST files
    /// on disk.
    pub fn flush_all(&self) -> Result<(), TypedStoreError> {
        match &self.storage {
            // See `flush_cf` for why the flushes run off the test thread under
            // the simulator.
            Storage::Rocks(rocks) => nondeterministic!({
                for cf_name in rocks.cf_names().map_err(|e| {
                    TypedStoreError::RocksDB(format!(
                        "Failed to list column families of {}: {e}",
                        rocks.underlying.path().display()
                    ))
                })? {
                    if let Some(cf) = rocks.underlying.cf_handle(&cf_name) {
                        rocks.underlying.flush_cf(&cf).map_err(|e| {
                            TypedStoreError::RocksDB(format!(
                                "Failed to flush column family {cf_name}: {e}"
                            ))
                        })?;
                    }
                }
                Ok(())
            }),
            // InMemory databases don't need flushing.
            Storage::InMemory(_) => Ok(()),
        }
    }

    pub fn write_opt(
        &self,
        batch: StorageWriteBatch,
        write_options: &rocksdb::WriteOptions,
    ) -> Result<(), TypedStoreError> {
        fail_point!("batch-write-before");
        let ret = match (&self.storage, batch) {
            (Storage::Rocks(rocks), StorageWriteBatch::Rocks(batch)) => rocks
                .underlying
                .write_opt(batch, write_options)
                .map_err(typed_store_err_from_rocks_err),
            (Storage::InMemory(db), StorageWriteBatch::InMemory(batch)) => {
                // InMemory doesn't support write options.
                db.write(batch)
            }
            _ => Err(TypedStoreError::RocksDB(
                "using invalid batch type for the database".to_string(),
            )),
        };
        fail_point!("batch-write-after");

        #[allow(clippy::let_and_return)]
        ret
    }

    pub(crate) fn compact_range_cf<K: AsRef<[u8]>>(
        &self,
        cf: &ColumnFamily,
        start: Option<K>,
        end: Option<K>,
    ) -> Result<(), TypedStoreError> {
        if let Storage::Rocks(rocksdb) = &self.storage {
            let handle = rocks_cf(rocksdb, cf.name())?;
            rocksdb.underlying.compact_range_cf(&handle, start, end);
        }
        Ok(())
    }

    /// How many sorted runs a point lookup in the column family `cf_name` may
    /// probe: one per level-0 file and one per other level holding files. `0`
    /// for a column family that does not exist and on the in-memory backend.
    pub fn sorted_runs(&self, cf_name: &str) -> Result<u64, TypedStoreError> {
        Ok(self
            .files_per_level(cf_name)?
            .iter()
            .enumerate()
            .map(|(level, &files)| {
                if level == 0 {
                    files
                } else {
                    u64::from(files > 0)
                }
            })
            .sum())
    }

    /// How many SST files the column family `cf_name` has. `0` for a column
    /// family that does not exist and on the in-memory backend.
    pub fn sst_file_count(&self, cf_name: &str) -> Result<u64, TypedStoreError> {
        Ok(self.files_per_level(cf_name)?.iter().sum())
    }

    /// The number of SST files on each level of the column family `cf_name`.
    fn files_per_level(&self, cf_name: &str) -> Result<Vec<u64>, TypedStoreError> {
        let Storage::Rocks(rocks) = &self.storage else {
            return Ok(Vec::new());
        };
        let Some(handle) = rocks.underlying.cf_handle(cf_name) else {
            return Ok(Vec::new());
        };
        let mut levels = Vec::new();
        // The property is absent for the levels past the column family's
        // last one.
        for level in 0.. {
            let Some(files) = rocks
                .underlying
                .property_value_cf(&handle, &format!("rocksdb.num-files-at-level{level}"))
                .map_err(typed_store_err_from_rocks_err)?
            else {
                break;
            };
            levels.push(files.trim().parse().map_err(|e| {
                TypedStoreError::RocksDB(format!(
                    "the file count of level {level} of {cf_name} is not a number: {e}"
                ))
            })?);
        }
        Ok(levels)
    }

    /// The size of the SST files of the column family `cf_name`, which a
    /// compaction of all of them needs again in free disk space while it runs.
    /// `0` for a column family that does not exist and on the in-memory
    /// backend.
    pub fn sst_files_size(&self, cf_name: &str) -> Result<u64, TypedStoreError> {
        let Storage::Rocks(rocks) = &self.storage else {
            return Ok(0);
        };
        let Some(handle) = rocks.underlying.cf_handle(cf_name) else {
            return Ok(0);
        };
        Ok(rocks
            .underlying
            .property_int_value_cf(&handle, rocksdb::properties::TOTAL_SST_FILES_SIZE)
            .map_err(typed_store_err_from_rocks_err)?
            .unwrap_or(0))
    }

    /// Compacts the whole column family `cf_name` into a single sorted run,
    /// flushing its memtable first. Blocks until done. Other column families
    /// keep compacting meanwhile. Does nothing on the in-memory backend.
    pub fn compact_cf(&self, cf_name: &str) -> Result<(), TypedStoreError> {
        if let Storage::Rocks(rocks) = &self.storage {
            let handle = rocks_cf(rocks, cf_name)?;
            let mut options = rocksdb::CompactOptions::default();
            options.set_exclusive_manual_compaction(false);
            options.set_bottommost_level_compaction(rocksdb::BottommostLevelCompaction::Force);
            nondeterministic!(rocks.underlying.compact_range_cf_opt(
                &handle,
                None::<&[u8]>,
                None::<&[u8]>,
                &options
            ));
        }
        Ok(())
    }

    /// A point-in-time snapshot of this database. See [`DbSnapshot`] for what
    /// holding one costs.
    ///
    /// # Panics
    ///
    /// Panics on the in-memory backend, which has no point-in-time reads.
    pub fn snapshot(&self) -> DbSnapshot<'_> {
        match &self.storage {
            Storage::Rocks(rocks) => DbSnapshot {
                snapshot: rocks.underlying.snapshot(),
            },
            Storage::InMemory(_) => unimplemented!("method is only supported for rocksdb backend"),
        }
    }

    pub fn get_sampling_interval(&self) -> SamplingInterval {
        self.metric_conf.read_sample_interval.new_from_self()
    }

    pub fn multiget_sampling_interval(&self) -> SamplingInterval {
        self.metric_conf.read_sample_interval.new_from_self()
    }

    pub fn write_sampling_interval(&self) -> SamplingInterval {
        self.metric_conf.write_sample_interval.new_from_self()
    }

    pub fn iter_sampling_interval(&self) -> SamplingInterval {
        self.metric_conf.iter_sample_interval.new_from_self()
    }

    pub(crate) fn db_name(&self) -> String {
        let name = &self.metric_conf.db_name;
        if name.is_empty() {
            "default".to_string()
        } else {
            name.clone()
        }
    }

    pub fn live_files(&self) -> Result<Vec<LiveFile>, Error> {
        match &self.storage {
            Storage::Rocks(rocks) => rocks.underlying.live_files(),
            _ => Ok(vec![]),
        }
    }

    pub(crate) fn try_catch_up_with_primary(&self) -> Result<(), TypedStoreError> {
        if let Storage::Rocks(rocks) = &self.storage {
            rocks
                .underlying
                .try_catch_up_with_primary()
                .map_err(typed_store_err_from_rocks_err)?;
        }
        Ok(())
    }
}

fn rocks_cf_from_db<'a>(
    db: &'a Database,
    cf_name: &str,
) -> Result<Arc<rocksdb::BoundColumnFamily<'a>>, TypedStoreError> {
    match &db.storage {
        Storage::Rocks(rocksdb) => rocks_cf(rocksdb, cf_name),
        _ => Err(TypedStoreError::RocksDB(
            "using invalid batch type for the database".to_string(),
        )),
    }
}

/// An interface to a rocksDB database, keyed by a columnfamily
#[derive(Clone, Debug)]
pub struct DBMap<K, V> {
    pub db: Arc<Database>,
    _phantom: PhantomData<fn(K) -> V>,
    column_family: ColumnFamily,
    pub opts: ReadWriteOptions,
    db_metrics: Arc<DBMetrics>,
    get_sample_interval: SamplingInterval,
    multiget_sample_interval: SamplingInterval,
    write_sample_interval: SamplingInterval,
    iter_sample_interval: SamplingInterval,
    _metrics_task_cancel_handle: Arc<oneshot::Sender<()>>,
    /// What the per-operation metrics are labelled with; the column family
    /// name unless [`Self::with_metrics_label`] set another.
    metrics_label: Arc<str>,
}

impl<K, V> DBMap<K, V> {
    pub(crate) fn new(
        db: Arc<Database>,
        opts: &ReadWriteOptions,
        column_family: ColumnFamily,
        skip_metrics_reporting: bool,
    ) -> Self {
        let db_cloned = Arc::downgrade(&db);
        let db_metrics = DBMetrics::get();
        let db_metrics_cloned = db_metrics.clone();
        let cf = column_family.name().to_string();

        let (sender, mut recv) = tokio::sync::oneshot::channel();
        if !skip_metrics_reporting && matches!(db.storage, Storage::Rocks(_)) {
            tokio::task::spawn(async move {
                let mut interval =
                    tokio::time::interval(Duration::from_secs(CF_METRICS_REPORT_PERIOD_SECS));
                loop {
                    tokio::select! {
                        _ = interval.tick() => {
                            if let Some(db) = db_cloned.upgrade() {
                                let cf = cf.clone();
                                let db_metrics = db_metrics.clone();
                                if let Err(e) = tokio::task::spawn_blocking(move || {
                                    Self::report_rocksdb_metrics(&db, &cf, &db_metrics);
                                }).await {
                                    error!("Failed to log metrics with error: {}", e);
                                }
                            } else {
                                break;
                            }
                        }
                        _ = &mut recv => break,
                    }
                }
                debug!("Returning the cf metric logging task for DBMap: {}", &cf);
            });
        }
        let metrics_label = Arc::from(column_family.name());
        DBMap {
            db: db.clone(),
            opts: opts.clone(),
            _phantom: PhantomData,
            column_family,
            metrics_label,
            db_metrics: db_metrics_cloned,
            _metrics_task_cancel_handle: Arc::new(sender),
            get_sample_interval: db.get_sampling_interval(),
            multiget_sample_interval: db.multiget_sampling_interval(),
            write_sample_interval: db.write_sampling_interval(),
            iter_sample_interval: db.iter_sampling_interval(),
        }
    }

    /// Reopens an open database as a typed map operating under a specific
    /// column family. if no column family is passed, the default column
    /// family is used.
    ///
    /// When `skip_metrics_reporting` is true, no periodic per-column-family
    /// metrics task is spawned; use this for deprecated tables and for large
    /// sets of rarely-touched column families.
    #[instrument(level = "debug", skip(db), err)]
    pub fn reopen(
        db: &Arc<Database>,
        opt_cf: Option<&str>,
        rw_options: &ReadWriteOptions,
        skip_metrics_reporting: bool,
    ) -> Result<Self, TypedStoreError> {
        let cf_key = opt_cf
            .unwrap_or(rocksdb::DEFAULT_COLUMN_FAMILY_NAME)
            .to_owned();
        if db.cf_handle(&cf_key).is_none() {
            return Err(TypedStoreError::UnregisteredColumn(cf_key));
        }

        let column_family = match &db.storage {
            Storage::Rocks(_) => ColumnFamily::Rocks(cf_key),
            Storage::InMemory(_) => ColumnFamily::InMemory(cf_key),
        };
        Ok(DBMap::new(
            db.clone(),
            rw_options,
            column_family,
            skip_metrics_reporting,
        ))
    }

    /// Labels the per-operation metrics with `label` instead of the column
    /// family name, so that the column families of one store, such as one per
    /// epoch, share their series.
    pub fn with_metrics_label(mut self, label: &str) -> Self {
        self.metrics_label = Arc::from(label);
        self
    }

    /// What the per-operation metrics are labelled with.
    pub fn metrics_label(&self) -> &str {
        &self.metrics_label
    }

    pub fn cf_name(&self) -> &str {
        self.column_family.name()
    }

    pub fn batch(&self) -> DBBatch {
        let batch = match &self.db.storage {
            Storage::Rocks(_) => StorageWriteBatch::Rocks(WriteBatch::default()),
            Storage::InMemory(_) => StorageWriteBatch::InMemory(InMemoryBatch::default()),
        };
        DBBatch::new(
            &self.db,
            batch,
            &self.db_metrics,
            &self.write_sample_interval,
        )
    }

    /// Moves the SST files at `paths`, written by
    /// [`SstFileBuffer::write_file`](crate::rocks::SstFileBuffer::write_file)
    /// with this column family's options, into this table. The files may
    /// overlap each other; a key in several of them keeps its value from the
    /// last.
    ///
    /// # Panics
    ///
    /// Panics on the in-memory backend.
    pub fn ingest_sst_files(&self, paths: &[PathBuf]) -> Result<(), TypedStoreError> {
        if paths.is_empty() {
            return Ok(());
        }
        match &self.db.storage {
            // Ingestion can wait on a memtable flush; see `Database::flush_cf`
            // for why that runs off the test thread under the simulator.
            Storage::Rocks(rocks) => nondeterministic!({
                let handle = rocks_cf(rocks, self.cf_name())?;
                let mut options = rocksdb::IngestExternalFileOptions::default();
                options.set_move_files(true);
                rocks
                    .underlying
                    .ingest_external_file_cf_opts(&handle, &options, paths.to_vec())
                    .map_err(typed_store_err_from_rocks_err)
            }),
            Storage::InMemory(_) => unimplemented!("method is only supported for rocksdb backend"),
        }
    }

    /// Flush the memtable of this table's column family to SST files on disk.
    pub fn flush(&self) -> Result<(), TypedStoreError> {
        self.db.flush_cf(&self.column_family)
    }

    /// Flush the memtables of every column family of the database to SST files
    /// on disk.
    pub fn flush_all(&self) -> Result<(), TypedStoreError> {
        self.db.flush_all()
    }

    pub fn compact_range<J: Serialize>(&self, start: &J, end: &J) -> Result<(), TypedStoreError> {
        let from_buf = be_fix_int_ser(start);
        let to_buf = be_fix_int_ser(end);
        self.db
            .compact_range_cf(&self.column_family, Some(from_buf), Some(to_buf))
    }

    pub fn compact_range_raw(
        &self,
        cf_name: &str,
        start: Vec<u8>,
        end: Vec<u8>,
    ) -> Result<(), TypedStoreError> {
        let cf = match &self.db.storage {
            Storage::Rocks(_) => ColumnFamily::Rocks(cf_name.to_string()),
            Storage::InMemory(_) => ColumnFamily::InMemory(cf_name.to_string()),
        };
        self.db.compact_range_cf(&cf, Some(start), Some(end))
    }

    /// Returns a vector of raw values corresponding to the keys provided.
    fn multi_get_pinned(
        &self,
        keys_bytes: impl IntoIterator<Item = Vec<u8>>,
    ) -> Result<Vec<Option<GetResult<'_>>>, TypedStoreError> {
        let _timer = self
            .db_metrics
            .op_metrics
            .rocksdb_multiget_latency_seconds
            .with_label_values(&[self.metrics_label()])
            .start_timer();
        let perf_ctx = if self.multiget_sample_interval.sample() {
            Some(RocksDBPerfContext)
        } else {
            None
        };
        let results: Result<Vec<_>, TypedStoreError> = self
            .db
            .multi_get(&self.column_family, keys_bytes, &self.opts.readopts())
            .into_iter()
            .collect();
        let entries = results?;
        let entry_size = entries
            .iter()
            .flatten()
            .map(|entry| entry.len())
            .sum::<usize>();
        self.db_metrics
            .op_metrics
            .rocksdb_multiget_bytes
            .with_label_values(&[self.metrics_label()])
            .observe(entry_size as f64);
        if perf_ctx.is_some() {
            self.db_metrics
                .read_perf_ctx_metrics
                .report_metrics(self.metrics_label());
        }
        Ok(entries)
    }

    pub fn table_summary(&self) -> eyre::Result<TableSummary>
    where
        K: Serialize + DeserializeOwned,
        V: Serialize + DeserializeOwned,
    {
        let mut num_keys = 0;
        let mut key_bytes_total = 0;
        let mut value_bytes_total = 0;
        let mut key_hist = hdrhistogram::Histogram::<u64>::new_with_max(100000, 2).unwrap();
        let mut value_hist = hdrhistogram::Histogram::<u64>::new_with_max(100000, 2).unwrap();
        for item in self.safe_iter() {
            let (key, value) = item?;
            num_keys += 1;
            let key_len = be_fix_int_ser(key.borrow()).len();
            let value_len = bcs::to_bytes(value.borrow())?.len();
            key_bytes_total += key_len;
            value_bytes_total += value_len;
            key_hist.record(key_len as u64)?;
            value_hist.record(value_len as u64)?;
        }
        Ok(TableSummary {
            num_keys,
            key_bytes_total,
            value_bytes_total,
            key_hist,
            value_hist,
        })
    }

    // Creates metrics and context for tracking an iterator usage and performance.
    fn create_iter_context(
        &self,
    ) -> (
        Option<HistogramTimer>,
        Option<Histogram>,
        Option<Histogram>,
        Option<RocksDBPerfContext>,
    ) {
        let timer = self
            .db_metrics
            .op_metrics
            .rocksdb_iter_latency_seconds
            .with_label_values(&[self.metrics_label()])
            .start_timer();
        let bytes_scanned = self
            .db_metrics
            .op_metrics
            .rocksdb_iter_bytes
            .with_label_values(&[self.metrics_label()]);
        let keys_scanned = self
            .db_metrics
            .op_metrics
            .rocksdb_iter_keys
            .with_label_values(&[self.metrics_label()]);
        let perf_ctx = if self.iter_sample_interval.sample() {
            Some(RocksDBPerfContext)
        } else {
            None
        };
        (
            Some(timer),
            Some(bytes_scanned),
            Some(keys_scanned),
            perf_ctx,
        )
    }

    /// Shared rocksdb `SafeIter` construction (raw iterator + scan metrics)
    /// used by both the forward and reverse raw iterators.
    fn rocks_safe_iter<'a>(
        &self,
        db: &'a RocksDB,
        readopts: ReadOptions,
    ) -> Result<SafeIter<'a, K, V>, TypedStoreError>
    where
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        let db_iter = db
            .underlying
            .raw_iterator_cf_opt(&rocks_cf(db, self.column_family.name())?, readopts);
        let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
        Ok(SafeIter::new(
            self.metrics_label().to_string(),
            db_iter,
            _timer,
            _perf_ctx,
            bytes_scanned,
            keys_scanned,
            Some(self.db_metrics.clone()),
        ))
    }

    /// Iterates the whole column family as of `db_snapshot` instead of the
    /// current state of the database.
    ///
    /// The blocks it reads are not added to the block cache.
    pub fn safe_iter_at_snapshot<'a>(
        &'a self,
        db_snapshot: &'a DbSnapshot<'a>,
    ) -> DbIterator<'a, (K, V)>
    where
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        match &self.db.storage {
            Storage::Rocks(db) => {
                let mut readopts = self.opts.readopts();
                readopts.fill_cache(false);
                readopts.set_snapshot(&db_snapshot.snapshot);
                match self.rocks_safe_iter(db, readopts) {
                    Ok(iter) => Box::new(iter),
                    Err(e) => error_iterator(e),
                }
            }
            Storage::InMemory(_) => {
                unreachable!("a DbSnapshot is only created over the rocksdb backend")
            }
        }
    }

    /// Forward iterator over the raw byte bounds `[lower_bound, upper_bound)`;
    /// both bounds (already serialized, `upper_bound` exclusive) are applied to
    /// the read options.
    fn iter_forward_raw(
        &self,
        lower_bound: Option<Vec<u8>>,
        upper_bound: Option<Vec<u8>>,
    ) -> DbIterator<'_, (K, V)>
    where
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        match &self.db.storage {
            Storage::Rocks(db) => {
                let readopts =
                    rocks_util::apply_range_bounds(self.opts.readopts(), lower_bound, upper_bound);
                match self.rocks_safe_iter(db, readopts) {
                    Ok(iter) => Box::new(iter),
                    Err(e) => error_iterator(e),
                }
            }
            Storage::InMemory(db) => {
                db.iterator(self.column_family.name(), lower_bound, upper_bound, false)
            }
        }
    }

    /// Reverse counterpart of [`Self::iter_forward_raw`], yielding the same
    /// keys in descending order. Only the lower bound is applied to the
    /// read options; the (exclusive) `upper_bound` is enforced by
    /// [`SafeRevIter`]'s initial seek, because setting an
    /// iterate-upper-bound would stop `seek_for_prev` from landing on the
    /// boundary key.
    fn iter_reversed_raw(
        &self,
        lower_bound: Option<Vec<u8>>,
        upper_bound: Option<Vec<u8>>,
    ) -> DbIterator<'_, (K, V)>
    where
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        match &self.db.storage {
            Storage::Rocks(db) => {
                let readopts =
                    rocks_util::apply_range_bounds(self.opts.readopts(), lower_bound, None);
                match self.rocks_safe_iter(db, readopts) {
                    Ok(iter) => Box::new(SafeRevIter::new(iter, upper_bound)),
                    Err(e) => error_iterator(e),
                }
            }
            Storage::InMemory(db) => {
                db.iterator(self.column_family.name(), lower_bound, upper_bound, true)
            }
        }
    }

    /// Forward iterator over every entry whose key begins with `prefix`.
    ///
    /// `prefix` is serialized with `be_fix_int_ser` and must form a prefix of
    /// the column family's key encoding — typically the leading field(s) of a
    /// tuple key. This avoids constructing artificial maximum keys to bound a
    /// composite-key scan.
    pub fn safe_iter_with_prefix<P>(&self, prefix: &P) -> DbIterator<'_, (K, V)>
    where
        P: ?Sized + Serialize,
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        let (lower_bound, upper_bound) = prefix_iterator_bounds(prefix);
        self.iter_forward_raw(lower_bound, upper_bound)
    }

    /// Forward iterator over entries whose key begins with `prefix`, starting
    /// at the position `lower` gives within the prefix rather than at the
    /// first key under the prefix.
    ///
    /// The bound value is **not** a full key: it is the entire remainder of
    /// the key that follows `prefix` (all the trailing fields of a tuple key,
    /// not a partial suffix). Both `prefix` and the bound value are
    /// serialized with `be_fix_int_ser` and concatenated, mirroring how a
    /// composite key is encoded. `Included` starts the scan at
    /// `prefix ++ lower`, `Excluded` strictly after that position (also when
    /// no such key exists), and `Unbounded` at the first key under the
    /// prefix. The scan bounds never leave the prefix (excluding the
    /// maximum remainder yields an empty scan, not the next prefix's keys),
    /// so a paginated scan can resume from a cursor without an artificial
    /// maximum key for the upper bound.
    pub fn safe_iter_with_prefix_from<P, C>(
        &self,
        prefix: &P,
        lower: Bound<&C>,
    ) -> DbIterator<'_, (K, V)>
    where
        P: ?Sized + Serialize,
        C: ?Sized + Serialize,
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        let (lower_bound, upper_bound) =
            prefix_iterator_bounds_with_range::<P, C>(prefix, (lower, Bound::Unbounded));
        self.iter_forward_raw(lower_bound, upper_bound)
    }

    /// Forward iterator over entries whose key begins with `prefix`, resuming
    /// strictly after `cursor`, or from the first key under the prefix when
    /// there is no cursor. `cursor` is the remainder of the key after
    /// `prefix`, as for [`Self::safe_iter_with_prefix_from`].
    pub fn safe_iter_with_prefix_after<P, C>(
        &self,
        prefix: &P,
        cursor: Option<&C>,
    ) -> DbIterator<'_, (K, V)>
    where
        P: ?Sized + Serialize,
        C: ?Sized + Serialize,
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        self.safe_iter_with_prefix_from(prefix, cursor.map_or(Bound::Unbounded, Bound::Excluded))
    }

    /// Reverse counterpart of [`Self::safe_iter_with_prefix`]: the matching
    /// entries in descending key order (e.g. for "latest entry under a
    /// prefix").
    pub fn safe_iter_with_prefix_reversed<P>(&self, prefix: &P) -> DbIterator<'_, (K, V)>
    where
        P: ?Sized + Serialize,
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        let (lower_bound, upper_bound) = prefix_iterator_bounds(prefix);
        self.iter_reversed_raw(lower_bound, upper_bound)
    }
}

/// Provides a mutable struct to form a collection of database write operations,
/// and execute them.
///
/// Batching write and delete operations is faster than performing them one by
/// one and ensures their atomicity,  ie. they are all written or none is.
/// This is also true of operations across column families in the same database.
///
/// Serializations / Deserialization, and naming of column families is performed
/// by passing a DBMap<K,V> with each operation.
///
/// ```
/// use core::fmt::Error;
/// use std::sync::Arc;
///
/// use prometheus_filtered::Registry;
/// use tempfile::tempdir;
/// use typed_store::{Map, metrics::DBMetrics, rocks::*};
///
/// #[tokio::main(flavor = "current_thread")]
/// async fn main() -> Result<(), Error> {
///     let rocks = open_cf_opts(
///         tempfile::tempdir().unwrap(),
///         None,
///         MetricConf::default(),
///         &[
///             ("First_CF", rocksdb::Options::default()),
///             ("Second_CF", rocksdb::Options::default()),
///         ],
///     )
///     .unwrap();
///
///     let db_cf_1 = DBMap::reopen(
///         &rocks,
///         Some("First_CF"),
///         &ReadWriteOptions::default(),
///         false,
///     )
///     .expect("Failed to open storage");
///     let keys_vals_1 = (1..100).map(|i| (i, i.to_string()));
///
///     let db_cf_2 = DBMap::reopen(
///         &rocks,
///         Some("Second_CF"),
///         &ReadWriteOptions::default(),
///         false,
///     )
///     .expect("Failed to open storage");
///     let keys_vals_2 = (1000..1100).map(|i| (i, i.to_string()));
///
///     let mut batch = db_cf_1.batch();
///     batch
///         .insert_batch(&db_cf_1, keys_vals_1.clone())
///         .expect("Failed to batch insert")
///         .insert_batch(&db_cf_2, keys_vals_2.clone())
///         .expect("Failed to batch insert");
///
///     let _ = batch.write().expect("Failed to execute batch");
///     for (k, v) in keys_vals_1 {
///         let val = db_cf_1.get(&k).expect("Failed to get inserted key");
///         assert_eq!(Some(v), val);
///     }
///
///     for (k, v) in keys_vals_2 {
///         let val = db_cf_2.get(&k).expect("Failed to get inserted key");
///         assert_eq!(Some(v), val);
///     }
///     Ok(())
/// }
/// ```
pub struct DBBatch {
    database: Arc<Database>,
    batch: StorageWriteBatch,
    db_metrics: Arc<DBMetrics>,
    write_sample_interval: SamplingInterval,
}

impl DBBatch {
    /// Create a new batch associated with a DB reference.
    ///
    /// Use `open_cf` to get the DB reference or an existing open database.
    pub fn new(
        dbref: &Arc<Database>,
        batch: StorageWriteBatch,
        db_metrics: &Arc<DBMetrics>,
        write_sample_interval: &SamplingInterval,
    ) -> Self {
        DBBatch {
            database: dbref.clone(),
            batch,
            db_metrics: db_metrics.clone(),
            write_sample_interval: write_sample_interval.clone(),
        }
    }

    /// Consume the batch and write its operations to the database
    #[instrument(level = "trace", skip_all, err)]
    pub fn write(self) -> Result<(), TypedStoreError> {
        self.write_opt(&rocksdb::WriteOptions::default())
    }

    /// Consume the batch and write its operations to the database with custom
    /// write options.
    ///
    /// With default options, an entry whose column family was dropped
    /// between staging and this write fails the whole write with a fatal
    /// RocksDB background error. Stores that drop column families at
    /// runtime must pass [`drop_tolerant_write_options`].
    #[instrument(level = "trace", skip_all, err)]
    pub fn write_opt(self, write_options: &rocksdb::WriteOptions) -> Result<(), TypedStoreError> {
        let db_name = self.database.db_name();
        let timer = self
            .db_metrics
            .op_metrics
            .rocksdb_batch_commit_latency_seconds
            .with_label_values(&[&db_name])
            .start_timer();
        let batch_size_bytes = self.size_in_bytes();

        let perf_ctx = if self.write_sample_interval.sample() {
            Some(RocksDBPerfContext)
        } else {
            None
        };
        self.database.write_opt(self.batch, write_options)?;
        self.db_metrics
            .op_metrics
            .rocksdb_batch_commit_bytes
            .with_label_values(&[&db_name])
            .observe(batch_size_bytes as f64);

        if perf_ctx.is_some() {
            self.db_metrics
                .write_perf_ctx_metrics
                .report_metrics(&db_name);
        }
        let elapsed_secs = timer.stop_and_record();
        let threshold_secs = very_slow_batch_write_threshold_secs(batch_size_bytes);
        if elapsed_secs > threshold_secs {
            warn!(
                elapsed_secs,
                threshold_secs,
                batch_size_bytes,
                ?db_name,
                "very slow batch write"
            );
            self.db_metrics
                .op_metrics
                .rocksdb_very_slow_batch_writes_count
                .with_label_values(&[&db_name])
                .inc();
            self.db_metrics
                .op_metrics
                .rocksdb_very_slow_batch_writes_duration_ms
                .with_label_values(&[&db_name])
                .inc_by((elapsed_secs * 1000.0) as u64);
        }
        Ok(())
    }

    pub fn size_in_bytes(&self) -> usize {
        match self.batch {
            StorageWriteBatch::Rocks(ref b) => b.size_in_bytes(),
            StorageWriteBatch::InMemory(_) => 0,
        }
    }

    pub fn delete_batch<J: Borrow<K>, K: Serialize, V>(
        &mut self,
        db: &DBMap<K, V>,
        purged_vals: impl IntoIterator<Item = J>,
    ) -> Result<(), TypedStoreError> {
        self.delete_batch_raw_keys(
            db,
            purged_vals
                .into_iter()
                .map(|key| be_fix_int_ser(key.borrow())),
        )
    }

    /// [`Self::delete_batch`] for keys already in the column family's
    /// `be_fix_int_ser` encoding.
    fn delete_batch_raw_keys<K, V>(
        &mut self,
        db: &DBMap<K, V>,
        purged_vals: impl IntoIterator<Item = Vec<u8>>,
    ) -> Result<(), TypedStoreError> {
        if !Arc::ptr_eq(&db.db, &self.database) {
            return Err(TypedStoreError::CrossDBBatch);
        }

        purged_vals
            .into_iter()
            .try_for_each::<_, Result<_, TypedStoreError>>(|k_buf| {
                match (&mut self.batch, &db.column_family) {
                    (StorageWriteBatch::Rocks(b), ColumnFamily::Rocks(name)) => {
                        b.delete_cf(&rocks_cf_from_db(&self.database, name)?, k_buf)
                    }
                    (StorageWriteBatch::InMemory(b), ColumnFamily::InMemory(name)) => {
                        b.delete_cf(name, k_buf)
                    }
                    _ => Err(TypedStoreError::RocksDB(
                        "typed store invariant violation".to_string(),
                    ))?,
                }
                Ok(())
            })?;
        Ok(())
    }

    /// Deletes a range of keys between `from` (inclusive) and `to`
    /// (non-inclusive) by writing a range delete tombstone in the db map.
    /// The effect of this write is visible immediately, i.e. you won't see
    /// old values when you do a lookup or scan.
    pub fn schedule_delete_range<K: Serialize, V>(
        &mut self,
        db: &DBMap<K, V>,
        from: &K,
        to: &K,
    ) -> Result<(), TypedStoreError> {
        self.schedule_delete_range_raw(db, be_fix_int_ser(from), be_fix_int_ser(to))
    }

    /// [`Self::schedule_delete_range`] for a range whose bounds are not keys,
    /// such as a whole key prefix. The bounds are raw bytes ordered like the
    /// map's keys, i.e. in the `be_fix_int_ser` encoding; `from` is inclusive
    /// and `to` exclusive.
    fn schedule_delete_range_raw<K, V>(
        &mut self,
        db: &DBMap<K, V>,
        from: Vec<u8>,
        to: Vec<u8>,
    ) -> Result<(), TypedStoreError> {
        if !Arc::ptr_eq(&db.db, &self.database) {
            return Err(TypedStoreError::CrossDBBatch);
        }

        if let StorageWriteBatch::Rocks(b) = &mut self.batch {
            b.delete_range_cf(&rocks_cf_from_db(&self.database, db.cf_name())?, from, to);
        }
        Ok(())
    }

    /// inserts a range of (key, value) pairs given as an iterator
    pub fn insert_batch<J: Borrow<K>, K: Serialize, U: Borrow<V>, V: Serialize>(
        &mut self,
        db: &DBMap<K, V>,
        new_vals: impl IntoIterator<Item = (J, U)>,
    ) -> Result<&mut Self, TypedStoreError> {
        self.insert_batch_raw_keys(
            db,
            new_vals
                .into_iter()
                .map(|(key, value)| (be_fix_int_ser(key.borrow()), value)),
        )
    }

    /// [`Self::insert_batch`] for keys already in the column family's
    /// `be_fix_int_ser` encoding.
    fn insert_batch_raw_keys<K, U: Borrow<V>, V: Serialize>(
        &mut self,
        db: &DBMap<K, V>,
        new_vals: impl IntoIterator<Item = (Vec<u8>, U)>,
    ) -> Result<&mut Self, TypedStoreError> {
        if !Arc::ptr_eq(&db.db, &self.database) {
            return Err(TypedStoreError::CrossDBBatch);
        }
        let mut total = 0usize;
        new_vals
            .into_iter()
            .try_for_each::<_, Result<_, TypedStoreError>>(|(k_buf, v)| {
                let v_buf = bcs::to_bytes(v.borrow()).map_err(typed_store_err_from_bcs_err)?;
                total += k_buf.len() + v_buf.len();
                if db.opts.log_value_hash {
                    let key_hash = default_hash(&k_buf);
                    let value_hash = default_hash(&v_buf);
                    debug!(
                        "Insert to DB table: {:?}, key_hash: {:?}, value_hash: {:?}",
                        db.cf_name(),
                        key_hash,
                        value_hash
                    );
                }
                match (&mut self.batch, &db.column_family) {
                    (StorageWriteBatch::Rocks(b), ColumnFamily::Rocks(name)) => {
                        b.put_cf(&rocks_cf_from_db(&self.database, name)?, k_buf, v_buf)
                    }
                    (StorageWriteBatch::InMemory(b), ColumnFamily::InMemory(name)) => {
                        b.put_cf(name, k_buf, v_buf)
                    }
                    _ => Err(TypedStoreError::RocksDB(
                        "typed store invariant violation".to_string(),
                    ))?,
                }
                Ok(())
            })?;
        self.db_metrics
            .op_metrics
            .rocksdb_batch_put_bytes
            .with_label_values(&[db.metrics_label()])
            .observe(total as f64);
        Ok(self)
    }

    /// [`Self::insert_batch`] for a map sharing its column family by tag.
    pub fn insert_batch_tagged<J: Borrow<K>, K: Serialize, U: Borrow<V>, V: Serialize>(
        &mut self,
        map: &TaggedDBMap<K, V>,
        new_vals: impl IntoIterator<Item = (J, U)>,
    ) -> Result<&mut Self, TypedStoreError> {
        self.insert_batch_raw_keys(
            &map.map,
            new_vals
                .into_iter()
                .map(|(key, value)| (be_fix_int_ser(&(map.tag, key.borrow())), value)),
        )
    }

    /// [`Self::delete_batch`] for a map sharing its column family by tag.
    pub fn delete_batch_tagged<J: Borrow<K>, K: Serialize, V>(
        &mut self,
        map: &TaggedDBMap<K, V>,
        purged_vals: impl IntoIterator<Item = J>,
    ) -> Result<(), TypedStoreError> {
        self.delete_batch_raw_keys(
            &map.map,
            purged_vals
                .into_iter()
                .map(|key| be_fix_int_ser(&(map.tag, key.borrow()))),
        )
    }
}

impl<K, V> DBMap<K, V> {
    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn contains_raw_key(&self, key_buf: Vec<u8>) -> Result<bool, TypedStoreError> {
        let readopts = self.opts.readopts();
        Ok(self
            .db
            .key_may_exist_cf(&self.column_family, &key_buf, &readopts)
            && self
                .db
                .get(&self.column_family, &key_buf, &readopts)?
                .is_some())
    }

    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn multi_contains_raw_keys(
        &self,
        keys: impl IntoIterator<Item = Vec<u8>>,
    ) -> Result<Vec<bool>, TypedStoreError> {
        let values = self.multi_get_pinned(keys)?;
        Ok(values.into_iter().map(|v| v.is_some()).collect())
    }

    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn get_raw_key(&self, key_buf: Vec<u8>) -> Result<Option<V>, TypedStoreError>
    where
        V: DeserializeOwned,
    {
        let _timer = self
            .db_metrics
            .op_metrics
            .rocksdb_get_latency_seconds
            .with_label_values(&[self.metrics_label()])
            .start_timer();
        let perf_ctx = if self.get_sample_interval.sample() {
            Some(RocksDBPerfContext)
        } else {
            None
        };
        let res = self
            .db
            .get(&self.column_family, &key_buf, &self.opts.readopts())?;
        self.db_metrics
            .op_metrics
            .rocksdb_get_bytes
            .with_label_values(&[self.metrics_label()])
            .observe(res.as_ref().map_or(0.0, |v| v.len() as f64));
        if perf_ctx.is_some() {
            self.db_metrics
                .read_perf_ctx_metrics
                .report_metrics(self.metrics_label());
        }
        match res {
            Some(data) => {
                let value = bcs::from_bytes(&data).map_err(typed_store_err_from_bcs_err);
                if value.is_err() {
                    let key_hash = default_hash(&key_buf);
                    let value_hash = default_hash(&data);
                    debug_fatal!(
                        "Failed to deserialize value from DB table {:?}, key_hash: {:?}, value_hash: {:?}, error: {:?}",
                        self.cf_name(),
                        key_hash,
                        value_hash,
                        value.as_ref().err().unwrap()
                    );
                }
                Ok(Some(value?))
            }
            None => Ok(None),
        }
    }

    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn insert_raw_key(&self, key_buf: Vec<u8>, value: &V) -> Result<(), TypedStoreError>
    where
        V: Serialize,
    {
        let timer = self
            .db_metrics
            .op_metrics
            .rocksdb_put_latency_seconds
            .with_label_values(&[self.metrics_label()])
            .start_timer();
        let perf_ctx = if self.write_sample_interval.sample() {
            Some(RocksDBPerfContext)
        } else {
            None
        };
        let value_buf = bcs::to_bytes(value).map_err(typed_store_err_from_bcs_err)?;
        self.db_metrics
            .op_metrics
            .rocksdb_put_bytes
            .with_label_values(&[self.metrics_label()])
            .observe((key_buf.len() + value_buf.len()) as f64);
        if perf_ctx.is_some() {
            self.db_metrics
                .write_perf_ctx_metrics
                .report_metrics(self.metrics_label());
        }
        self.db.put_cf(&self.column_family, key_buf, value_buf)?;

        let elapsed_secs = timer.stop_and_record();
        if elapsed_secs > 1.0 {
            warn!(elapsed_secs, cf = ?self.cf_name(), "very slow insert");
            self.db_metrics
                .op_metrics
                .rocksdb_very_slow_puts_count
                .with_label_values(&[self.metrics_label()])
                .inc();
            self.db_metrics
                .op_metrics
                .rocksdb_very_slow_puts_duration_ms
                .with_label_values(&[self.metrics_label()])
                .inc_by((elapsed_secs * 1000.0) as u64);
        }

        Ok(())
    }

    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn remove_raw_key(&self, key_buf: Vec<u8>) -> Result<(), TypedStoreError> {
        let _timer = self
            .db_metrics
            .op_metrics
            .rocksdb_delete_latency_seconds
            .with_label_values(&[self.metrics_label()])
            .start_timer();
        let perf_ctx = if self.write_sample_interval.sample() {
            Some(RocksDBPerfContext)
        } else {
            None
        };
        self.db.delete_cf(&self.column_family, key_buf)?;
        self.db_metrics
            .op_metrics
            .rocksdb_deletes
            .with_label_values(&[self.metrics_label()])
            .inc();
        if perf_ctx.is_some() {
            self.db_metrics
                .write_perf_ctx_metrics
                .report_metrics(self.metrics_label());
        }
        Ok(())
    }

    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn multi_get_raw_keys(
        &self,
        keys: impl IntoIterator<Item = Vec<u8>>,
    ) -> Result<Vec<Option<V>>, TypedStoreError>
    where
        V: DeserializeOwned,
    {
        let results = self.multi_get_pinned(keys)?;
        let values_parsed: Result<Vec<_>, TypedStoreError> = results
            .into_iter()
            .map(|value_byte| match value_byte {
                Some(data) => Ok(Some(
                    bcs::from_bytes(&data).map_err(typed_store_err_from_bcs_err)?,
                )),
                None => Ok(None),
            })
            .collect();

        values_parsed
    }
}

impl<'a, K, V> Map<'a, K, V> for DBMap<K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Error = TypedStoreError;

    fn contains_key(&self, key: &K) -> Result<bool, TypedStoreError> {
        self.contains_raw_key(be_fix_int_ser(key))
    }

    fn multi_contains_keys<J>(
        &self,
        keys: impl IntoIterator<Item = J>,
    ) -> Result<Vec<bool>, Self::Error>
    where
        J: Borrow<K>,
    {
        self.multi_contains_raw_keys(keys.into_iter().map(|k| be_fix_int_ser(k.borrow())))
    }

    fn get(&self, key: &K) -> Result<Option<V>, TypedStoreError> {
        self.get_raw_key(be_fix_int_ser(key))
    }

    fn insert(&self, key: &K, value: &V) -> Result<(), TypedStoreError> {
        self.insert_raw_key(be_fix_int_ser(key), value)
    }

    fn remove(&self, key: &K) -> Result<(), TypedStoreError> {
        self.remove_raw_key(be_fix_int_ser(key))
    }

    /// Writes a range delete tombstone to delete all entries in the db map.
    /// The effect of this write is visible immediately, i.e. you won't see
    /// old values when you do a lookup or scan.
    #[instrument(level = "trace", skip_all, err)]
    fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
        let Some(last_key) = self
            .safe_range_iter_reversed(..)
            .next()
            .transpose()?
            .map(|(k, _v)| k)
        else {
            return Ok(());
        };
        // The tombstone excludes its upper bound, so it has to end past the last key:
        // appending a zero byte gives the first key sorting above it.
        let mut to = be_fix_int_ser(&last_key);
        to.push(0);
        // The empty key sorts below every key, so the range starts before the first
        // entry.
        let mut batch = self.batch();
        batch.schedule_delete_range_raw(self, Vec::new(), to)?;
        batch.write()
    }

    fn is_empty(&self) -> bool {
        self.safe_iter().next().is_none()
    }

    fn safe_iter(&'a self) -> DbIterator<'a, (K, V)> {
        match &self.db.storage {
            Storage::Rocks(db) => match self.rocks_safe_iter(db, self.opts.readopts()) {
                Ok(iter) => Box::new(iter),
                Err(e) => error_iterator(e),
            },
            Storage::InMemory(db) => db.iterator(self.column_family.name(), None, None, false),
        }
    }

    fn safe_iter_with_bounds(
        &'a self,
        lower_bound: Option<K>,
        upper_bound: Option<K>,
    ) -> DbIterator<'a, (K, V)> {
        let range = (
            lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded),
            upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
        );
        self.safe_range_iter(range)
    }

    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
        let (lower_bound, upper_bound) = iterator_bounds_with_range(range);
        self.iter_forward_raw(lower_bound, upper_bound)
    }

    /// Both directions derive their bounds from the same
    /// `iterator_bounds_with_range`, so they are guaranteed to cover the
    /// identical set of keys regardless of the bound inclusivity.
    fn safe_range_iter_reversed(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
        let (lower_bound, upper_bound) = iterator_bounds_with_range(range);
        self.iter_reversed_raw(lower_bound, upper_bound)
    }

    /// Returns a vector of values corresponding to the keys provided.
    fn multi_get<J>(
        &self,
        keys: impl IntoIterator<Item = J>,
    ) -> Result<Vec<Option<V>>, TypedStoreError>
    where
        J: Borrow<K>,
    {
        self.multi_get_raw_keys(keys.into_iter().map(|k| be_fix_int_ser(k.borrow())))
    }

    /// Convenience method for batch insertion
    #[instrument(level = "trace", skip_all, err)]
    fn multi_insert<J, U>(
        &self,
        key_val_pairs: impl IntoIterator<Item = (J, U)>,
    ) -> Result<(), Self::Error>
    where
        J: Borrow<K>,
        U: Borrow<V>,
    {
        let mut batch = self.batch();
        batch.insert_batch(self, key_val_pairs)?;
        batch.write()
    }

    /// Convenience method for batch removal
    #[instrument(level = "trace", skip_all, err)]
    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), Self::Error>
    where
        J: Borrow<K>,
    {
        let mut batch = self.batch();
        batch.delete_batch(self, keys)?;
        batch.write()
    }

    /// Try to catch up with primary when running as secondary
    #[instrument(level = "trace", skip_all, err)]
    fn try_catch_up_with_primary(&self) -> Result<(), Self::Error> {
        self.db.try_catch_up_with_primary()
    }
}

/// A typed map sharing one column family with other typed maps,
/// distinguished by a tag byte prefixed to every key: the big-endian fixint
/// serialization of `(tag, key)` keeps each map a contiguous key range of
/// the column family. Useful for sets of tables that are always created and
/// dropped together — one column family instead of one per table keeps the
/// column-family and SST-file counts low.
///
/// Every read is scoped to the map's tag: full scans iterate the tag's key
/// prefix, and ranges are two-sided within it, so rows of other tags never
/// surface and their (differently typed) values are never deserialized.
///
/// Tags of existing maps must never change or be reused, or old rows would
/// be read under the wrong types.
///
/// Metrics are labelled by column family, so all the maps of one column
/// family report together, with no way to tell the tags apart.
///
/// # Examples
///
/// ```
/// use tempfile::tempdir;
/// use typed_store::{Map, rocks::*};
///
/// #[tokio::main(flavor = "current_thread")]
/// async fn main() {
///     let db = open_cf_opts(
///         tempdir().unwrap(),
///         None,
///         MetricConf::default(),
///         &[("shared", rocksdb::Options::default())],
///     )
///     .unwrap();
///
///     let numbers: TaggedDBMap<u32, String> =
///         TaggedDBMap::reopen(&db, "shared", 0, &ReadWriteOptions::default(), false).unwrap();
///     let words: TaggedDBMap<String, u64> =
///         TaggedDBMap::reopen(&db, "shared", 1, &ReadWriteOptions::default(), false).unwrap();
///
///     let mut batch = numbers.batch();
///     batch
///         .insert_batch_tagged(&numbers, [(1, "one".to_string())])
///         .unwrap()
///         .insert_batch_tagged(&words, [("one".to_string(), 1)])
///         .unwrap();
///     batch.write().unwrap();
///
///     assert_eq!(numbers.get(&1).unwrap(), Some("one".to_string()));
///     assert_eq!(words.get(&"one".to_string()).unwrap(), Some(1));
///     // Each map sees only its own row.
///     assert_eq!(numbers.safe_iter().count(), 1);
/// }
/// ```
pub struct TaggedDBMap<K, V> {
    tag: u8,
    map: DBMap<(u8, K), V>,
}

impl<K, V> TaggedDBMap<K, V> {
    /// Drops the tag prefixed to a key of the underlying map, turning its row
    /// into a row of this map.
    fn strip_tag(row: Result<((u8, K), V), TypedStoreError>) -> Result<(K, V), TypedStoreError> {
        row.map(|((_, key), value)| (key, value))
    }

    pub(crate) fn tag(&self) -> u8 {
        self.tag
    }

    /// [`DBMap::ingest_sst_files`] for this map's column family, which the
    /// files may hold rows of any table sharing it for.
    pub fn ingest_sst_files(&self, paths: &[PathBuf]) -> Result<(), TypedStoreError> {
        self.map.ingest_sst_files(paths)
    }

    /// Reopens an open database as a tagged map operating under `cf_name`.
    /// See [`DBMap::reopen`] for the remaining parameters.
    pub fn reopen(
        db: &Arc<Database>,
        cf_name: &str,
        tag: u8,
        rw_options: &ReadWriteOptions,
        skip_metrics_reporting: bool,
    ) -> Result<Self, TypedStoreError> {
        Ok(Self {
            tag,
            map: DBMap::reopen(db, Some(cf_name), rw_options, skip_metrics_reporting)?,
        })
    }

    /// Create a batch associated with the map's database.
    pub fn batch(&self) -> DBBatch {
        self.map.batch()
    }

    /// See [`DBMap::with_metrics_label`].
    pub fn with_metrics_label(mut self, label: &str) -> Self {
        self.map = self.map.with_metrics_label(label);
        self
    }
}

impl<'a, K, V> Map<'a, K, V> for TaggedDBMap<K, V>
where
    K: Serialize + DeserializeOwned,
    V: Serialize + DeserializeOwned,
{
    type Error = TypedStoreError;

    fn contains_key(&self, key: &K) -> Result<bool, TypedStoreError> {
        self.map.contains_raw_key(be_fix_int_ser(&(self.tag, key)))
    }

    fn multi_contains_keys<J>(
        &self,
        keys: impl IntoIterator<Item = J>,
    ) -> Result<Vec<bool>, TypedStoreError>
    where
        J: Borrow<K>,
    {
        self.map.multi_contains_raw_keys(
            keys.into_iter()
                .map(|key| be_fix_int_ser(&(self.tag, key.borrow()))),
        )
    }

    fn get(&self, key: &K) -> Result<Option<V>, TypedStoreError> {
        self.map.get_raw_key(be_fix_int_ser(&(self.tag, key)))
    }

    fn multi_get<J>(
        &self,
        keys: impl IntoIterator<Item = J>,
    ) -> Result<Vec<Option<V>>, TypedStoreError>
    where
        J: Borrow<K>,
    {
        self.map.multi_get_raw_keys(
            keys.into_iter()
                .map(|key| be_fix_int_ser(&(self.tag, key.borrow()))),
        )
    }

    fn insert(&self, key: &K, value: &V) -> Result<(), TypedStoreError> {
        self.map
            .insert_raw_key(be_fix_int_ser(&(self.tag, key)), value)
    }

    fn remove(&self, key: &K) -> Result<(), TypedStoreError> {
        self.map.remove_raw_key(be_fix_int_ser(&(self.tag, key)))
    }

    #[instrument(level = "trace", skip_all, err)]
    fn multi_insert<J, U>(
        &self,
        key_val_pairs: impl IntoIterator<Item = (J, U)>,
    ) -> Result<(), TypedStoreError>
    where
        J: Borrow<K>,
        U: Borrow<V>,
    {
        let mut batch = self.batch();
        batch.insert_batch_tagged(self, key_val_pairs)?;
        batch.write()
    }

    #[instrument(level = "trace", skip_all, err)]
    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), TypedStoreError>
    where
        J: Borrow<K>,
    {
        let mut batch = self.batch();
        batch.delete_batch_tagged(self, keys)?;
        batch.write()
    }

    /// Deletes every entry of the map with a single range tombstone over the
    /// tag's key range, leaving the other tags of the column family untouched.
    ///
    /// The entries stay on disk until compaction, and until then iteration
    /// is slowed because of them.
    #[instrument(level = "trace", skip_all, err)]
    fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
        let from = be_fix_int_ser(&self.tag);
        let to = match prefix_iterator_bounds(&self.tag).1 {
            Some(to) => {
                if self.is_empty() {
                    return Ok(());
                }
                to
            }
            // The maximum tag has no following tag to bound the range
            // against, so end it just past the last entry. That read is why a
            // row written in between survives the clear, and why a value that
            // does not deserialize fails it.
            None => match self.safe_range_iter_reversed(..).next().transpose()? {
                Some((last_key, _)) => {
                    let mut to = be_fix_int_ser(&(self.tag, last_key));
                    to.push(0);
                    to
                }
                None => return Ok(()),
            },
        };

        let mut batch = self.batch();
        batch.schedule_delete_range_raw(&self.map, from, to)?;
        batch.write()
    }

    fn is_empty(&self) -> bool {
        self.safe_iter().next().is_none()
    }

    fn safe_iter(&'a self) -> DbIterator<'a, (K, V)> {
        Box::new(
            self.map
                .safe_iter_with_prefix(&self.tag)
                .map(Self::strip_tag),
        )
    }

    fn safe_iter_with_bounds(
        &'a self,
        lower_bound: Option<K>,
        upper_bound: Option<K>,
    ) -> DbIterator<'a, (K, V)> {
        let range = (
            lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded),
            upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
        );
        self.safe_range_iter(range)
    }

    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
        let (lower_bound, upper_bound) = prefix_iterator_bounds_with_range(&self.tag, range);
        Box::new(
            self.map
                .iter_forward_raw(lower_bound, upper_bound)
                .map(Self::strip_tag),
        )
    }

    fn safe_range_iter_reversed(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
        let (lower_bound, upper_bound) = prefix_iterator_bounds_with_range(&self.tag, range);
        Box::new(
            self.map
                .iter_reversed_raw(lower_bound, upper_bound)
                .map(Self::strip_tag),
        )
    }

    fn try_catch_up_with_primary(&self) -> Result<(), TypedStoreError> {
        self.map.try_catch_up_with_primary()
    }
}

/// Waits until the database of `db` has been closed, every handle to it
/// dropped and RocksDB's files released, so its directory can be reopened,
/// moved, or deleted. Returns `false` if it is still open after 30 seconds.
pub async fn wait_for_database_close(db: DatabaseClose) -> bool {
    // The deadline is read off the same clock the sleep advances, so a test
    // pausing time does not turn the wait into a busy loop.
    let deadline = Instant::now() + Duration::from_secs(30);
    while db.0.strong_count() != 0 {
        if Instant::now() > deadline {
            return false;
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
    true
}

/// Returns an iterator whose only item is `Err(error)`.
pub(crate) fn error_iterator<'a, T: 'a>(error: TypedStoreError) -> DbIterator<'a, T> {
    Box::new(std::iter::once(Err(error)))
}

/// Write options that discard the batch entries of a column family dropped
/// after staging, instead of failing the write; see [`DBBatch::write_opt`].
pub fn drop_tolerant_write_options() -> rocksdb::WriteOptions {
    let mut options = rocksdb::WriteOptions::default();
    options.set_ignore_missing_column_families(true);
    options
}

fn default_hash(value: &[u8]) -> Digest<32> {
    let mut hasher = fastcrypto::hash::Blake2b256::default();
    hasher.update(value);
    hasher.finalize()
}

// The throughput floor is used to compute the threshold for when a batch write
// is considered "very slow".
const THROUGHPUT_FLOOR_BYTES_PER_SEC: f64 = 32.0 * 1024.0 * 1024.0; // 32 MiB/s

// A batch write that took less than this is never reported as very slow.
const MIN_VERY_SLOW_BATCH_WRITE_SECS: f64 = 1.0;

/// Returns the elapsed time above which a batch write of `batch_size_bytes` is
/// reported as very slow: what it would take at
/// [`THROUGHPUT_FLOOR_BYTES_PER_SEC`], and never less than
/// [`MIN_VERY_SLOW_BATCH_WRITE_SECS`].
fn very_slow_batch_write_threshold_secs(batch_size_bytes: usize) -> f64 {
    (batch_size_bytes as f64 / THROUGHPUT_FLOOR_BYTES_PER_SEC).max(MIN_VERY_SLOW_BATCH_WRITE_SECS)
}

#[cfg(test)]
mod tests {
    use super::very_slow_batch_write_threshold_secs;

    const MIB: usize = 1024 * 1024;

    #[test]
    fn very_slow_batch_write_threshold_never_below_one_second() {
        assert_eq!(very_slow_batch_write_threshold_secs(0), 1.0);
        assert_eq!(very_slow_batch_write_threshold_secs(1), 1.0);
        assert_eq!(very_slow_batch_write_threshold_secs(MIB), 1.0);
        assert_eq!(very_slow_batch_write_threshold_secs(32 * MIB), 1.0);
    }

    #[test]
    fn very_slow_batch_write_threshold_grows_with_batch_size() {
        assert_eq!(very_slow_batch_write_threshold_secs(48 * MIB), 1.5);
        assert_eq!(very_slow_batch_write_threshold_secs(64 * MIB), 2.0);
        assert_eq!(very_slow_batch_write_threshold_secs(320 * MIB), 10.0);
    }
}