forge-orchestration 0.6.0

//! Consensus-backed [`StateStore`](crate::storage::StateStore) using
//! [openraft](https://docs.rs/openraft) 0.9.
//!
//! ## Table of Contents
//! - **RaftStateStore**: a [`StateStore`] whose `set`/`delete` are committed through Raft
//!   consensus (`Raft::client_write`) and whose `get`/`list_prefix` are served from the
//!   replicated in-memory key/value state machine.
//!
//! ## Scope (honest)
//! Writes always go through the real Raft `client_write` path (proposed -> appended to the log
//! -> committed -> applied to the state machine). Reads are served from the local applied state
//! machine. Three deployment shapes are supported:
//!
//! - **Single-node, in-memory** ([`RaftStateStore::bootstrap_single_node`], feature `raft`): the
//!   node bootstraps itself as the sole voter and elects itself leader. The loopback
//!   `RaftNetwork` is never invoked because there are no peers.
//! - **Single-node, crash-durable** ([`RaftStateStore::open_persistent`], feature `raft-persist`):
//!   the log + vote + applied state machine are persisted with [fjall](https://docs.rs/fjall)
//!   (pure Rust, no C toolchain). Committed keys survive a process restart; on reopen the node
//!   recovers from disk.
//! - **Multi-node over HTTP** ([`RaftStateStore::start_node`] +
//!   [`RaftStateStore::initialize_cluster`], feature `raft`): a real reqwest/axum transport
//!   replicates entries between peers, elects leaders, and survives leader failure.
//!
//! ## Storage internals
//! The in-memory variant keeps the log and state machine in `BTreeMap`s, modelled on the
//! official openraft `raft-kv-memstore` example for the v0.9 line. The persistent variant stores
//! the log (key = big-endian u64 index -> serialized `Entry`), vote/committed/purged pointers,
//! and a full state-machine snapshot in fjall partitions, fsyncing to disk before signalling that
//! a log append is durable. Both use a [`TypeConfig`] whose request is [`KvRequest`] (`Set`/`Delete`)
//! and whose response is [`KvResponse`] (the prior value).

use std::collections::BTreeMap;
use std::fmt;
use std::io::Cursor;
use std::ops::RangeBounds;
use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use openraft::storage::LogFlushed;
use openraft::storage::LogState;
use openraft::storage::RaftLogStorage;
use openraft::storage::RaftStateMachine;
use openraft::storage::Snapshot;
use openraft::Config;
use openraft::Entry;
use openraft::EntryPayload;
use openraft::LogId;
use openraft::OptionalSend;
use openraft::RaftLogReader;
use openraft::RaftSnapshotBuilder;
use openraft::SnapshotMeta;
use openraft::StorageError;
use openraft::StorageIOError;
use openraft::StoredMembership;
use openraft::Vote;
use serde::Deserialize;
use serde::Serialize;
use tokio::sync::RwLock;
use tracing::info;

use crate::error::{ForgeError, Result};
use crate::storage::StateStore;

#[cfg(feature = "raft-persist")]
use std::path::Path;

// ---------------------------------------------------------------------------
// Type config: request / response and the openraft `TypeConfig`.
// ---------------------------------------------------------------------------

/// A mutating request applied to the replicated key/value state machine.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum KvRequest {
    /// Insert or overwrite a key with the given value.
    Set {
        /// Key to write.
        key: String,
        /// Value bytes to store.
        value: Vec<u8>,
    },
    /// Remove a key.
    Delete {
        /// Key to remove.
        key: String,
    },
}

/// The response produced by applying a [`KvRequest`].
///
/// Carries the value that was present *before* the operation, mirroring the prior value so a
/// caller could implement compare-and-swap style semantics on top of it.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct KvResponse {
    /// The value previously stored under the key (if any).
    pub prev: Option<Vec<u8>>,
}

openraft::declare_raft_types!(
    /// openraft type configuration for the Forge key/value store.
    pub TypeConfig:
        D = KvRequest,
        R = KvResponse,
        NodeId = u64,
        Node = openraft::BasicNode,
        Entry = openraft::Entry<TypeConfig>,
        SnapshotData = Cursor<Vec<u8>>,
);

/// Convenience alias for this crate's openraft handle.
type ForgeRaft = openraft::Raft<TypeConfig>;

// ---------------------------------------------------------------------------
// State machine (replicated key/value map + snapshotting).
// ---------------------------------------------------------------------------

/// The serializable contents of the state machine, used both as the live state and as the
/// snapshot body.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct StateMachineData {
    last_applied: Option<LogId<u64>>,
    last_membership: StoredMembership<u64, openraft::BasicNode>,
    /// The replicated key/value map.
    kv: BTreeMap<String, Vec<u8>>,
}

/// A stored snapshot: metadata plus the serialized [`StateMachineData`].
#[derive(Debug, Clone)]
struct StoredSnapshot {
    meta: SnapshotMeta<u64, openraft::BasicNode>,
    data: Vec<u8>,
}

/// In-memory state machine. Shared (via `Arc`) so reads can observe applied state directly.
#[derive(Debug, Default)]
struct StateMachineStore {
    data: RwLock<StateMachineData>,
    /// Monotonic counter used to assign snapshot ids.
    snapshot_idx: RwLock<u64>,
    /// The most recently built/installed snapshot.
    current_snapshot: RwLock<Option<StoredSnapshot>>,
}

impl StateMachineStore {
    /// Read a key directly from the applied state machine.
    async fn get(&self, key: &str) -> Option<Vec<u8>> {
        let d = self.data.read().await;
        d.kv.get(key).cloned()
    }

    /// List keys with the given prefix from the applied state machine.
    async fn list_prefix(&self, prefix: &str) -> Vec<String> {
        let d = self.data.read().await;
        d.kv.range(prefix.to_string()..).take_while(|(k, _)| k.starts_with(prefix)).map(|(k, _)| k.clone()).collect()
    }
}

impl RaftSnapshotBuilder<TypeConfig> for Arc<StateMachineStore> {
    async fn build_snapshot(&mut self) -> std::result::Result<Snapshot<TypeConfig>, StorageError<u64>> {
        let (data, last_applied, last_membership) = {
            let d = self.data.read().await;
            (d.clone(), d.last_applied, d.last_membership.clone())
        };

        let bytes = serde_json::to_vec(&data).map_err(|e| StorageIOError::read_state_machine(&e))?;

        let snapshot_id = {
            let mut idx = self.snapshot_idx.write().await;
            *idx += 1;
            if let Some(last) = last_applied {
                format!("{}-{}-{}", last.leader_id, last.index, *idx)
            } else {
                format!("--{}", *idx)
            }
        };

        let meta = SnapshotMeta { last_log_id: last_applied, last_membership, snapshot_id };

        let stored = StoredSnapshot { meta: meta.clone(), data: bytes.clone() };
        *self.current_snapshot.write().await = Some(stored);

        Ok(Snapshot { meta, snapshot: Box::new(Cursor::new(bytes)) })
    }
}

impl RaftStateMachine<TypeConfig> for Arc<StateMachineStore> {
    type SnapshotBuilder = Arc<StateMachineStore>;

    async fn applied_state(
        &mut self,
    ) -> std::result::Result<(Option<LogId<u64>>, StoredMembership<u64, openraft::BasicNode>), StorageError<u64>> {
        let d = self.data.read().await;
        Ok((d.last_applied, d.last_membership.clone()))
    }

    async fn apply<I>(&mut self, entries: I) -> std::result::Result<Vec<KvResponse>, StorageError<u64>>
    where
        I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
        I::IntoIter: OptionalSend,
    {
        let mut d = self.data.write().await;
        let mut responses = Vec::new();

        for entry in entries {
            d.last_applied = Some(entry.log_id);

            match entry.payload {
                EntryPayload::Blank => {
                    responses.push(KvResponse::default());
                }
                EntryPayload::Normal(req) => {
                    let resp = match req {
                        KvRequest::Set { key, value } => {
                            let prev = d.kv.insert(key, value);
                            KvResponse { prev }
                        }
                        KvRequest::Delete { key } => {
                            let prev = d.kv.remove(&key);
                            KvResponse { prev }
                        }
                    };
                    responses.push(resp);
                }
                EntryPayload::Membership(mem) => {
                    d.last_membership = StoredMembership::new(Some(entry.log_id), mem);
                    responses.push(KvResponse::default());
                }
            }
        }

        Ok(responses)
    }

    async fn get_snapshot_builder(&mut self) -> Self::SnapshotBuilder {
        self.clone()
    }

    async fn begin_receiving_snapshot(
        &mut self,
    ) -> std::result::Result<Box<Cursor<Vec<u8>>>, StorageError<u64>> {
        Ok(Box::new(Cursor::new(Vec::new())))
    }

    async fn install_snapshot(
        &mut self,
        meta: &SnapshotMeta<u64, openraft::BasicNode>,
        snapshot: Box<Cursor<Vec<u8>>>,
    ) -> std::result::Result<(), StorageError<u64>> {
        let bytes = snapshot.into_inner();

        let new_data: StateMachineData =
            serde_json::from_slice(&bytes).map_err(|e| StorageIOError::read_snapshot(Some(meta.signature()), &e))?;

        {
            let mut d = self.data.write().await;
            *d = new_data;
        }

        *self.current_snapshot.write().await =
            Some(StoredSnapshot { meta: meta.clone(), data: bytes });

        Ok(())
    }

    async fn get_current_snapshot(
        &mut self,
    ) -> std::result::Result<Option<Snapshot<TypeConfig>>, StorageError<u64>> {
        let guard = self.current_snapshot.read().await;
        match &*guard {
            Some(s) => Ok(Some(Snapshot {
                meta: s.meta.clone(),
                snapshot: Box::new(Cursor::new(s.data.clone())),
            })),
            None => Ok(None),
        }
    }
}

// ---------------------------------------------------------------------------
// Log store (in-memory Raft log + vote).
// ---------------------------------------------------------------------------

/// In-memory Raft log store.
#[derive(Debug, Default)]
struct LogStoreInner {
    /// The Raft log, indexed by log index.
    log: BTreeMap<u64, Entry<TypeConfig>>,
    /// The last persisted vote.
    vote: Option<Vote<u64>>,
    /// The last committed log id (best-effort, persisted in memory).
    committed: Option<LogId<u64>>,
    /// The greatest log id that has been purged.
    last_purged_log_id: Option<LogId<u64>>,
}

/// Cloneable handle to the in-memory log store.
#[derive(Debug, Clone, Default)]
struct LogStore {
    inner: Arc<RwLock<LogStoreInner>>,
}

impl LogStore {
    async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + fmt::Debug + OptionalSend>(
        &self,
        range: RB,
    ) -> std::result::Result<Vec<Entry<TypeConfig>>, StorageError<u64>> {
        let inner = self.inner.read().await;
        Ok(inner.log.range(range).map(|(_, e)| e.clone()).collect())
    }
}

impl RaftLogReader<TypeConfig> for LogStore {
    async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + fmt::Debug + OptionalSend>(
        &mut self,
        range: RB,
    ) -> std::result::Result<Vec<Entry<TypeConfig>>, StorageError<u64>> {
        LogStore::try_get_log_entries(self, range).await
    }
}

impl RaftLogStorage<TypeConfig> for LogStore {
    type LogReader = LogStore;

    async fn get_log_state(&mut self) -> std::result::Result<LogState<TypeConfig>, StorageError<u64>> {
        let inner = self.inner.read().await;
        let last = inner.log.iter().next_back().map(|(_, e)| e.log_id);
        let last_log_id = match last {
            Some(id) => Some(id),
            None => inner.last_purged_log_id,
        };
        Ok(LogState { last_purged_log_id: inner.last_purged_log_id, last_log_id })
    }

    async fn get_log_reader(&mut self) -> Self::LogReader {
        self.clone()
    }

    async fn save_vote(&mut self, vote: &Vote<u64>) -> std::result::Result<(), StorageError<u64>> {
        let mut inner = self.inner.write().await;
        inner.vote = Some(*vote);
        Ok(())
    }

    async fn read_vote(&mut self) -> std::result::Result<Option<Vote<u64>>, StorageError<u64>> {
        let inner = self.inner.read().await;
        Ok(inner.vote)
    }

    async fn save_committed(
        &mut self,
        committed: Option<LogId<u64>>,
    ) -> std::result::Result<(), StorageError<u64>> {
        let mut inner = self.inner.write().await;
        inner.committed = committed;
        Ok(())
    }

    async fn read_committed(&mut self) -> std::result::Result<Option<LogId<u64>>, StorageError<u64>> {
        let inner = self.inner.read().await;
        Ok(inner.committed)
    }

    async fn append<I>(
        &mut self,
        entries: I,
        callback: LogFlushed<TypeConfig>,
    ) -> std::result::Result<(), StorageError<u64>>
    where
        I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
        I::IntoIter: OptionalSend,
    {
        {
            let mut inner = self.inner.write().await;
            for entry in entries {
                inner.log.insert(entry.log_id.index, entry);
            }
        }
        // Entries are durable (in memory) by the time we return; report completion.
        callback.log_io_completed(Ok(()));
        Ok(())
    }

    async fn truncate(&mut self, log_id: LogId<u64>) -> std::result::Result<(), StorageError<u64>> {
        let mut inner = self.inner.write().await;
        let keys: Vec<u64> = inner.log.range(log_id.index..).map(|(k, _)| *k).collect();
        for k in keys {
            inner.log.remove(&k);
        }
        Ok(())
    }

    async fn purge(&mut self, log_id: LogId<u64>) -> std::result::Result<(), StorageError<u64>> {
        let mut inner = self.inner.write().await;
        inner.last_purged_log_id = Some(log_id);
        let keys: Vec<u64> = inner.log.range(..=log_id.index).map(|(k, _)| *k).collect();
        for k in keys {
            inner.log.remove(&k);
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Persistent (fjall-backed) log store + state machine (feature `raft-persist`).
// ---------------------------------------------------------------------------
//
// Crash-durability strategy:
//   * A single fjall `Keyspace` holds three partitions: `logs` (key = big-endian u64
//     index -> serialized `Entry`), `meta` (vote / committed / last-purged), and `sm`
//     (the serialized state-machine snapshot).
//   * `append` writes the entries and then calls `keyspace.persist(SyncAll)` BEFORE
//     invoking the `LogFlushed` callback, so by the time openraft considers an entry
//     flushed it is genuinely fsync'd to disk.
//   * `save_vote` persists before returning (openraft requires the vote durable on return).
//   * The state machine persists a full snapshot of its applied state on every `apply`,
//     fsyncing to disk, so applied data survives a restart even with an empty/purged log.
//
// On reopen, `open_persistent` recovers the vote, committed/purged pointers, the log
// entries, and the state machine contents directly from the fjall partitions.

#[cfg(feature = "raft-persist")]
mod persist {
    use super::*;

    /// Partition name for serialized log entries (key = big-endian u64 index).
    pub(super) const TREE_LOGS: &str = "raft_logs";
    /// Partition name for vote / committed / last-purged pointers.
    pub(super) const TREE_META: &str = "raft_meta";
    /// Partition name for the persisted state-machine snapshot.
    pub(super) const TREE_SM: &str = "raft_sm";

    pub(super) const KEY_VOTE: &[u8] = b"vote";
    pub(super) const KEY_COMMITTED: &[u8] = b"committed";
    pub(super) const KEY_PURGED: &[u8] = b"last_purged";
    pub(super) const KEY_SM_DATA: &[u8] = b"sm_data";

    /// Map any fjall error into an openraft write `StorageError`.
    pub(super) fn fjall_write_err<E: std::error::Error + Send + Sync + 'static>(e: E) -> StorageError<u64> {
        StorageIOError::write(&e).into()
    }

    /// Map any fjall error into an openraft read-logs `StorageError`.
    pub(super) fn fjall_read_err<E: std::error::Error + Send + Sync + 'static>(e: E) -> StorageError<u64> {
        StorageIOError::read_logs(&e).into()
    }

    /// A crash-durable, fjall-backed Raft log store implementing both
    /// [`RaftLogReader`] and [`RaftLogStorage`].
    #[derive(Clone)]
    pub(super) struct FjallLogStore {
        pub(super) keyspace: fjall::Keyspace,
        pub(super) logs: fjall::PartitionHandle,
        pub(super) meta: fjall::PartitionHandle,
    }

    impl fmt::Debug for FjallLogStore {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("FjallLogStore").finish()
        }
    }

    impl FjallLogStore {
        pub(super) fn new(keyspace: fjall::Keyspace) -> std::result::Result<Self, fjall::Error> {
            let logs = keyspace.open_partition(TREE_LOGS, fjall::PartitionCreateOptions::default())?;
            let meta = keyspace.open_partition(TREE_META, fjall::PartitionCreateOptions::default())?;
            Ok(Self { keyspace, logs, meta })
        }

        async fn get_range<RB: RangeBounds<u64> + Clone + fmt::Debug + OptionalSend>(
            &self,
            range: RB,
        ) -> std::result::Result<Vec<Entry<TypeConfig>>, StorageError<u64>> {
            use std::ops::Bound;

            // Translate the u64 bounds into big-endian byte key bounds for fjall.
            let start: Bound<[u8; 8]> = match range.start_bound() {
                Bound::Included(i) => Bound::Included(i.to_be_bytes()),
                Bound::Excluded(i) => Bound::Excluded(i.to_be_bytes()),
                Bound::Unbounded => Bound::Unbounded,
            };
            let end: Bound<[u8; 8]> = match range.end_bound() {
                Bound::Included(i) => Bound::Included(i.to_be_bytes()),
                Bound::Excluded(i) => Bound::Excluded(i.to_be_bytes()),
                Bound::Unbounded => Bound::Unbounded,
            };

            let mut out = Vec::new();
            for kv in self.logs.range((start, end)) {
                let (_k, v) = kv.map_err(fjall_read_err)?;
                let entry: Entry<TypeConfig> =
                    serde_json::from_slice(&v).map_err(|e| StorageIOError::read_logs(&e))?;
                out.push(entry);
            }
            Ok(out)
        }
    }

    impl RaftLogReader<TypeConfig> for FjallLogStore {
        async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + fmt::Debug + OptionalSend>(
            &mut self,
            range: RB,
        ) -> std::result::Result<Vec<Entry<TypeConfig>>, StorageError<u64>> {
            self.get_range(range).await
        }
    }

    impl RaftLogStorage<TypeConfig> for FjallLogStore {
        type LogReader = FjallLogStore;

        async fn get_log_state(&mut self) -> std::result::Result<LogState<TypeConfig>, StorageError<u64>> {
            let last_purged_log_id: Option<LogId<u64>> = match self.meta.get(KEY_PURGED).map_err(fjall_read_err)? {
                Some(v) => serde_json::from_slice(&v).map_err(|e| StorageIOError::read_logs(&e))?,
                None => None,
            };

            // The last present log entry, if any.
            let last_in_log: Option<LogId<u64>> = match self.logs.last_key_value().map_err(fjall_read_err)? {
                Some((_k, v)) => {
                    let entry: Entry<TypeConfig> =
                        serde_json::from_slice(&v).map_err(|e| StorageIOError::read_logs(&e))?;
                    Some(entry.log_id)
                }
                None => None,
            };

            let last_log_id = last_in_log.or(last_purged_log_id);
            Ok(LogState { last_purged_log_id, last_log_id })
        }

        async fn get_log_reader(&mut self) -> Self::LogReader {
            self.clone()
        }

        async fn save_vote(&mut self, vote: &Vote<u64>) -> std::result::Result<(), StorageError<u64>> {
            let bytes = serde_json::to_vec(vote).map_err(|e| StorageIOError::write_vote(&e))?;
            self.meta.insert(KEY_VOTE, bytes).map_err(fjall_write_err)?;
            // The vote MUST be durable before returning.
            self.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_write_err)?;
            Ok(())
        }

        async fn read_vote(&mut self) -> std::result::Result<Option<Vote<u64>>, StorageError<u64>> {
            match self.meta.get(KEY_VOTE).map_err(fjall_read_err)? {
                Some(v) => Ok(Some(serde_json::from_slice(&v).map_err(|e| StorageIOError::read_vote(&e))?)),
                None => Ok(None),
            }
        }

        async fn save_committed(
            &mut self,
            committed: Option<LogId<u64>>,
        ) -> std::result::Result<(), StorageError<u64>> {
            let bytes = serde_json::to_vec(&committed).map_err(|e| StorageIOError::write(&e))?;
            self.meta.insert(KEY_COMMITTED, bytes).map_err(fjall_write_err)?;
            self.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_write_err)?;
            Ok(())
        }

        async fn read_committed(&mut self) -> std::result::Result<Option<LogId<u64>>, StorageError<u64>> {
            match self.meta.get(KEY_COMMITTED).map_err(fjall_read_err)? {
                Some(v) => Ok(serde_json::from_slice(&v).map_err(|e| StorageIOError::read(&e))?),
                None => Ok(None),
            }
        }

        async fn append<I>(
            &mut self,
            entries: I,
            callback: LogFlushed<TypeConfig>,
        ) -> std::result::Result<(), StorageError<u64>>
        where
            I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
            I::IntoIter: OptionalSend,
        {
            for entry in entries {
                let key = entry.log_id.index.to_be_bytes();
                let val = serde_json::to_vec(&entry).map_err(|e| StorageIOError::write_logs(&e))?;
                self.logs.insert(key, val).map_err(fjall_write_err)?;
            }
            // CRITICAL: fsync to disk BEFORE signalling completion so a committed entry is durable.
            self.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_write_err)?;
            callback.log_io_completed(Ok(()));
            Ok(())
        }

        async fn truncate(&mut self, log_id: LogId<u64>) -> std::result::Result<(), StorageError<u64>> {
            // Remove all entries with index >= log_id.index.
            let from = log_id.index.to_be_bytes();
            let keys: Vec<fjall::Slice> = self
                .logs
                .range(from..)
                .map(|kv| kv.map(|(k, _)| k))
                .collect::<std::result::Result<_, _>>()
                .map_err(fjall_write_err)?;
            for k in keys {
                self.logs.remove(k).map_err(fjall_write_err)?;
            }
            self.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_write_err)?;
            Ok(())
        }

        async fn purge(&mut self, log_id: LogId<u64>) -> std::result::Result<(), StorageError<u64>> {
            // Record the new purge watermark, then drop all entries with index <= log_id.index.
            let bytes = serde_json::to_vec(&Some(log_id)).map_err(|e| StorageIOError::write(&e))?;
            self.meta.insert(KEY_PURGED, bytes).map_err(fjall_write_err)?;

            let to = log_id.index.to_be_bytes();
            let keys: Vec<fjall::Slice> = self
                .logs
                .range(..=to)
                .map(|kv| kv.map(|(k, _)| k))
                .collect::<std::result::Result<_, _>>()
                .map_err(fjall_write_err)?;
            for k in keys {
                self.logs.remove(k).map_err(fjall_write_err)?;
            }
            self.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_write_err)?;
            Ok(())
        }
    }

    /// A crash-durable, fjall-backed state machine. Applied state and the current snapshot
    /// are persisted to disk; reads are served from an in-memory mirror for speed.
    #[derive(Clone)]
    pub(super) struct FjallStateMachine {
        pub(super) inner: Arc<FjallStateMachineInner>,
    }

    pub(super) struct FjallStateMachineInner {
        /// In-memory mirror of the applied state (kept in sync with disk).
        pub(super) data: RwLock<StateMachineData>,
        /// Monotonic snapshot id counter.
        pub(super) snapshot_idx: RwLock<u64>,
        /// The most recently built/installed snapshot (also persisted on disk).
        pub(super) current_snapshot: RwLock<Option<StoredSnapshot>>,
        /// fjall handles.
        pub(super) keyspace: fjall::Keyspace,
        pub(super) sm: fjall::PartitionHandle,
    }

    impl fmt::Debug for FjallStateMachine {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("FjallStateMachine").finish()
        }
    }

    impl FjallStateMachine {
        pub(super) fn new(keyspace: fjall::Keyspace) -> std::result::Result<Self, fjall::Error> {
            let sm = keyspace.open_partition(TREE_SM, fjall::PartitionCreateOptions::default())?;

            // Recover applied state from disk, if present.
            let data: StateMachineData = match sm.get(KEY_SM_DATA)? {
                Some(v) => serde_json::from_slice(&v).unwrap_or_default(),
                None => StateMachineData::default(),
            };
            // Recover the persisted snapshot, if present.
            let current_snapshot = build_recovered_snapshot(&data);

            Ok(Self {
                inner: Arc::new(FjallStateMachineInner {
                    data: RwLock::new(data),
                    snapshot_idx: RwLock::new(0),
                    current_snapshot: RwLock::new(current_snapshot),
                    keyspace,
                    sm,
                }),
            })
        }

        /// Read a key from the applied state machine (in-memory mirror).
        pub(super) async fn get(&self, key: &str) -> Option<Vec<u8>> {
            let d = self.inner.data.read().await;
            d.kv.get(key).cloned()
        }

        /// List keys with the given prefix.
        pub(super) async fn list_prefix(&self, prefix: &str) -> Vec<String> {
            let d = self.inner.data.read().await;
            d.kv.range(prefix.to_string()..).take_while(|(k, _)| k.starts_with(prefix)).map(|(k, _)| k.clone()).collect()
        }

        /// Persist the current applied state to disk and fsync.
        async fn persist_data(&self, d: &StateMachineData) -> std::result::Result<(), StorageError<u64>> {
            let bytes = serde_json::to_vec(d).map_err(|e| StorageIOError::write_state_machine(&e))?;
            self.inner.sm.insert(KEY_SM_DATA, bytes).map_err(fjall_write_err)?;
            self.inner.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_write_err)?;
            Ok(())
        }
    }

    /// Build a [`StoredSnapshot`] from recovered state-machine data (used on reopen).
    fn build_recovered_snapshot(data: &StateMachineData) -> Option<StoredSnapshot> {
        let bytes = serde_json::to_vec(data).ok()?;
        let snapshot_id = match data.last_applied {
            Some(last) => format!("{}-{}-recovered", last.leader_id, last.index),
            None => return None,
        };
        let meta = SnapshotMeta {
            last_log_id: data.last_applied,
            last_membership: data.last_membership.clone(),
            snapshot_id,
        };
        Some(StoredSnapshot { meta, data: bytes })
    }

    impl RaftSnapshotBuilder<TypeConfig> for FjallStateMachine {
        async fn build_snapshot(&mut self) -> std::result::Result<Snapshot<TypeConfig>, StorageError<u64>> {
            let (data, last_applied, last_membership) = {
                let d = self.inner.data.read().await;
                (d.clone(), d.last_applied, d.last_membership.clone())
            };

            let bytes = serde_json::to_vec(&data).map_err(|e| StorageIOError::read_state_machine(&e))?;

            let snapshot_id = {
                let mut idx = self.inner.snapshot_idx.write().await;
                *idx += 1;
                if let Some(last) = last_applied {
                    format!("{}-{}-{}", last.leader_id, last.index, *idx)
                } else {
                    format!("--{}", *idx)
                }
            };

            let meta = SnapshotMeta { last_log_id: last_applied, last_membership, snapshot_id };
            let stored = StoredSnapshot { meta: meta.clone(), data: bytes.clone() };
            *self.inner.current_snapshot.write().await = Some(stored);

            Ok(Snapshot { meta, snapshot: Box::new(Cursor::new(bytes)) })
        }
    }

    impl RaftStateMachine<TypeConfig> for FjallStateMachine {
        type SnapshotBuilder = FjallStateMachine;

        async fn applied_state(
            &mut self,
        ) -> std::result::Result<(Option<LogId<u64>>, StoredMembership<u64, openraft::BasicNode>), StorageError<u64>>
        {
            let d = self.inner.data.read().await;
            Ok((d.last_applied, d.last_membership.clone()))
        }

        async fn apply<I>(&mut self, entries: I) -> std::result::Result<Vec<KvResponse>, StorageError<u64>>
        where
            I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
            I::IntoIter: OptionalSend,
        {
            let mut responses = Vec::new();
            let snapshot_after = {
                let mut d = self.inner.data.write().await;
                for entry in entries {
                    d.last_applied = Some(entry.log_id);
                    match entry.payload {
                        EntryPayload::Blank => responses.push(KvResponse::default()),
                        EntryPayload::Normal(req) => {
                            let resp = match req {
                                KvRequest::Set { key, value } => KvResponse { prev: d.kv.insert(key, value) },
                                KvRequest::Delete { key } => KvResponse { prev: d.kv.remove(&key) },
                            };
                            responses.push(resp);
                        }
                        EntryPayload::Membership(mem) => {
                            d.last_membership = StoredMembership::new(Some(entry.log_id), mem);
                            responses.push(KvResponse::default());
                        }
                    }
                }
                d.clone()
            };
            // Persist applied state to disk before returning so it survives a restart.
            self.persist_data(&snapshot_after).await?;
            Ok(responses)
        }

        async fn get_snapshot_builder(&mut self) -> Self::SnapshotBuilder {
            self.clone()
        }

        async fn begin_receiving_snapshot(
            &mut self,
        ) -> std::result::Result<Box<Cursor<Vec<u8>>>, StorageError<u64>> {
            Ok(Box::new(Cursor::new(Vec::new())))
        }

        async fn install_snapshot(
            &mut self,
            meta: &SnapshotMeta<u64, openraft::BasicNode>,
            snapshot: Box<Cursor<Vec<u8>>>,
        ) -> std::result::Result<(), StorageError<u64>> {
            let bytes = snapshot.into_inner();
            let new_data: StateMachineData = serde_json::from_slice(&bytes)
                .map_err(|e| StorageIOError::read_snapshot(Some(meta.signature()), &e))?;

            {
                let mut d = self.inner.data.write().await;
                *d = new_data.clone();
            }
            self.persist_data(&new_data).await?;
            *self.inner.current_snapshot.write().await =
                Some(StoredSnapshot { meta: meta.clone(), data: bytes });
            Ok(())
        }

        async fn get_current_snapshot(
            &mut self,
        ) -> std::result::Result<Option<Snapshot<TypeConfig>>, StorageError<u64>> {
            let guard = self.inner.current_snapshot.read().await;
            match &*guard {
                Some(s) => Ok(Some(Snapshot {
                    meta: s.meta.clone(),
                    snapshot: Box::new(Cursor::new(s.data.clone())),
                })),
                None => Ok(None),
            }
        }
    }
}

// ---------------------------------------------------------------------------
// No-op loopback network (single-node only).
// ---------------------------------------------------------------------------

/// A network factory that produces loopback connections. In a single-node cluster these
/// connections are never used because there are no peers to contact.
#[derive(Debug, Clone, Default)]
struct LoopbackNetworkFactory;

/// A single loopback connection. All RPC methods report the target as unreachable, which is
/// correct: in a single-node cluster openraft never sends RPCs.
#[derive(Debug, Clone)]
struct LoopbackNetwork {
    target: u64,
}

impl openraft::network::RaftNetworkFactory<TypeConfig> for LoopbackNetworkFactory {
    type Network = LoopbackNetwork;

    async fn new_client(&mut self, target: u64, _node: &openraft::BasicNode) -> Self::Network {
        LoopbackNetwork { target }
    }
}

impl openraft::network::RaftNetwork<TypeConfig> for LoopbackNetwork {
    async fn append_entries(
        &mut self,
        _rpc: openraft::raft::AppendEntriesRequest<TypeConfig>,
        _option: openraft::network::RPCOption,
    ) -> std::result::Result<
        openraft::raft::AppendEntriesResponse<u64>,
        openraft::error::RPCError<u64, openraft::BasicNode, openraft::error::RaftError<u64>>,
    > {
        Err(unreachable_rpc(self.target))
    }

    async fn install_snapshot(
        &mut self,
        _rpc: openraft::raft::InstallSnapshotRequest<TypeConfig>,
        _option: openraft::network::RPCOption,
    ) -> std::result::Result<
        openraft::raft::InstallSnapshotResponse<u64>,
        openraft::error::RPCError<
            u64,
            openraft::BasicNode,
            openraft::error::RaftError<u64, openraft::error::InstallSnapshotError>,
        >,
    > {
        Err(unreachable_rpc(self.target))
    }

    async fn vote(
        &mut self,
        _rpc: openraft::raft::VoteRequest<u64>,
        _option: openraft::network::RPCOption,
    ) -> std::result::Result<
        openraft::raft::VoteResponse<u64>,
        openraft::error::RPCError<u64, openraft::BasicNode, openraft::error::RaftError<u64>>,
    > {
        Err(unreachable_rpc(self.target))
    }
}

/// Build an `Unreachable` RPC error for the given target node.
fn unreachable_rpc<E>(target: u64) -> openraft::error::RPCError<u64, openraft::BasicNode, E>
where
    E: std::error::Error,
{
    openraft::error::RPCError::Unreachable(openraft::error::Unreachable::new(
        &std::io::Error::new(std::io::ErrorKind::NotConnected, format!("no peer {target} in single-node cluster")),
    ))
}

// ---------------------------------------------------------------------------
// Real multi-node HTTP transport (feature `raft`).
// ---------------------------------------------------------------------------
//
// A node exposes three POST endpoints (`/raft/append`, `/raft/vote`, `/raft/snapshot`) served by
// an axum router. Each handler deserializes the openraft RPC request, calls the LOCAL Raft handle
// (`append_entries` / `vote` / `install_snapshot`), and returns the serialized response.
//
// On the client side, `HttpNetworkFactory` creates one `HttpNetwork` per peer. The peer's HTTP
// base URL is taken from `BasicNode::addr` (e.g. "127.0.0.1:9001"). RPCs POST the bincode-free
// JSON-serialized request and deserialize the JSON response. Connection failures map to
// `Unreachable` so openraft backs off and retries (important during elections / when a node is
// briefly down).

mod http_transport {
    use super::*;
    use axum::extract::State;
    use axum::http::StatusCode;
    use axum::response::IntoResponse;
    use axum::routing::post;
    use axum::Router;
    use openraft::error::{NetworkError, RPCError, Unreachable};
    use openraft::network::{RPCOption, RaftNetwork, RaftNetworkFactory};
    use openraft::raft::{
        AppendEntriesRequest, AppendEntriesResponse, InstallSnapshotRequest, InstallSnapshotResponse,
        VoteRequest, VoteResponse,
    };

    /// Shared axum state: the local Raft handle used to serve incoming RPCs.
    #[derive(Clone)]
    pub(super) struct RaftHttpState {
        pub(super) raft: ForgeRaft,
    }

    /// Build the axum router that serves this node's Raft RPC endpoints.
    pub(super) fn router(raft: ForgeRaft) -> Router {
        Router::new()
            .route("/raft/append", post(append_handler))
            .route("/raft/vote", post(vote_handler))
            .route("/raft/snapshot", post(snapshot_handler))
            .with_state(RaftHttpState { raft })
    }

    /// Generic "deserialize body, run, serialize response" helper for the handlers.
    async fn run_rpc<Req, Resp, Fut>(
        body: axum::body::Bytes,
        f: impl FnOnce(Req) -> Fut,
    ) -> std::result::Result<Vec<u8>, (StatusCode, String)>
    where
        Req: for<'de> Deserialize<'de>,
        Resp: Serialize,
        Fut: std::future::Future<Output = std::result::Result<Resp, String>>,
    {
        let req: Req = serde_json::from_slice(&body)
            .map_err(|e| (StatusCode::BAD_REQUEST, format!("decode request: {e}")))?;
        let resp = f(req).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
        serde_json::to_vec(&resp).map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("encode response: {e}")))
    }

    async fn append_handler(
        State(st): State<RaftHttpState>,
        body: axum::body::Bytes,
    ) -> impl IntoResponse {
        match run_rpc::<AppendEntriesRequest<TypeConfig>, AppendEntriesResponse<u64>, _>(body, |rpc| async move {
            st.raft.append_entries(rpc).await.map_err(|e| e.to_string())
        })
        .await
        {
            Ok(bytes) => (StatusCode::OK, bytes).into_response(),
            Err((code, msg)) => (code, msg).into_response(),
        }
    }

    async fn vote_handler(State(st): State<RaftHttpState>, body: axum::body::Bytes) -> impl IntoResponse {
        match run_rpc::<VoteRequest<u64>, VoteResponse<u64>, _>(body, |rpc| async move {
            st.raft.vote(rpc).await.map_err(|e| e.to_string())
        })
        .await
        {
            Ok(bytes) => (StatusCode::OK, bytes).into_response(),
            Err((code, msg)) => (code, msg).into_response(),
        }
    }

    async fn snapshot_handler(State(st): State<RaftHttpState>, body: axum::body::Bytes) -> impl IntoResponse {
        match run_rpc::<InstallSnapshotRequest<TypeConfig>, InstallSnapshotResponse<u64>, _>(body, |rpc| async move {
            #[allow(deprecated)]
            st.raft.install_snapshot(rpc).await.map_err(|e| e.to_string())
        })
        .await
        {
            Ok(bytes) => (StatusCode::OK, bytes).into_response(),
            Err((code, msg)) => (code, msg).into_response(),
        }
    }

    /// Factory that produces one [`HttpNetwork`] client per peer node.
    #[derive(Clone)]
    pub(super) struct HttpNetworkFactory {
        client: reqwest::Client,
    }

    impl HttpNetworkFactory {
        pub(super) fn new() -> Self {
            let client = reqwest::Client::builder()
                .timeout(Duration::from_secs(3))
                .build()
                .expect("reqwest client");
            Self { client }
        }
    }

    /// A client connection to a single peer node, addressed by its HTTP base URL.
    #[derive(Clone)]
    pub(super) struct HttpNetwork {
        client: reqwest::Client,
        target: u64,
        base_url: String,
    }

    impl HttpNetwork {
        /// POST `req` (JSON) to `/{path}` on the peer and deserialize the JSON response.
        async fn post<Req, Resp, E>(
            &self,
            path: &str,
            req: &Req,
        ) -> std::result::Result<Resp, RPCError<u64, openraft::BasicNode, E>>
        where
            Req: Serialize,
            Resp: for<'de> Deserialize<'de>,
            E: std::error::Error,
        {
            let url = format!("{}/{}", self.base_url, path);
            let body = serde_json::to_vec(req).map_err(|e| RPCError::Network(NetworkError::new(&e)))?;

            let resp = self
                .client
                .post(&url)
                .header("content-type", "application/json")
                .body(body)
                .send()
                .await
                // A connection failure is transient (peer down / not yet up): Unreachable so
                // openraft applies backoff and retries rather than treating it as fatal.
                .map_err(|e| {
                    RPCError::Unreachable(Unreachable::new(&std::io::Error::new(
                        std::io::ErrorKind::NotConnected,
                        format!("peer {} ({}): {e}", self.target, self.base_url),
                    )))
                })?;

            if !resp.status().is_success() {
                let code = resp.status();
                let text = resp.text().await.unwrap_or_default();
                return Err(RPCError::Network(NetworkError::new(&std::io::Error::other(format!(
                    "peer {} returned {code}: {text}",
                    self.target
                )))));
            }

            let bytes = resp
                .bytes()
                .await
                .map_err(|e| RPCError::Network(NetworkError::new(&e)))?;
            serde_json::from_slice(&bytes).map_err(|e| RPCError::Network(NetworkError::new(&e)))
        }
    }

    impl RaftNetworkFactory<TypeConfig> for HttpNetworkFactory {
        type Network = HttpNetwork;

        async fn new_client(&mut self, target: u64, node: &openraft::BasicNode) -> Self::Network {
            // `node.addr` is the peer's "host:port"; build a base URL from it.
            let base_url = if node.addr.starts_with("http://") || node.addr.starts_with("https://") {
                node.addr.clone()
            } else {
                format!("http://{}", node.addr)
            };
            HttpNetwork { client: self.client.clone(), target, base_url }
        }
    }

    impl RaftNetwork<TypeConfig> for HttpNetwork {
        async fn append_entries(
            &mut self,
            rpc: AppendEntriesRequest<TypeConfig>,
            _option: RPCOption,
        ) -> std::result::Result<
            AppendEntriesResponse<u64>,
            RPCError<u64, openraft::BasicNode, openraft::error::RaftError<u64>>,
        > {
            self.post("raft/append", &rpc).await
        }

        async fn install_snapshot(
            &mut self,
            rpc: InstallSnapshotRequest<TypeConfig>,
            _option: RPCOption,
        ) -> std::result::Result<
            InstallSnapshotResponse<u64>,
            RPCError<u64, openraft::BasicNode, openraft::error::RaftError<u64, openraft::error::InstallSnapshotError>>,
        > {
            self.post("raft/snapshot", &rpc).await
        }

        async fn vote(
            &mut self,
            rpc: VoteRequest<u64>,
            _option: RPCOption,
        ) -> std::result::Result<
            VoteResponse<u64>,
            RPCError<u64, openraft::BasicNode, openraft::error::RaftError<u64>>,
        > {
            self.post("raft/vote", &rpc).await
        }
    }
}

// ---------------------------------------------------------------------------
// RaftStateStore: the public StateStore implementation.
// ---------------------------------------------------------------------------

/// The read path for a [`RaftStateStore`]: either the in-memory state machine or, when the
/// `raft-persist` feature is used via [`RaftStateStore::open_persistent`], the fjall-backed one.
/// Reads (`get`/`list_prefix`) are served from the locally-applied state machine in both cases.
enum ReadPath {
    /// In-memory state machine (single-node bootstrap, multi-node HTTP cluster).
    Memory(Arc<StateMachineStore>),
    /// Disk-backed state machine (feature `raft-persist`).
    #[cfg(feature = "raft-persist")]
    Persistent(persist::FjallStateMachine),
}

impl ReadPath {
    async fn get(&self, key: &str) -> Option<Vec<u8>> {
        match self {
            ReadPath::Memory(sm) => sm.get(key).await,
            #[cfg(feature = "raft-persist")]
            ReadPath::Persistent(sm) => sm.get(key).await,
        }
    }

    async fn list_prefix(&self, prefix: &str) -> Vec<String> {
        match self {
            ReadPath::Memory(sm) => sm.list_prefix(prefix).await,
            #[cfg(feature = "raft-persist")]
            ReadPath::Persistent(sm) => sm.list_prefix(prefix).await,
        }
    }
}

/// A [`StateStore`] backed by an openraft consensus group.
///
/// `set` and `delete` are committed through Raft (`client_write`); `get` and `list_prefix`
/// read from the locally-applied state machine. See the module docs for the consistency model.
///
/// Construct one of three ways:
/// - [`RaftStateStore::bootstrap_single_node`] — single in-memory voter (default `raft`).
/// - [`RaftStateStore::open_persistent`] — single voter with a crash-durable fjall log + snapshot
///   that survives process restart (feature `raft-persist`).
/// - [`RaftStateStore::start_node`] + [`RaftStateStore::initialize_cluster`] — a multi-node cluster
///   over real HTTP transport (feature `raft`).
pub struct RaftStateStore {
    node_id: u64,
    raft: ForgeRaft,
    read_path: ReadPath,
    /// For a multi-node HTTP node, the background axum server task (aborted on shutdown).
    server_task: Option<tokio::task::JoinHandle<()>>,
    /// For a multi-node HTTP node, the actual bound local address (useful for ephemeral ports).
    local_addr: Option<std::net::SocketAddr>,
}

impl fmt::Debug for RaftStateStore {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RaftStateStore").field("node_id", &self.node_id).finish()
    }
}

impl RaftStateStore {
    /// Bootstrap a single-node Raft cluster and return a ready-to-use store.
    ///
    /// This creates the Raft node, initializes it as the sole voter, and waits until the node
    /// has elected itself leader (by polling metrics, not by blind sleeping). The returned
    /// store is immediately usable for `set`/`get`/`delete`/`list_prefix`.
    pub async fn bootstrap_single_node(node_id: u64) -> Result<Self> {
        // Short timeouts so a single node elects itself promptly in tests and startup.
        let config = Config {
            cluster_name: "forge".to_string(),
            election_timeout_min: 150,
            election_timeout_max: 300,
            heartbeat_interval: 50,
            ..Default::default()
        };
        let config = Arc::new(config.validate().map_err(|e| ForgeError::Consensus(format!("invalid raft config: {e}")))?);

        let log_store = LogStore::default();
        let state_machine = Arc::new(StateMachineStore::default());
        let network = LoopbackNetworkFactory;

        let raft = openraft::Raft::new(node_id, config, network, log_store, state_machine.clone())
            .await
            .map_err(|e| ForgeError::Consensus(format!("failed to create raft node: {e}")))?;

        // Bootstrap: make this node the sole voter in the cluster.
        let mut members = BTreeMap::new();
        members.insert(node_id, openraft::BasicNode::default());

        match raft.initialize(members).await {
            Ok(()) => {}
            Err(openraft::error::RaftError::APIError(openraft::error::InitializeError::NotAllowed(_))) => {
                // Already initialized (e.g. recovered) -- that's fine.
            }
            Err(e) => {
                return Err(ForgeError::Consensus(format!("failed to initialize raft cluster: {e}")));
            }
        }

        let store = Self { node_id, raft, read_path: ReadPath::Memory(state_machine), server_task: None, local_addr: None };
        store.wait_for_leadership().await?;
        info!(node_id, "RaftStateStore bootstrapped and leader elected");
        Ok(store)
    }

    /// Open (or create) a crash-durable, disk-backed single-node Raft store rooted at `dir`.
    ///
    /// The log and snapshot are persisted with [fjall](https://docs.rs/fjall) (pure Rust, no C
    /// toolchain). On first open the node bootstraps itself as the sole voter. On a subsequent
    /// reopen of the same `dir`, the vote, log, and applied state machine are recovered from disk
    /// (so previously-committed keys survive a process restart); openraft replays as needed and
    /// the node resumes leadership without re-initializing.
    ///
    /// Requires the `raft-persist` feature.
    #[cfg(feature = "raft-persist")]
    pub async fn open_persistent(node_id: u64, dir: &Path) -> Result<Self> {
        let config = Config {
            cluster_name: "forge".to_string(),
            election_timeout_min: 150,
            election_timeout_max: 300,
            heartbeat_interval: 50,
            ..Default::default()
        };
        let config = Arc::new(config.validate().map_err(|e| ForgeError::Consensus(format!("invalid raft config: {e}")))?);

        let keyspace = fjall::Config::new(dir)
            .open()
            .map_err(|e| ForgeError::Consensus(format!("failed to open fjall keyspace at {dir:?}: {e}")))?;

        let log_store = persist::FjallLogStore::new(keyspace.clone())
            .map_err(|e| ForgeError::Consensus(format!("failed to open fjall log store: {e}")))?;
        let state_machine = persist::FjallStateMachine::new(keyspace.clone())
            .map_err(|e| ForgeError::Consensus(format!("failed to open fjall state machine: {e}")))?;
        let network = LoopbackNetworkFactory;

        let raft = openraft::Raft::new(node_id, config, network, log_store, state_machine.clone())
            .await
            .map_err(|e| ForgeError::Consensus(format!("failed to create raft node: {e}")))?;

        // Bootstrap on first open; a recovered node is already initialized.
        let mut members = BTreeMap::new();
        members.insert(node_id, openraft::BasicNode::default());
        match raft.initialize(members).await {
            Ok(()) => {}
            Err(openraft::error::RaftError::APIError(openraft::error::InitializeError::NotAllowed(_))) => {
                // Already initialized from a prior run -- recovered from disk.
            }
            Err(e) => {
                return Err(ForgeError::Consensus(format!("failed to initialize raft cluster: {e}")));
            }
        }

        let store = Self { node_id, raft, read_path: ReadPath::Persistent(state_machine), server_task: None, local_addr: None };
        store.wait_for_leadership().await?;
        info!(node_id, ?dir, "RaftStateStore opened (persistent) and leader elected");
        Ok(store)
    }

    /// Start a multi-node Raft node: create the Raft instance with the real HTTP transport and
    /// spawn an axum server bound to `bind_addr` that serves this node's `/raft/*` RPC endpoints.
    ///
    /// This does NOT initialize a cluster. After starting all nodes, call
    /// [`RaftStateStore::initialize_cluster`] on exactly one of them with the full member set
    /// (node id -> "host:port"). Requires the `raft` feature.
    ///
    /// The node listens on `bind_addr` but advertises itself to peers via the address in the
    /// member map passed to `initialize_cluster` / `add_learner`.
    pub async fn start_node(node_id: u64, bind_addr: std::net::SocketAddr) -> Result<Self> {
        let config = Config {
            cluster_name: "forge".to_string(),
            // Slightly longer, well-separated election window for multi-node stability.
            election_timeout_min: 300,
            election_timeout_max: 600,
            heartbeat_interval: 100,
            ..Default::default()
        };
        let config = Arc::new(config.validate().map_err(|e| ForgeError::Consensus(format!("invalid raft config: {e}")))?);

        let log_store = LogStore::default();
        let state_machine = Arc::new(StateMachineStore::default());
        let network = http_transport::HttpNetworkFactory::new();

        let raft = openraft::Raft::new(node_id, config, network, log_store, state_machine.clone())
            .await
            .map_err(|e| ForgeError::Consensus(format!("failed to create raft node: {e}")))?;

        // Bind first (so the port is held) then spawn the server task. Binding to port 0 yields
        // an ephemeral port, which we read back via `local_addr` for tests.
        let listener = tokio::net::TcpListener::bind(bind_addr)
            .await
            .map_err(|e| ForgeError::Consensus(format!("failed to bind raft http server on {bind_addr}: {e}")))?;
        let local_addr = listener
            .local_addr()
            .map_err(|e| ForgeError::Consensus(format!("failed to read local addr: {e}")))?;
        let router = http_transport::router(raft.clone());
        let server_task = tokio::spawn(async move {
            if let Err(e) = axum::serve(listener, router).await {
                tracing::error!(error = %e, "raft http server stopped");
            }
        });

        info!(node_id, %local_addr, "RaftStateStore node started (multi-node HTTP transport)");
        Ok(Self {
            node_id,
            raft,
            read_path: ReadPath::Memory(state_machine),
            server_task: Some(server_task),
            local_addr: Some(local_addr),
        })
    }

    /// The actual bound HTTP address of this multi-node node (resolved ephemeral port).
    pub fn local_addr(&self) -> Option<std::net::SocketAddr> {
        self.local_addr
    }

    /// Initialize a multi-node cluster from this node, declaring all initial voters.
    ///
    /// `members` maps each node id to its advertised HTTP address ("host:port"). Call this on
    /// exactly one node after all nodes have been started with [`RaftStateStore::start_node`].
    pub async fn initialize_cluster(&self, members: BTreeMap<u64, String>) -> Result<()> {
        let nodes: BTreeMap<u64, openraft::BasicNode> =
            members.into_iter().map(|(id, addr)| (id, openraft::BasicNode::new(addr))).collect();

        match self.raft.initialize(nodes).await {
            Ok(()) => Ok(()),
            Err(openraft::error::RaftError::APIError(openraft::error::InitializeError::NotAllowed(_))) => Ok(()),
            Err(e) => Err(ForgeError::Consensus(format!("failed to initialize cluster: {e}"))),
        }
    }

    /// Add a learner (non-voting replica) to the cluster. `addr` is the learner's HTTP address.
    /// If `blocking`, waits until the learner has caught up. Must be called on the leader.
    pub async fn add_learner(&self, id: u64, addr: String, blocking: bool) -> Result<()> {
        self.raft
            .add_learner(id, openraft::BasicNode::new(addr), blocking)
            .await
            .map_err(|e| ForgeError::Consensus(format!("add_learner({id}) failed: {e}")))?;
        Ok(())
    }

    /// Change the voter membership of the cluster. Must be called on the leader.
    pub async fn change_membership(&self, voters: std::collections::BTreeSet<u64>, retain: bool) -> Result<()> {
        self.raft
            .change_membership(voters, retain)
            .await
            .map_err(|e| ForgeError::Consensus(format!("change_membership failed: {e}")))?;
        Ok(())
    }

    /// The node id this node currently believes is the leader, if any.
    pub async fn current_leader(&self) -> Option<u64> {
        self.raft.current_leader().await
    }

    /// The greatest log id applied to this node's local state machine (for replication polling).
    pub fn last_applied_index(&self) -> Option<u64> {
        self.raft.metrics().borrow().last_applied.map(|l| l.index)
    }

    /// Borrow a clone of the current metrics snapshot (test/diagnostic helper).
    pub fn server_state(&self) -> openraft::ServerState {
        self.raft.metrics().borrow().state
    }

    /// Poll metrics until this node reports itself as the current leader.
    async fn wait_for_leadership(&self) -> Result<()> {
        self.raft
            .wait(Some(Duration::from_secs(10)))
            .current_leader(self.node_id, "wait for self to become leader")
            .await
            .map_err(|e| ForgeError::Consensus(format!("node did not become leader: {e}")))?;
        Ok(())
    }

    /// Commit a request through Raft consensus and return the applied response.
    async fn write(&self, req: KvRequest) -> Result<KvResponse> {
        let resp = self
            .raft
            .client_write(req)
            .await
            .map_err(|e| ForgeError::Consensus(format!("raft client_write failed: {e}")))?;
        Ok(resp.data)
    }

    /// Number of voters/learners this node knows about (1 for single-node).
    pub fn node_id(&self) -> u64 {
        self.node_id
    }

    /// Gracefully shut down the underlying Raft node (and the HTTP server task, if any).
    pub async fn shutdown(&self) -> Result<()> {
        if let Some(task) = &self.server_task {
            task.abort();
        }
        self.raft
            .shutdown()
            .await
            .map_err(|e| ForgeError::Consensus(format!("raft shutdown failed: {e}")))
    }
}

#[async_trait]
impl StateStore for RaftStateStore {
    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
        // This is a leader-local read of applied state.
        Ok(self.read_path.get(key).await)
    }

    async fn set(&self, key: &str, value: Vec<u8>) -> Result<()> {
        self.write(KvRequest::Set { key: key.to_string(), value }).await?;
        Ok(())
    }

    async fn delete(&self, key: &str) -> Result<()> {
        self.write(KvRequest::Delete { key: key.to_string() }).await?;
        Ok(())
    }

    async fn list_prefix(&self, prefix: &str) -> Result<Vec<String>> {
        Ok(self.read_path.list_prefix(prefix).await)
    }

    fn name(&self) -> &str {
        "raft"
    }
}

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

    #[tokio::test]
    async fn test_raft_single_node_roundtrip() {
        let store = RaftStateStore::bootstrap_single_node(1)
            .await
            .expect("single node should bootstrap and elect itself leader");

        assert_eq!(store.name(), "raft");

        // set -> get
        store.set("forge/jobs/a", b"alpha".to_vec()).await.expect("set a");
        store.set("forge/jobs/b", b"bravo".to_vec()).await.expect("set b");
        store.set("forge/other/c", b"charlie".to_vec()).await.expect("set c");

        assert_eq!(store.get("forge/jobs/a").await.unwrap(), Some(b"alpha".to_vec()));
        assert_eq!(store.get("forge/jobs/b").await.unwrap(), Some(b"bravo".to_vec()));
        assert_eq!(store.get("missing").await.unwrap(), None);

        // overwrite returns the new value on read
        store.set("forge/jobs/a", b"alpha2".to_vec()).await.expect("overwrite a");
        assert_eq!(store.get("forge/jobs/a").await.unwrap(), Some(b"alpha2".to_vec()));

        // list_prefix
        let mut keys = store.list_prefix("forge/jobs/").await.unwrap();
        keys.sort();
        assert_eq!(keys, vec!["forge/jobs/a".to_string(), "forge/jobs/b".to_string()]);

        let all = store.list_prefix("forge/").await.unwrap();
        assert_eq!(all.len(), 3);

        // delete -> get returns None
        store.delete("forge/jobs/a").await.expect("delete a");
        assert_eq!(store.get("forge/jobs/a").await.unwrap(), None);

        let keys_after = store.list_prefix("forge/jobs/").await.unwrap();
        assert_eq!(keys_after, vec!["forge/jobs/b".to_string()]);

        store.shutdown().await.expect("clean shutdown");
    }

    #[tokio::test]
    async fn test_get_and_list_on_empty_store() {
        let store = RaftStateStore::bootstrap_single_node(7).await.unwrap();
        assert_eq!(store.get("nope").await.unwrap(), None);
        assert!(store.list_prefix("anything/").await.unwrap().is_empty());
        store.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_delete_missing_key_is_idempotent() {
        let store = RaftStateStore::bootstrap_single_node(8).await.unwrap();
        // Deleting a never-set key must still commit cleanly through Raft.
        store.delete("forge/jobs/ghost").await.unwrap();
        assert_eq!(store.get("forge/jobs/ghost").await.unwrap(), None);
        // Delete-after-delete is also fine (idempotent).
        store.set("forge/jobs/x", b"v".to_vec()).await.unwrap();
        store.delete("forge/jobs/x").await.unwrap();
        store.delete("forge/jobs/x").await.unwrap();
        assert_eq!(store.get("forge/jobs/x").await.unwrap(), None);
        store.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_overwrite_latest_wins_and_empty_value_distinct_from_absent() {
        let store = RaftStateStore::bootstrap_single_node(9).await.unwrap();
        store.set("k", b"v1".to_vec()).await.unwrap();
        store.set("k", b"v2".to_vec()).await.unwrap();
        store.set("k", b"v3".to_vec()).await.unwrap();
        assert_eq!(store.get("k").await.unwrap(), Some(b"v3".to_vec()));
        // An empty value is a real value, distinct from a missing key.
        store.set("empty", Vec::new()).await.unwrap();
        assert_eq!(store.get("empty").await.unwrap(), Some(Vec::new()));
        assert_eq!(store.get("never-set").await.unwrap(), None);
        store.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_list_prefix_respects_boundaries() {
        let store = RaftStateStore::bootstrap_single_node(10).await.unwrap();
        store.set("forge/jobs/b", b"1".to_vec()).await.unwrap();
        store.set("forge/jobs/a", b"1".to_vec()).await.unwrap();
        store.set("forge/nodes/x", b"1".to_vec()).await.unwrap();
        // Shares the text "forge/jobs" but not the "forge/jobs/" boundary.
        store.set("forge/jobsicle", b"1".to_vec()).await.unwrap();

        let jobs = store.list_prefix("forge/jobs/").await.unwrap();
        assert_eq!(jobs, vec!["forge/jobs/a".to_string(), "forge/jobs/b".to_string()]);
        assert!(!jobs.iter().any(|k| k == "forge/jobsicle"));
        assert_eq!(store.list_prefix("forge/").await.unwrap().len(), 4);
        store.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_many_keys_commit_then_partial_delete() {
        let store = RaftStateStore::bootstrap_single_node(11).await.unwrap();
        for i in 0..100u32 {
            store
                .set(&format!("forge/jobs/{i:03}"), i.to_le_bytes().to_vec())
                .await
                .unwrap();
        }
        assert_eq!(store.list_prefix("forge/jobs/").await.unwrap().len(), 100);
        assert_eq!(
            store.get("forge/jobs/042").await.unwrap(),
            Some(42u32.to_le_bytes().to_vec())
        );
        for i in 0..50u32 {
            store.delete(&format!("forge/jobs/{i:03}")).await.unwrap();
        }
        assert_eq!(store.list_prefix("forge/jobs/").await.unwrap().len(), 50);
        assert_eq!(store.get("forge/jobs/000").await.unwrap(), None);
        assert!(store.get("forge/jobs/099").await.unwrap().is_some());
        store.shutdown().await.unwrap();
    }

    // -----------------------------------------------------------------------
    // GOAL A: durability — survive a real process-internal restart from disk.
    // -----------------------------------------------------------------------
    #[cfg(feature = "raft-persist")]
    #[tokio::test]
    async fn test_raft_persist_survives_restart() {
        let dir = tempfile::tempdir().expect("temp dir");
        let path = dir.path().to_path_buf();

        // --- First lifetime: open at the dir, write several keys, then fully drop. ---
        {
            let store = RaftStateStore::open_persistent(100, &path)
                .await
                .expect("open persistent store (first time)");

            store.set("forge/jobs/a", b"alpha".to_vec()).await.expect("set a");
            store.set("forge/jobs/b", b"bravo".to_vec()).await.expect("set b");
            store.set("forge/nodes/x", b"node-x".to_vec()).await.expect("set x");
            store.set("forge/jobs/a", b"alpha-v2".to_vec()).await.expect("overwrite a");
            store.delete("forge/jobs/b").await.expect("delete b");

            // Sanity within the first lifetime.
            assert_eq!(store.get("forge/jobs/a").await.unwrap(), Some(b"alpha-v2".to_vec()));
            assert_eq!(store.get("forge/jobs/b").await.unwrap(), None);

            // Shut the node down and DROP it (and the fjall keyspace handle) so nothing is
            // reused in memory. append/apply already fsync'd to disk via PersistMode::SyncAll.
            store.shutdown().await.expect("clean shutdown");
            drop(store);
        }

        // --- Second lifetime: REOPEN the SAME dir from disk. State must have survived. ---
        let store2 = RaftStateStore::open_persistent(100, &path)
            .await
            .expect("reopen persistent store from disk");

        assert_eq!(
            store2.get("forge/jobs/a").await.unwrap(),
            Some(b"alpha-v2".to_vec()),
            "overwritten value must survive restart"
        );
        assert_eq!(
            store2.get("forge/jobs/b").await.unwrap(),
            None,
            "deleted key must stay deleted across restart"
        );
        assert_eq!(
            store2.get("forge/nodes/x").await.unwrap(),
            Some(b"node-x".to_vec()),
            "untouched key must survive restart"
        );

        // A write after reopen must still commit (proves the recovered node is a working leader).
        store2.set("forge/jobs/c", b"charlie".to_vec()).await.expect("set c after reopen");
        assert_eq!(store2.get("forge/jobs/c").await.unwrap(), Some(b"charlie".to_vec()));

        store2.shutdown().await.expect("clean shutdown of reopened store");
    }

    // -----------------------------------------------------------------------
    // GOAL B: multi-node HTTP transport — replication and failover.
    // -----------------------------------------------------------------------

    /// Poll a synchronous predicate until it is true or `timeout` elapses.
    async fn poll_until<F>(timeout: Duration, mut f: F) -> bool
    where
        F: FnMut() -> bool,
    {
        let deadline = tokio::time::Instant::now() + timeout;
        loop {
            if f() {
                return true;
            }
            if tokio::time::Instant::now() >= deadline {
                return false;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    }

    /// Poll until `store.get(key)` equals `want`, or `timeout` elapses (async-safe read path).
    async fn poll_until_value(store: &RaftStateStore, key: &str, want: &[u8], timeout: Duration) -> bool {
        let deadline = tokio::time::Instant::now() + timeout;
        loop {
            if store.get(key).await.ok().flatten().as_deref() == Some(want) {
                return true;
            }
            if tokio::time::Instant::now() >= deadline {
                return false;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    }

    /// Start 3 nodes on ephemeral localhost ports and initialize them as a single cluster.
    /// Returns the three stores and the index (0..3) of the elected leader.
    async fn start_three_node_cluster() -> (Vec<RaftStateStore>, usize) {
        let loopback: std::net::SocketAddr = ([127, 0, 0, 1], 0).into();

        let mut stores = Vec::new();
        for id in 1u64..=3 {
            stores.push(RaftStateStore::start_node(id, loopback).await.expect("start node"));
        }

        // Build the advertised member map from the actual bound ephemeral ports.
        let mut members = BTreeMap::new();
        for s in &stores {
            members.insert(s.node_id(), s.local_addr().unwrap().to_string());
        }

        // Initialize the cluster from node 1.
        stores[0].initialize_cluster(members).await.expect("initialize cluster");

        // Wait for a leader to be elected somewhere in the cluster.
        let elected = poll_until(Duration::from_secs(15), || {
            stores.iter().any(|s| s.server_state() == openraft::ServerState::Leader)
        })
        .await;
        assert!(elected, "a leader should be elected within timeout");

        let leader_idx = stores
            .iter()
            .position(|s| s.server_state() == openraft::ServerState::Leader)
            .expect("leader present");
        (stores, leader_idx)
    }

    #[tokio::test]
    async fn test_raft_multinode_replication() {
        let (stores, leader_idx) = start_three_node_cluster().await;

        // Write several keys on the leader.
        let leader = &stores[leader_idx];
        leader.set("forge/jobs/a", b"alpha".to_vec()).await.expect("set a on leader");
        leader.set("forge/jobs/b", b"bravo".to_vec()).await.expect("set b on leader");
        leader.set("forge/jobs/a", b"alpha2".to_vec()).await.expect("overwrite a on leader");

        // Pick a FOLLOWER and poll its locally-applied state until the value replicates.
        let follower_idx = (0..3).find(|&i| i != leader_idx).expect("a follower exists");
        let follower = &stores[follower_idx];

        let replicated = poll_until_value(follower, "forge/jobs/a", b"alpha2", Duration::from_secs(15)).await;
        assert!(
            replicated,
            "follower {} should observe the replicated value from leader {}",
            stores[follower_idx].node_id(),
            stores[leader_idx].node_id()
        );

        // The other written key should also be present on the follower by now.
        assert_eq!(follower.get("forge/jobs/b").await.unwrap(), Some(b"bravo".to_vec()));

        for s in &stores {
            s.shutdown().await.expect("shutdown");
        }
    }

    #[tokio::test]
    async fn test_raft_multinode_failover() {
        let (stores, leader_idx) = start_three_node_cluster().await;

        // Commit an initial value so we know replication is live before failover.
        stores[leader_idx].set("forge/before", b"v0".to_vec()).await.expect("initial write");

        // Wait until at least one follower has applied the initial value (quorum is replicating).
        let mut followers: Vec<usize> = (0..3).filter(|&i| i != leader_idx).collect();
        let mut replicated = false;
        for &i in &followers {
            if poll_until_value(&stores[i], "forge/before", b"v0", Duration::from_secs(15)).await {
                replicated = true;
                break;
            }
        }
        assert!(replicated, "initial value should replicate before failover");

        let old_leader_id = stores[leader_idx].node_id();

        // Kill the leader: shut it down so the remaining 2/3 must elect a new leader.
        stores[leader_idx].shutdown().await.expect("shut down old leader");

        // Among the two survivors, poll for a NEW leader (different node) to emerge.
        let new_leader = poll_until(Duration::from_secs(20), || {
            followers
                .iter()
                .any(|&i| stores[i].server_state() == openraft::ServerState::Leader)
        })
        .await;
        assert!(new_leader, "a new leader should be elected from the surviving quorum");

        let new_leader_idx = followers
            .iter()
            .copied()
            .find(|&i| stores[i].server_state() == openraft::ServerState::Leader)
            .expect("new leader present");
        assert_ne!(stores[new_leader_idx].node_id(), old_leader_id, "the new leader must differ from the killed one");

        // A subsequent write on the new leader must commit, and the surviving follower must see it.
        stores[new_leader_idx].set("forge/after", b"v1".to_vec()).await.expect("write on new leader");

        followers.retain(|&i| i != new_leader_idx);
        let survivor = followers[0];
        let saw_after = poll_until_value(&stores[survivor], "forge/after", b"v1", Duration::from_secs(15)).await;
        assert!(saw_after, "post-failover write should replicate to the surviving follower");

        // And the new leader can read back its own write.
        assert_eq!(stores[new_leader_idx].get("forge/after").await.unwrap(), Some(b"v1".to_vec()));

        for &i in &[new_leader_idx, survivor] {
            stores[i].shutdown().await.expect("shutdown survivor");
        }
    }
}