forge-orchestration 0.5.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)
//! This module implements a **single-node** Raft cluster. The node bootstraps itself as the
//! sole voter, wins the election, and becomes leader. Writes 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.
//!
//! Multi-node networking is intentionally **out of scope**: the [`RaftNetwork`] implementation
//! here is a loopback no-op. Its methods are never invoked in a single-node cluster (there are
//! no peers to talk to), so they return an `Unreachable` RPC error / `unreachable!()` rather
//! than performing any real RPC.
//!
//! ## Storage internals
//! Both the log store and the state machine are kept fully in memory (a `BTreeMap` log and a
//! `BTreeMap` key/value map). They are modelled on the official openraft `raft-kv-memstore`
//! example for the v0.9 line, adapted to 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::RaftTypeConfig;
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;

// ---------------------------------------------------------------------------
// 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(())
    }
}

// ---------------------------------------------------------------------------
// 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")),
    ))
}

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

/// A [`StateStore`] backed by a single-node 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.
pub struct RaftStateStore {
    node_id: u64,
    raft: ForgeRaft,
    state_machine: Arc<StateMachineStore>,
}

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, state_machine };
        store.wait_for_leadership().await?;
        info!(node_id, "RaftStateStore bootstrapped and leader elected");
        Ok(store)
    }

    /// 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.
    pub async fn shutdown(&self) -> Result<()> {
        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>>> {
        // See `read_consistency_note`: this is a leader-local read of applied state.
        Ok(self.state_machine.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.state_machine.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");
    }
}