tsoracle-paxos-toolkit 0.2.3

//
//  ░▀█▀░█▀▀░█▀█░█▀▄░█▀█░█▀▀░█░░░█▀▀
//  ░░█░░▀▀█░█░█░█▀▄░█▀█░█░░░█░░░█▀▀
//  ░░▀░░▀▀▀░▀▀▀░▀░▀░▀░▀░▀▀▀░▀▀▀░▀▀▀
//
//  tsoracle — Distributed Timestamp Oracle
//
//  Copyright (c) 2026 Prisma Risk
//  Licensed under the Apache License, Version 2.0
//  https://github.com/prisma-risk/tsoracle
//

//! RocksDB-backed `omnipaxos::storage::Storage` implementation.
//!
//! Log keys are absolute; the `compacted_idx` offset is tracked separately
//! via the meta column. `get_log_len()` returns the physical remaining
//! count, matching `omnipaxos_storage::persistent_storage`. OmniPaxos pairs
//! `trim(idx)` with `set_compacted_idx(idx)`; the two are independent here.

pub mod key_space;
pub mod meta;

#[cfg(feature = "rocksdb-storage")]
use std::marker::PhantomData;
#[cfg(feature = "rocksdb-storage")]
use std::sync::Arc;

#[cfg(feature = "rocksdb-storage")]
use omnipaxos::storage::Entry;
#[cfg(feature = "rocksdb-storage")]
use rocksdb::{BoundColumnFamily, DB, WriteBatch, WriteOptions};

#[cfg(feature = "rocksdb-storage")]
#[derive(Debug, thiserror::Error)]
pub enum StorageError {
    #[error("column family `{0}` not found in the supplied DB handle")]
    ColumnFamilyNotFound(String),
    #[error("rocksdb error: {0}")]
    Rocksdb(#[from] rocksdb::Error),
    #[error("meta serialization error: {0}")]
    Meta(#[from] crate::storage::meta::MetaError),
    #[error("codec error: {0}")]
    Codec(#[from] crate::codec::CodecError),
    #[error("log integrity violation: {0}")]
    LogIntegrity(String),
}

/// RocksDB-backed implementation of [`omnipaxos::storage::Storage`].
///
/// The caller owns the `Arc<DB>` and the lifetime of the column family. This
/// crate never opens or closes a database; constructing a `RocksdbStorage`
/// only validates that the requested column family is present. Sharing the
/// same `Arc<DB>` with other column families (e.g., a snapshot store) is
/// supported and is the intended pattern.
///
/// The type parameter `T` is the OmniPaxos `Entry` type the log will hold;
/// the implementation persists log entries via `postcard` and meta fields
/// via fixed-shape encoders in [`crate::storage::meta`].
///
/// Construct with [`RocksdbStorage::open_in`].
#[cfg(feature = "rocksdb-storage")]
pub struct RocksdbStorage<T: Entry> {
    db: Arc<DB>,
    cf_name: String,
    /// Cached next absolute write index. Holds `(highest_log_key + 1).max(compacted_idx)`
    /// when the log is non-empty, else `compacted_idx` — the same value the old
    /// per-append reverse scan computed. Seeded once at [`RocksdbStorage::open_in`]
    /// and maintained in lockstep with every mutating write (always *after* the
    /// write succeeds), so the append hot path discovers the next index from
    /// memory rather than from a RocksDB iterator. The `&mut self` on every
    /// mutating `Storage` method makes this sound without interior mutability;
    /// correctness assumes a single `RocksdbStorage` owns its column family,
    /// which is the documented usage.
    next_idx: u64,
    /// Cached compaction offset, a mirror of the persisted `M/compacted_idx`.
    /// Only [`omnipaxos::storage::Storage::set_compacted_idx`] writes it.
    compacted_idx: u64,
    _marker: PhantomData<T>,
}

#[cfg(feature = "rocksdb-storage")]
impl<T: Entry> RocksdbStorage<T> {
    /// Build a storage handle backed by an existing column family.
    ///
    /// The column family must already exist in `db`. The constructor only
    /// validates presence — it does NOT create the CF. Callers typically
    /// open the database with the appropriate `ColumnFamilyDescriptor`s
    /// for both the paxos log and any peer column families (snapshot
    /// store, application state) before invoking this.
    ///
    /// Returns [`StorageError::ColumnFamilyNotFound`] if `cf_name` is not
    /// registered on the supplied database.
    pub fn open_in(db: Arc<DB>, cf_name: &str) -> Result<Self, StorageError> {
        if db.cf_handle(cf_name).is_none() {
            return Err(StorageError::ColumnFamilyNotFound(cf_name.to_string()));
        }
        let mut storage = Self {
            db,
            cf_name: cf_name.to_string(),
            next_idx: 0,
            compacted_idx: 0,
            _marker: PhantomData,
        };
        // Seed the in-memory cursors from disk exactly once. This is the only
        // place the log is scanned for its tail; every later mutation keeps the
        // cursors current, so appends never re-scan. A reopen after a crash
        // re-derives both from the persisted state here.
        storage.compacted_idx = storage.read_compacted_idx()?;
        storage.next_idx = storage.scan_next_log_idx(storage.compacted_idx)?;
        Ok(storage)
    }

    /// Reverse-scan the log column family for the highest absolute index,
    /// returning `(highest + 1).max(compacted)`, or `compacted` when the log
    /// is empty. Used only by [`RocksdbStorage::open_in`] to seed `next_idx`;
    /// the floor matches the on-disk invariant a fresh append must satisfy
    /// (`next >= compacted_idx`) so a snapshot that trims every surviving log
    /// key cannot let an append land below the compaction floor.
    fn scan_next_log_idx(&self, compacted: u64) -> Result<u64, StorageError> {
        use crate::storage::key_space::{LOG_PREFIX, log_key_range, parse_log_key};
        use rocksdb::IteratorMode;

        let cf = self.cf()?;
        let (_, upper) = log_key_range();
        let iter = self
            .db
            .iterator_cf(&cf, IteratorMode::From(&upper, rocksdb::Direction::Reverse));
        for result in iter {
            let (key, _value) = result?;
            if key.starts_with(LOG_PREFIX) {
                let idx = parse_log_key(&key).ok_or_else(|| {
                    StorageError::LogIntegrity(format!("malformed log key: {key:?}"))
                })?;
                return Ok((idx + 1).max(compacted));
            }
        }
        Ok(compacted)
    }

    /// Read the persisted compaction offset (`M/compacted_idx`), defaulting to
    /// `0` when absent. Used only to seed `compacted_idx` at
    /// [`RocksdbStorage::open_in`]; thereafter the cached field is authoritative.
    fn read_compacted_idx(&self) -> Result<u64, StorageError> {
        use crate::storage::key_space::meta_compacted_idx_key;
        use crate::storage::meta::decode_u64;
        let cf = self.cf()?;
        match self.db.get_cf(&cf, meta_compacted_idx_key())? {
            Some(bytes) => Ok(decode_u64(&bytes)?),
            None => Ok(0),
        }
    }

    pub(crate) fn cf(&self) -> Result<Arc<BoundColumnFamily<'_>>, StorageError> {
        self.db
            .cf_handle(&self.cf_name)
            .ok_or_else(|| StorageError::ColumnFamilyNotFound(self.cf_name.clone()))
    }

    /// Write options for Paxos safety-critical writes.
    ///
    /// `sync = true` forces fsync of the WAL on every call. Used for
    /// [`Storage::set_promise`], [`Storage::set_accepted_round`],
    /// [`Storage::append_entry`], [`Storage::append_entries`], and
    /// [`Storage::append_on_prefix`] — a node that promises ballot N or
    /// accepts a log entry must not forget that promise across a crash, or
    /// Paxos safety is violated.
    pub(crate) fn write_sync_opts() -> WriteOptions {
        let mut opts = WriteOptions::default();
        opts.set_sync(true);
        opts
    }

    /// Write options for non-safety-critical writes.
    ///
    /// Used for [`Storage::set_decided_idx`], [`Storage::set_compacted_idx`],
    /// [`Storage::trim`], [`Storage::set_snapshot`], and
    /// [`Storage::set_stopsign`] — these fields can be reconstructed or
    /// recovered after a crash (decided_idx is monotonic from the log
    /// contents; compacted_idx is bounded by the trimmed range; snapshot
    /// and stopsign can be re-derived). Avoiding fsync on the hot apply
    /// path keeps throughput sane.
    pub(crate) fn write_async_opts() -> WriteOptions {
        WriteOptions::default()
    }

    pub(crate) fn batch_sync<F>(&self, f: F) -> Result<(), StorageError>
    where
        F: FnOnce(Arc<BoundColumnFamily<'_>>, &mut WriteBatch) -> Result<(), StorageError>,
    {
        let cf = self.cf()?;
        let mut batch = WriteBatch::default();
        f(cf, &mut batch)?;
        self.db.write_opt(batch, &Self::write_sync_opts())?;
        Ok(())
    }

    /// Run a non-synced batch write, tagged with the OmniPaxos storage
    /// operation `op` driving it (`set_decided_idx`, `set_compacted_idx`,
    /// `trim`, `set_snapshot`, `set_stopsign`). On failure, the error is
    /// counted under `tsoracle.paxos.storage.async_write_failures.total`
    /// (labelled by `op`) before being returned, so a recurring async-write
    /// fault surfaces as an operational signal instead of only flowing up as a
    /// `StorageResult` the caller may swallow. The counter lives on the error
    /// path, so the hot apply-path write pays only a branch.
    pub(crate) fn batch_async<F>(&self, op: &'static str, f: F) -> Result<(), StorageError>
    where
        F: FnOnce(Arc<BoundColumnFamily<'_>>, &mut WriteBatch) -> Result<(), StorageError>,
    {
        // `op` is read only by the failpoint message and the failure counter;
        // keep it live when both of those are compiled out.
        #[cfg(not(any(feature = "metrics", feature = "failpoints")))]
        let _ = op;
        let result = (|| {
            tsoracle_failpoint::failpoint!("paxos_toolkit::storage::async_write", |_| Err(
                StorageError::LogIntegrity(format!("async_write failpoint armed for {op}"))
            ));
            let cf = self.cf()?;
            let mut batch = WriteBatch::default();
            f(cf, &mut batch)?;
            self.db.write_opt(batch, &Self::write_async_opts())?;
            Ok(())
        })();
        #[cfg(feature = "metrics")]
        if result.is_err() {
            metrics::counter!("tsoracle.paxos.storage.async_write_failures.total", "op" => op)
                .increment(1);
        }
        result
    }
}

/// Box a storage-side error for OmniPaxos's [`omnipaxos::storage::StorageResult`].
///
/// OmniPaxos fixes the error type at `Box<dyn Error>` (no `Send + Sync`), and
/// every call site funnels straight into a `StorageResult` via `?`, so that is
/// the return type here — `?` narrows through the `From<E>` conversion, never
/// the unsize coercion a wider box would require. The `E: Send + Sync` bound on
/// the *input* keeps this consistent with the project-wide
/// `Box<dyn Error + Send + Sync>` convention and statically guarantees every
/// error we box is thread-movable at the source, even though OmniPaxos's
/// signature erases that capability at the trait boundary.
#[cfg(feature = "rocksdb-storage")]
fn box_err<E: std::error::Error + Send + Sync + 'static>(err: E) -> Box<dyn std::error::Error> {
    Box::new(err)
}

#[cfg(feature = "rocksdb-storage")]
impl<T> RocksdbStorage<T>
where
    T: Entry + serde::Serialize + serde::de::DeserializeOwned + Clone + 'static,
{
    fn store_log_entry(
        &self,
        batch: &mut WriteBatch,
        cf: &Arc<BoundColumnFamily<'_>>,
        idx: u64,
        entry: &T,
    ) -> Result<(), StorageError> {
        use crate::codec::{SCHEMA_VERSION, encode as codec_encode};
        use crate::storage::key_space::log_key;
        let key = log_key(idx);
        let value = codec_encode(SCHEMA_VERSION, entry)?;
        batch.put_cf(cf, key, value);
        Ok(())
    }
}

#[cfg(feature = "rocksdb-storage")]
impl<T> omnipaxos::storage::Storage<T> for RocksdbStorage<T>
where
    T: Entry + serde::Serialize + serde::de::DeserializeOwned + Clone + 'static,
    T::Snapshot: serde::Serialize + serde::de::DeserializeOwned + Clone,
{
    fn append_entry(&mut self, entry: T) -> omnipaxos::storage::StorageResult<u64> {
        tsoracle_failpoint::failpoint!("paxos_toolkit::storage::append_entry");
        let next = self.next_idx;
        self.batch_sync(|cf, batch| self.store_log_entry(batch, &cf, next, &entry))
            .map_err(box_err)?;
        // Advance the cursor only after the write commits, so a failed write
        // leaves the cache matching disk.
        self.next_idx = next + 1;
        Ok(self.next_idx.saturating_sub(self.compacted_idx))
    }

    fn append_entries(&mut self, entries: Vec<T>) -> omnipaxos::storage::StorageResult<u64> {
        let start = self.next_idx;
        let count = entries.len() as u64;
        self.batch_sync(|cf, batch| {
            for (offset, entry) in entries.iter().enumerate() {
                self.store_log_entry(batch, &cf, start + offset as u64, entry)?;
            }
            Ok(())
        })
        .map_err(box_err)?;
        self.next_idx = start + count;
        Ok(self.next_idx.saturating_sub(self.compacted_idx))
    }

    fn append_on_prefix(
        &mut self,
        from_idx: u64,
        entries: Vec<T>,
    ) -> omnipaxos::storage::StorageResult<u64> {
        use crate::storage::key_space::{log_key, log_key_range};
        let (_, upper) = log_key_range();
        let count = entries.len() as u64;
        self.batch_sync(|cf, batch| {
            let lower = log_key(from_idx);
            batch.delete_range_cf(&cf, &lower, &upper);
            for (offset, entry) in entries.iter().enumerate() {
                self.store_log_entry(batch, &cf, from_idx + offset as u64, entry)?;
            }
            Ok(())
        })
        .map_err(box_err)?;
        // The write truncated every key `>= from_idx` and re-appended `count`
        // of them, so the new tail is `from_idx + count`, floored at the
        // compaction offset just as a fresh append would be.
        self.next_idx = (from_idx + count).max(self.compacted_idx);
        Ok(self.next_idx.saturating_sub(self.compacted_idx))
    }

    fn set_promise(
        &mut self,
        n_prom: omnipaxos::ballot_leader_election::Ballot,
    ) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::meta_promise_key;
        use crate::storage::meta::encode_ballot;
        let bytes = encode_ballot(&n_prom).map_err(box_err)?;
        self.batch_sync(|cf, batch| {
            batch.put_cf(&cf, meta_promise_key(), bytes);
            Ok(())
        })
        .map_err(box_err)?;
        Ok(())
    }

    fn set_decided_idx(&mut self, ld: u64) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::meta_decided_idx_key;
        use crate::storage::meta::encode_u64;
        let bytes = encode_u64(ld);
        self.batch_async("set_decided_idx", |cf, batch| {
            batch.put_cf(&cf, meta_decided_idx_key(), bytes);
            Ok(())
        })
        .map_err(box_err)?;
        Ok(())
    }

    fn get_decided_idx(&self) -> omnipaxos::storage::StorageResult<u64> {
        use crate::storage::key_space::meta_decided_idx_key;
        use crate::storage::meta::decode_u64;
        let cf = self.cf().map_err(box_err)?;
        match self
            .db
            .get_cf(&cf, meta_decided_idx_key())
            .map_err(box_err)?
        {
            Some(bytes) => Ok(decode_u64(&bytes).map_err(box_err)?),
            None => Ok(0),
        }
    }

    fn set_accepted_round(
        &mut self,
        na: omnipaxos::ballot_leader_election::Ballot,
    ) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::meta_accepted_round_key;
        use crate::storage::meta::encode_ballot;
        let bytes = encode_ballot(&na).map_err(box_err)?;
        self.batch_sync(|cf, batch| {
            batch.put_cf(&cf, meta_accepted_round_key(), bytes);
            Ok(())
        })
        .map_err(box_err)?;
        Ok(())
    }

    fn get_accepted_round(
        &self,
    ) -> omnipaxos::storage::StorageResult<Option<omnipaxos::ballot_leader_election::Ballot>> {
        use crate::storage::key_space::meta_accepted_round_key;
        use crate::storage::meta::decode_ballot;
        let cf = self.cf().map_err(box_err)?;
        match self
            .db
            .get_cf(&cf, meta_accepted_round_key())
            .map_err(box_err)?
        {
            Some(bytes) => Ok(Some(decode_ballot(&bytes).map_err(box_err)?)),
            None => Ok(None),
        }
    }

    fn get_entries(&self, from: u64, to: u64) -> omnipaxos::storage::StorageResult<Vec<T>> {
        use crate::codec::{SCHEMA_VERSION, decode as codec_decode};
        use crate::storage::key_space::{LOG_PREFIX, log_key};
        use rocksdb::IteratorMode;

        if from >= to {
            return Ok(Vec::new());
        }
        let cf = self.cf().map_err(box_err)?;
        let start = log_key(from);
        let end = log_key(to);
        let expected = (to - from) as usize;
        let iter = self
            .db
            .iterator_cf(&cf, IteratorMode::From(&start, rocksdb::Direction::Forward));
        let mut out = Vec::with_capacity(expected);
        for result in iter {
            let (key, value) = result.map_err(box_err)?;
            if !key.starts_with(LOG_PREFIX) || key.as_ref() >= end.as_slice() {
                break;
            }
            let entry: T = codec_decode(SCHEMA_VERSION, &value).map_err(box_err)?;
            out.push(entry);
            if out.len() == expected {
                break;
            }
        }
        // Gap contract: any missing index in [from, to) yields an empty Vec
        // rather than a partial prefix. The iterator returns entries in
        // ascending log-index order, so if the count is short, at least
        // one index in the requested range is missing.
        if out.len() < expected {
            return Ok(Vec::new());
        }
        Ok(out)
    }

    fn get_log_len(&self) -> omnipaxos::storage::StorageResult<u64> {
        Ok(self.next_idx.saturating_sub(self.compacted_idx))
    }

    fn get_suffix(&self, from: u64) -> omnipaxos::storage::StorageResult<Vec<T>> {
        use crate::codec::{SCHEMA_VERSION, decode as codec_decode};
        use crate::storage::key_space::{LOG_PREFIX, log_key};
        use rocksdb::IteratorMode;

        // `IteratorMode::From(log_key(from), Forward)` seeks to the first
        // key `>= log_key(from)`. Because log keys are big-endian u64,
        // every key yielded by the iterator has an index `>= from`, so we
        // do not need to re-check the index after decoding.
        let cf = self.cf().map_err(box_err)?;
        let start = log_key(from);
        let iter = self
            .db
            .iterator_cf(&cf, IteratorMode::From(&start, rocksdb::Direction::Forward));
        let mut out = Vec::new();
        for result in iter {
            let (key, value) = result.map_err(box_err)?;
            if !key.starts_with(LOG_PREFIX) {
                break;
            }
            let entry: T = codec_decode(SCHEMA_VERSION, &value).map_err(box_err)?;
            out.push(entry);
        }
        Ok(out)
    }

    fn get_promise(
        &self,
    ) -> omnipaxos::storage::StorageResult<Option<omnipaxos::ballot_leader_election::Ballot>> {
        use crate::storage::key_space::meta_promise_key;
        use crate::storage::meta::decode_ballot;
        let cf = self.cf().map_err(box_err)?;
        match self.db.get_cf(&cf, meta_promise_key()).map_err(box_err)? {
            Some(bytes) => Ok(Some(decode_ballot(&bytes).map_err(box_err)?)),
            None => Ok(None),
        }
    }

    fn set_stopsign(
        &mut self,
        stopsign: Option<omnipaxos::storage::StopSign>,
    ) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::meta_stopsign_key;
        use crate::storage::meta::encode_postcard;
        match stopsign {
            Some(inner) => {
                let bytes = encode_postcard(&inner).map_err(box_err)?;
                self.batch_async("set_stopsign", |cf, batch| {
                    batch.put_cf(&cf, meta_stopsign_key(), bytes);
                    Ok(())
                })
                .map_err(box_err)?;
            }
            None => {
                self.batch_async("set_stopsign", |cf, batch| {
                    batch.delete_cf(&cf, meta_stopsign_key());
                    Ok(())
                })
                .map_err(box_err)?;
            }
        }
        Ok(())
    }

    fn get_stopsign(
        &self,
    ) -> omnipaxos::storage::StorageResult<Option<omnipaxos::storage::StopSign>> {
        use crate::storage::key_space::meta_stopsign_key;
        use crate::storage::meta::decode_postcard;
        let cf = self.cf().map_err(box_err)?;
        match self.db.get_cf(&cf, meta_stopsign_key()).map_err(box_err)? {
            Some(bytes) => Ok(Some(decode_postcard(&bytes).map_err(box_err)?)),
            None => Ok(None),
        }
    }

    fn trim(&mut self, idx: u64) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::log_key;
        self.batch_async("trim", |cf, batch| {
            let lower = log_key(0);
            let upper = log_key(idx);
            batch.delete_range_cf(&cf, &lower, &upper);
            Ok(())
        })
        .map_err(box_err)?;
        // Trimming a prefix leaves the tail untouched unless it removes every
        // live key (`idx` past the highest one). A now-empty log falls back to
        // the compaction floor, matching what a reverse scan would find.
        if idx >= self.next_idx {
            self.next_idx = self.compacted_idx;
        }
        Ok(())
    }

    fn set_compacted_idx(&mut self, idx: u64) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::meta_compacted_idx_key;
        use crate::storage::meta::encode_u64;
        let bytes = encode_u64(idx);
        self.batch_async("set_compacted_idx", |cf, batch| {
            batch.put_cf(&cf, meta_compacted_idx_key(), bytes);
            Ok(())
        })
        .map_err(box_err)?;
        self.compacted_idx = idx;
        // Raising the floor past every live key floats the next write up to it,
        // preserving the `next >= compacted_idx` invariant.
        self.next_idx = self.next_idx.max(idx);
        Ok(())
    }

    fn get_compacted_idx(&self) -> omnipaxos::storage::StorageResult<u64> {
        Ok(self.compacted_idx)
    }

    fn set_snapshot(
        &mut self,
        snapshot: Option<T::Snapshot>,
    ) -> omnipaxos::storage::StorageResult<()> {
        use crate::storage::key_space::meta_snapshot_key;
        use crate::storage::meta::encode_postcard;
        match snapshot {
            Some(snap) => {
                let bytes = encode_postcard(&snap).map_err(box_err)?;
                self.batch_async("set_snapshot", |cf, batch| {
                    batch.put_cf(&cf, meta_snapshot_key(), bytes);
                    Ok(())
                })
                .map_err(box_err)?;
            }
            None => {
                self.batch_async("set_snapshot", |cf, batch| {
                    batch.delete_cf(&cf, meta_snapshot_key());
                    Ok(())
                })
                .map_err(box_err)?;
            }
        }
        Ok(())
    }

    fn get_snapshot(&self) -> omnipaxos::storage::StorageResult<Option<T::Snapshot>> {
        use crate::storage::key_space::meta_snapshot_key;
        use crate::storage::meta::decode_postcard;
        let cf = self.cf().map_err(box_err)?;
        match self.db.get_cf(&cf, meta_snapshot_key()).map_err(box_err)? {
            Some(bytes) => Ok(Some(decode_postcard(&bytes).map_err(box_err)?)),
            None => Ok(None),
        }
    }
}

#[cfg(all(test, feature = "rocksdb-storage"))]
pub(crate) mod open_in_tests {
    use super::*;
    use rocksdb::{ColumnFamilyDescriptor, DB, Options};
    use std::sync::Arc;
    use tempfile::TempDir;

    pub(crate) fn open_db(dir: &TempDir, cf_name: &str) -> Arc<DB> {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);
        let cf = ColumnFamilyDescriptor::new(cf_name, Options::default());
        Arc::new(DB::open_cf_descriptors(&opts, dir.path(), vec![cf]).expect("open db"))
    }

    #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
    pub(crate) struct TestEntry {
        pub value: u64,
    }
    impl omnipaxos::storage::Entry for TestEntry {
        type Snapshot = TestSnapshot;
    }
    #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
    pub(crate) struct TestSnapshot {
        pub value: u64,
    }
    impl omnipaxos::storage::Snapshot<TestEntry> for TestSnapshot {
        fn create(entries: &[TestEntry]) -> Self {
            Self {
                value: entries.iter().map(|e| e.value).max().unwrap_or(0),
            }
        }
        fn merge(&mut self, other: Self) {
            self.value = self.value.max(other.value);
        }
        fn use_snapshots() -> bool {
            true
        }
    }

    #[test]
    fn open_in_returns_storage_for_existing_cf() {
        let dir = TempDir::new().expect("tempdir");
        let db = open_db(&dir, "tso_paxos");
        let storage: RocksdbStorage<TestEntry> =
            RocksdbStorage::open_in(db.clone(), "tso_paxos").expect("open_in");
        assert!(
            Arc::strong_count(&db) >= 2,
            "storage should hold an Arc to DB"
        );
        drop(storage);
    }

    #[test]
    fn open_in_rejects_missing_cf() {
        let dir = TempDir::new().expect("tempdir");
        let db = open_db(&dir, "tso_paxos");
        let result: Result<RocksdbStorage<TestEntry>, _> =
            RocksdbStorage::open_in(db, "missing_cf");
        assert!(matches!(result, Err(StorageError::ColumnFamilyNotFound(_))));
    }
}

#[cfg(all(test, feature = "rocksdb-storage"))]
mod log_tests {
    use super::open_in_tests::{TestEntry, open_db};
    use super::*;
    use omnipaxos::storage::Storage;
    use tempfile::TempDir;

    pub(super) fn fresh_storage(dir: &TempDir) -> RocksdbStorage<TestEntry> {
        let db = open_db(dir, "tso_paxos");
        RocksdbStorage::open_in(db, "tso_paxos").expect("open_in")
    }

    #[test]
    fn empty_log_has_zero_len() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert_eq!(storage.get_log_len().unwrap(), 0);
    }

    #[test]
    fn append_entry_returns_new_log_len() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        let len = storage.append_entry(TestEntry { value: 10 }).unwrap();
        assert_eq!(len, 1);
        let len = storage.append_entry(TestEntry { value: 20 }).unwrap();
        assert_eq!(len, 2);
    }

    #[test]
    fn append_entries_returns_new_log_len() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        let len = storage
            .append_entries(vec![
                TestEntry { value: 1 },
                TestEntry { value: 2 },
                TestEntry { value: 3 },
            ])
            .unwrap();
        assert_eq!(len, 3);
    }

    #[test]
    fn get_entries_returns_range() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![
                TestEntry { value: 10 },
                TestEntry { value: 20 },
                TestEntry { value: 30 },
            ])
            .unwrap();
        let got = storage.get_entries(0, 2).unwrap();
        assert_eq!(got.len(), 2);
        assert_eq!(got[0].value, 10);
        assert_eq!(got[1].value, 20);
    }

    #[test]
    fn get_entries_empty_on_gap() {
        // Contract: any index missing in [from, to) returns an empty Vec.
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![TestEntry { value: 1 }, TestEntry { value: 2 }])
            .unwrap();
        // Range [0, 5) crosses missing indices 2, 3, 4.
        assert!(storage.get_entries(0, 5).unwrap().is_empty());
    }

    #[test]
    fn get_suffix_returns_from_index_to_end() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![
                TestEntry { value: 100 },
                TestEntry { value: 200 },
                TestEntry { value: 300 },
            ])
            .unwrap();
        let got = storage.get_suffix(1).unwrap();
        assert_eq!(got.len(), 2);
        assert_eq!(got[0].value, 200);
    }

    #[test]
    fn append_on_prefix_truncates_then_appends() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![
                TestEntry { value: 1 },
                TestEntry { value: 2 },
                TestEntry { value: 3 },
            ])
            .unwrap();
        let len = storage
            .append_on_prefix(1, vec![TestEntry { value: 9 }, TestEntry { value: 8 }])
            .unwrap();
        assert_eq!(len, 3);
        let all = storage.get_suffix(0).unwrap();
        assert_eq!(all[0].value, 1);
        assert_eq!(all[1].value, 9);
        assert_eq!(all[2].value, 8);
    }
}

#[cfg(all(test, feature = "rocksdb-storage"))]
mod ballot_tests {
    use super::log_tests::fresh_storage;
    use omnipaxos::ballot_leader_election::Ballot;
    use omnipaxos::storage::Storage;
    use tempfile::TempDir;

    fn ballot(config_id: u32, n: u32, pid: u64) -> Ballot {
        Ballot {
            config_id,
            n,
            priority: 0,
            pid,
        }
    }

    #[test]
    fn promise_is_none_initially() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert!(storage.get_promise().unwrap().is_none());
    }

    #[test]
    fn set_promise_round_trip() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        let expected = ballot(1, 5, 2);
        storage.set_promise(expected).unwrap();
        let got = storage.get_promise().unwrap().expect("present");
        assert_eq!(got.n, expected.n);
        assert_eq!(got.pid, expected.pid);
        assert_eq!(got.config_id, expected.config_id);
    }

    #[test]
    fn accepted_round_is_none_initially() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert!(storage.get_accepted_round().unwrap().is_none());
    }

    #[test]
    fn set_accepted_round_round_trip() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        let expected = ballot(1, 7, 3);
        storage.set_accepted_round(expected).unwrap();
        let got = storage.get_accepted_round().unwrap().expect("present");
        assert_eq!(got.n, expected.n);
    }
}

#[cfg(all(test, feature = "rocksdb-storage"))]
mod decided_compacted_tests {
    use super::log_tests::fresh_storage;
    use super::open_in_tests::TestEntry;
    use omnipaxos::storage::Storage;
    use tempfile::TempDir;

    #[test]
    fn decided_idx_defaults_to_zero() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert_eq!(storage.get_decided_idx().unwrap(), 0);
    }

    #[test]
    fn set_decided_idx_round_trip() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage.set_decided_idx(42).unwrap();
        assert_eq!(storage.get_decided_idx().unwrap(), 42);
    }

    #[test]
    fn compacted_idx_defaults_to_zero() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert_eq!(storage.get_compacted_idx().unwrap(), 0);
    }

    #[test]
    fn set_compacted_idx_round_trip() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage.set_compacted_idx(17).unwrap();
        assert_eq!(storage.get_compacted_idx().unwrap(), 17);
    }

    #[test]
    fn append_floors_at_compacted_idx_raised_past_live_keys() {
        // Advancing compacted_idx past every live key WITHOUT trimming must
        // float the next write index up to the compaction floor: the next
        // append lands at `compacted_idx`, not just past the highest key.
        // Otherwise the entry would persist below the floor and be
        // unreachable via get_suffix(compacted_idx). (This is RocksdbStorage's
        // own contract — the `.max(compacted)` floor — and is intentionally a
        // Rocksdb-only test: MemStorage floors only on an empty log, so the
        // two impls diverge in this OmniPaxos-unreachable corner.)
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![TestEntry { value: 1 }, TestEntry { value: 2 }])
            .unwrap();
        storage.set_compacted_idx(10).unwrap();
        let new_len = storage.append_entry(TestEntry { value: 7 }).unwrap();
        assert_eq!(new_len, 1, "physical remaining = (10 + 1) - 10");
        let at_floor = storage.get_entries(10, 11).unwrap();
        assert_eq!(at_floor.len(), 1);
        assert_eq!(at_floor[0].value, 7, "append floors at compacted_idx = 10");
    }

    #[test]
    fn trim_then_compacted_yields_physical_remaining_length() {
        // trim takes an absolute index and deletes [0, idx). set_compacted_idx
        // is the paired call that advances the offset. After both, get_log_len
        // returns `next_abs - compacted = physical remaining`.
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![
                TestEntry { value: 1 },
                TestEntry { value: 2 },
                TestEntry { value: 3 },
                TestEntry { value: 4 },
            ])
            .unwrap();
        storage.trim(2).unwrap();
        storage.set_compacted_idx(2).unwrap();
        let suffix = storage.get_suffix(2).unwrap();
        assert_eq!(suffix.len(), 2);
        assert_eq!(suffix[0].value, 3);
        // Physical remaining = next_abs (4) - compacted (2) = 2.
        assert_eq!(storage.get_log_len().unwrap(), 2);
    }

    #[test]
    fn append_after_full_trim_writes_at_compacted_idx() {
        // After a snapshot trims every physical log key, the next append
        // must continue at the absolute index `compacted_idx`, not reset
        // to 0. Otherwise OmniPaxos's in-memory accepted_idx diverges
        // from the on-disk key space and the entry becomes unreadable
        // via get_suffix/get_entries.
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .append_entries(vec![
                TestEntry { value: 1 },
                TestEntry { value: 2 },
                TestEntry { value: 3 },
                TestEntry { value: 4 },
            ])
            .unwrap();
        storage.trim(4).unwrap();
        storage.set_compacted_idx(4).unwrap();

        let new_len = storage.append_entry(TestEntry { value: 99 }).unwrap();
        // physical remaining = (compacted + 1) - compacted = 1
        assert_eq!(new_len, 1);
        assert_eq!(storage.get_log_len().unwrap(), 1);

        let suffix = storage.get_suffix(4).unwrap();
        assert_eq!(suffix.len(), 1, "entry must be reachable at absolute idx 4");
        assert_eq!(suffix[0].value, 99);

        let range = storage.get_entries(4, 5).unwrap();
        assert_eq!(range.len(), 1);
        assert_eq!(range[0].value, 99);

        // The phantom-write bug stored the entry at L/0 even though
        // compacted_idx is 4, so a get_suffix(0) would surface it.
        // The fix means the entry lives at L/4, so get_suffix(0)
        // returns exactly the same single entry (no phantom at L/0).
        let from_zero = storage.get_suffix(0).unwrap();
        assert_eq!(from_zero.len(), 1);
    }

    #[test]
    fn append_after_full_trim_survives_reopen() {
        // Crash-recovery shape: after full compaction + append, drop the
        // storage and reopen on the same path. The persisted entry must
        // still live at the correct absolute index, and the next append
        // must continue from there.
        use super::open_in_tests::open_db;
        use crate::storage::RocksdbStorage;

        let dir = TempDir::new().unwrap();
        {
            let db = open_db(&dir, "tso_paxos");
            let mut storage: RocksdbStorage<TestEntry> =
                RocksdbStorage::open_in(db.clone(), "tso_paxos").unwrap();
            storage
                .append_entries(vec![
                    TestEntry { value: 10 },
                    TestEntry { value: 20 },
                    TestEntry { value: 30 },
                ])
                .unwrap();
            storage.trim(3).unwrap();
            storage.set_compacted_idx(3).unwrap();
            storage.append_entry(TestEntry { value: 77 }).unwrap();
            drop(storage);
            drop(db);
        }

        let db = open_db(&dir, "tso_paxos");
        let mut storage: RocksdbStorage<TestEntry> =
            RocksdbStorage::open_in(db, "tso_paxos").unwrap();
        assert_eq!(storage.get_compacted_idx().unwrap(), 3);
        let recovered = storage.get_suffix(3).unwrap();
        assert_eq!(recovered.len(), 1);
        assert_eq!(recovered[0].value, 77);

        let next_len = storage.append_entry(TestEntry { value: 88 }).unwrap();
        assert_eq!(next_len, 2);
        let recovered = storage.get_suffix(3).unwrap();
        assert_eq!(recovered.len(), 2);
        assert_eq!(recovered[0].value, 77);
        assert_eq!(recovered[1].value, 88);
        // The new entry must land at absolute idx 4, not L/1.
        let single = storage.get_entries(4, 5).unwrap();
        assert_eq!(single.len(), 1);
        assert_eq!(single[0].value, 88);
    }
}

#[cfg(all(test, feature = "rocksdb-storage"))]
mod crash_recovery_conformance {
    //! Crash-conformance for the snapshot/compaction sequence
    //! (`set_snapshot` / `trim` / `set_compacted_idx`).
    //!
    //! These writes are non-synced (`batch_async`), so a crash can lose a
    //! suffix of them. RocksDB's WAL is append-only and recovery replays a
    //! *prefix* of the issued writes — a crash can drop the tail but never
    //! resurrect an earlier write while losing a later one. Because OmniPaxos
    //! (and this crate's own pairing tests) issue `trim(idx)` before
    //! `set_compacted_idx(idx)`, the realistic partial-crash state is "trim
    //! persisted, compacted-index update lost", which
    //! `lost_compacted_index_after_trim_recovers_forward_only` in
    //! `tests/failpoints.rs` reproduces faithfully via the `async_write`
    //! failpoint.
    //!
    //! fsync itself is not in-process observable (a clean drop + reopen flushes
    //! everything), so a lost write is modeled by making that specific write
    //! fail, not by toggling the sync flag.

    use super::open_in_tests::{TestEntry, TestSnapshot, open_db};
    use super::*;
    use omnipaxos::storage::Storage;
    use tempfile::TempDir;

    #[test]
    fn full_ordered_compaction_sequence_survives_reopen() {
        // set_snapshot -> trim -> set_compacted_idx, all persisted, then reopen.
        // The snapshot, the compacted offset, and the surviving suffix must all
        // recover consistently, and the next append must continue past them.
        let dir = TempDir::new().unwrap();
        {
            let db = open_db(&dir, "tso_paxos");
            let mut storage: RocksdbStorage<TestEntry> =
                RocksdbStorage::open_in(db, "tso_paxos").unwrap();
            storage
                .append_entries(vec![
                    TestEntry { value: 10 },
                    TestEntry { value: 20 },
                    TestEntry { value: 30 },
                    TestEntry { value: 40 },
                ])
                .unwrap();
            storage
                .set_snapshot(Some(TestSnapshot { value: 30 }))
                .unwrap();
            storage.trim(3).unwrap();
            storage.set_compacted_idx(3).unwrap();
        }

        let db = open_db(&dir, "tso_paxos");
        let mut storage: RocksdbStorage<TestEntry> =
            RocksdbStorage::open_in(db, "tso_paxos").unwrap();
        assert_eq!(storage.get_compacted_idx().unwrap(), 3);
        assert_eq!(storage.get_snapshot().unwrap().expect("snapshot").value, 30);
        let suffix = storage.get_suffix(3).unwrap();
        assert_eq!(suffix.len(), 1);
        assert_eq!(
            suffix[0].value, 40,
            "the single surviving entry lives at idx 3"
        );

        let new_len = storage.append_entry(TestEntry { value: 99 }).unwrap();
        assert_eq!(new_len, 2, "physical remaining = (4+1) - 3");
        let single = storage.get_entries(4, 5).unwrap();
        assert_eq!(single.len(), 1);
        assert_eq!(
            single[0].value, 99,
            "next append continues at absolute idx 4"
        );
    }
}

#[cfg(all(test, feature = "rocksdb-storage"))]
mod snapshot_stopsign_tests {
    use super::log_tests::fresh_storage;
    use super::open_in_tests::TestSnapshot;
    use omnipaxos::ClusterConfig;
    use omnipaxos::storage::{StopSign, Storage};
    use tempfile::TempDir;

    fn stopsign() -> StopSign {
        StopSign::with(
            ClusterConfig {
                configuration_id: 2,
                nodes: vec![1, 2, 3],
                flexible_quorum: None,
            },
            None,
        )
    }

    #[test]
    fn snapshot_is_none_initially() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert!(storage.get_snapshot().unwrap().is_none());
    }

    #[test]
    fn set_snapshot_round_trip() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .set_snapshot(Some(TestSnapshot { value: 99 }))
            .unwrap();
        let got = storage.get_snapshot().unwrap().expect("present");
        assert_eq!(got.value, 99);
    }

    #[test]
    fn set_snapshot_none_clears() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage
            .set_snapshot(Some(TestSnapshot { value: 5 }))
            .unwrap();
        storage.set_snapshot(None).unwrap();
        assert!(storage.get_snapshot().unwrap().is_none());
    }

    #[test]
    fn stopsign_is_none_initially() {
        let dir = TempDir::new().unwrap();
        let storage = fresh_storage(&dir);
        assert!(storage.get_stopsign().unwrap().is_none());
    }

    #[test]
    fn set_stopsign_round_trip() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage.set_stopsign(Some(stopsign())).unwrap();
        let got = storage.get_stopsign().unwrap().expect("present");
        assert_eq!(got.next_config.configuration_id, 2);
        assert_eq!(got.next_config.nodes, vec![1, 2, 3]);
    }

    #[test]
    fn set_stopsign_none_clears() {
        let dir = TempDir::new().unwrap();
        let mut storage = fresh_storage(&dir);
        storage.set_stopsign(Some(stopsign())).unwrap();
        storage.set_stopsign(None).unwrap();
        assert!(storage.get_stopsign().unwrap().is_none());
    }
}