clawstor/claw-store/src/cluster/wal_mutation.rs

//! Phase 4d (2026-07-13): typed mutation records for the WAL.
//!
//! The WAL itself (`cluster::wal`) treats payloads as opaque bytes.
//! Callers need a stable, self-describing encoding so a replay can
//! decide which RPC to re-issue on reconnect. This module owns that
//! encoding.
//!
//! Frame:
//! ```text
//! version : u8            = 0x01 today
//! kind    : u8            — one of `Kind`
//! body    : [u8]          — kind-specific
//! ```
//!
//! Body encodings deliberately mirror the existing on-wire shapes
//! (`refs::StampedRef::to_bytes`, `tags::encode_stamped_record`,
//! ...) so a future consumer can splice a WAL record straight into
//! an RPC payload with a memcpy.
//!
//! Blob-put mutations are NOT modeled here. Blob data is too large
//! to keep in the WAL; the roaming-client design stages blobs on
//! local disk and records a reference to them once the local
//! `BlobPutStream` completes. Ref/tag updates *are* the ordered,
//! small mutations the WAL was built for.

use crate::cluster::refs::{RefKey, RefValue, StampedRef};
use crate::cluster::tags::StampedTagValue;
use anyhow::{Context, Result};

/// Frame version. Bumps only on incompatible changes; adding a new
/// `Kind` is a compatible change (old readers surface it as
/// [`WalMutationError::UnknownKind`]).
pub const FRAME_VERSION: u8 = 0x01;

/// One-byte discriminator for the variants below. Values are the
/// authoritative wire order — never renumber.
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Kind {
    PutRef = 0x01,
    PutRefVersioned = 0x02,
    PutTag = 0x03,
    PutTagVersioned = 0x04,
    DeleteTag = 0x05,
    SetTagExpiry = 0x06,
}

impl Kind {
    pub fn from_byte(b: u8) -> Option<Self> {
        match b {
            0x01 => Some(Kind::PutRef),
            0x02 => Some(Kind::PutRefVersioned),
            0x03 => Some(Kind::PutTag),
            0x04 => Some(Kind::PutTagVersioned),
            0x05 => Some(Kind::DeleteTag),
            0x06 => Some(Kind::SetTagExpiry),
            _ => None,
        }
    }
    pub fn as_byte(self) -> u8 {
        self as u8
    }
}

/// A single client-mode mutation ready to durably record and later
/// replay. `PartialEq` + `Clone` are on the variants so tests + the
/// replay path can compare records without ceremony.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WalMutation {
    PutRef { key: RefKey, value: RefValue },
    PutRefVersioned { key: RefKey, stamped: StampedRef },
    PutTag { key: String, value: [u8; 32] },
    PutTagVersioned { key: String, stamped: StampedTagValue },
    DeleteTag { key: String },
    SetTagExpiry { key: String, expires_at_unix: u64 },
}

/// Distinct decode failures so callers can classify (unknown-kind
/// records get logged & skipped; malformed ones abort replay).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WalMutationError {
    Empty,
    BadVersion(u8, u8),
    UnknownKind(u8),
    Malformed(String),
}

impl std::fmt::Display for WalMutationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            WalMutationError::Empty => write!(f, "WAL mutation frame is empty"),
            WalMutationError::BadVersion(got, want) => write!(
                f,
                "WAL mutation frame version {got:#04x} is not supported (expected {want:#04x})"
            ),
            WalMutationError::UnknownKind(k) => {
                write!(f, "WAL mutation frame kind {k:#04x} is unknown")
            }
            WalMutationError::Malformed(msg) => {
                write!(f, "WAL mutation body is malformed: {msg}")
            }
        }
    }
}

impl std::error::Error for WalMutationError {}

impl WalMutation {
    pub fn kind(&self) -> Kind {
        match self {
            WalMutation::PutRef { .. } => Kind::PutRef,
            WalMutation::PutRefVersioned { .. } => Kind::PutRefVersioned,
            WalMutation::PutTag { .. } => Kind::PutTag,
            WalMutation::PutTagVersioned { .. } => Kind::PutTagVersioned,
            WalMutation::DeleteTag { .. } => Kind::DeleteTag,
            WalMutation::SetTagExpiry { .. } => Kind::SetTagExpiry,
        }
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(64);
        out.push(FRAME_VERSION);
        out.push(self.kind().as_byte());
        match self {
            WalMutation::PutRef { key, value } => {
                out.extend_from_slice(key);
                out.extend_from_slice(value);
            }
            WalMutation::PutRefVersioned { key, stamped } => {
                out.extend_from_slice(key);
                out.extend_from_slice(&stamped.to_bytes());
            }
            WalMutation::PutTag { key, value } => {
                write_len_prefixed_key(&mut out, key);
                out.extend_from_slice(value);
            }
            WalMutation::PutTagVersioned { key, stamped } => {
                write_len_prefixed_key(&mut out, key);
                out.extend_from_slice(&stamped.to_bytes());
            }
            WalMutation::DeleteTag { key } => {
                write_len_prefixed_key(&mut out, key);
            }
            WalMutation::SetTagExpiry {
                key,
                expires_at_unix,
            } => {
                write_len_prefixed_key(&mut out, key);
                out.extend_from_slice(&expires_at_unix.to_le_bytes());
            }
        }
        out
    }

    pub fn decode(bytes: &[u8]) -> Result<Self, WalMutationError> {
        if bytes.len() < 2 {
            return Err(WalMutationError::Empty);
        }
        if bytes[0] != FRAME_VERSION {
            return Err(WalMutationError::BadVersion(bytes[0], FRAME_VERSION));
        }
        let kind = Kind::from_byte(bytes[1])
            .ok_or(WalMutationError::UnknownKind(bytes[1]))?;
        let body = &bytes[2..];
        match kind {
            Kind::PutRef => {
                let (key, value) = split_pair::<32, 32>(body, "PutRef")?;
                Ok(WalMutation::PutRef { key, value })
            }
            Kind::PutRefVersioned => {
                if body.len() != 32 + StampedRef::ENCODED_LEN {
                    return Err(malformed(format!(
                        "PutRefVersioned body is {} bytes, need {}",
                        body.len(),
                        32 + StampedRef::ENCODED_LEN
                    )));
                }
                let mut key = [0u8; 32];
                key.copy_from_slice(&body[..32]);
                let stamped = StampedRef::from_bytes(&body[32..])
                    .map_err(|e| malformed(format!("PutRefVersioned stamped: {e}")))?;
                Ok(WalMutation::PutRefVersioned { key, stamped })
            }
            Kind::PutTag => {
                let (key, rest) = read_len_prefixed_key(body, "PutTag")?;
                if rest.len() != 32 {
                    return Err(malformed(format!(
                        "PutTag value is {} bytes, need 32",
                        rest.len()
                    )));
                }
                let mut value = [0u8; 32];
                value.copy_from_slice(rest);
                Ok(WalMutation::PutTag { key, value })
            }
            Kind::PutTagVersioned => {
                let (key, rest) = read_len_prefixed_key(body, "PutTagVersioned")?;
                if rest.len() != StampedTagValue::ENCODED_LEN {
                    return Err(malformed(format!(
                        "PutTagVersioned stamped is {} bytes, need {}",
                        rest.len(),
                        StampedTagValue::ENCODED_LEN
                    )));
                }
                let stamped = StampedTagValue::from_bytes(rest)
                    .map_err(|e| malformed(format!("PutTagVersioned stamped: {e}")))?;
                Ok(WalMutation::PutTagVersioned { key, stamped })
            }
            Kind::DeleteTag => {
                let (key, rest) = read_len_prefixed_key(body, "DeleteTag")?;
                if !rest.is_empty() {
                    return Err(malformed(format!(
                        "DeleteTag has trailing {} bytes",
                        rest.len()
                    )));
                }
                Ok(WalMutation::DeleteTag { key })
            }
            Kind::SetTagExpiry => {
                let (key, rest) = read_len_prefixed_key(body, "SetTagExpiry")?;
                if rest.len() != 8 {
                    return Err(malformed(format!(
                        "SetTagExpiry expires_at is {} bytes, need 8",
                        rest.len()
                    )));
                }
                let expires_at_unix = u64::from_le_bytes(rest.try_into().unwrap());
                Ok(WalMutation::SetTagExpiry {
                    key,
                    expires_at_unix,
                })
            }
        }
    }
}

fn write_len_prefixed_key(out: &mut Vec<u8>, key: &str) {
    let bytes = key.as_bytes();
    out.extend_from_slice(&(bytes.len() as u16).to_le_bytes());
    out.extend_from_slice(bytes);
}

fn read_len_prefixed_key<'a>(
    body: &'a [u8],
    ctx: &'static str,
) -> Result<(String, &'a [u8]), WalMutationError> {
    if body.len() < 2 {
        return Err(malformed(format!("{ctx}: missing key length prefix")));
    }
    let key_len = u16::from_le_bytes([body[0], body[1]]) as usize;
    let start: usize = 2;
    let end = start
        .checked_add(key_len)
        .ok_or_else(|| malformed(format!("{ctx}: key length overflow")))?;
    if body.len() < end {
        return Err(malformed(format!(
            "{ctx}: declared key_len={key_len} but body has {} bytes after prefix",
            body.len() - 2
        )));
    }
    let key = std::str::from_utf8(&body[start..end])
        .map_err(|e| malformed(format!("{ctx}: key is not valid UTF-8: {e}")))?
        .to_string();
    Ok((key, &body[end..]))
}

fn split_pair<const A: usize, const B: usize>(
    body: &[u8],
    ctx: &'static str,
) -> Result<([u8; A], [u8; B]), WalMutationError> {
    if body.len() != A + B {
        return Err(malformed(format!(
            "{ctx} body is {} bytes, need {}",
            body.len(),
            A + B
        )));
    }
    let mut a = [0u8; A];
    let mut b = [0u8; B];
    a.copy_from_slice(&body[..A]);
    b.copy_from_slice(&body[A..]);
    Ok((a, b))
}

fn malformed(msg: String) -> WalMutationError {
    WalMutationError::Malformed(msg)
}

/// Convenience: encode + append to a WAL, returning the assigned
/// seq. Kept out of `WriteAheadLog` itself so the WAL primitive
/// stays payload-agnostic.
pub async fn append_mutation(
    wal: &mut crate::cluster::wal::WriteAheadLog,
    mutation: &WalMutation,
) -> anyhow::Result<u64> {
    let bytes = mutation.encode();
    wal.append(&bytes)
        .await
        .with_context(|| format!("appending {:?} to WAL", mutation.kind()))
}

/// Convenience: replay every WAL record from `start_seq`, decoding
/// each to a `WalMutation`. Unknown-kind records are surfaced so
/// callers can log-and-skip (forward-compat) rather than aborting
/// the whole replay.
pub async fn replay_mutations(
    wal: &crate::cluster::wal::WriteAheadLog,
    start_seq: u64,
) -> anyhow::Result<Vec<(u64, Result<WalMutation, WalMutationError>)>> {
    let recs = wal.iter_from(start_seq).await?;
    Ok(recs
        .into_iter()
        .map(|r| (r.seq, WalMutation::decode(&r.payload)))
        .collect())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cluster::refs::node_stamp_for;

    fn sample_stamped_ref() -> StampedRef {
        StampedRef {
            value: [0x11; 32],
            clock: 42,
            node: node_stamp_for("test-node"),
        }
    }

    fn sample_stamped_tag() -> StampedTagValue {
        StampedTagValue {
            value: [0x22; 32],
            clock: 99,
            node: node_stamp_for("other-node"),
        }
    }

    fn all_variants() -> Vec<WalMutation> {
        vec![
            WalMutation::PutRef {
                key: [0xAA; 32],
                value: [0xBB; 32],
            },
            WalMutation::PutRefVersioned {
                key: [0xCC; 32],
                stamped: sample_stamped_ref(),
            },
            WalMutation::PutTag {
                key: "clawverse:main".into(),
                value: [0xDD; 32],
            },
            WalMutation::PutTagVersioned {
                key: "clawverse:main:latest".into(),
                stamped: sample_stamped_tag(),
            },
            WalMutation::DeleteTag {
                key: "clawverse:main:pr-1".into(),
            },
            WalMutation::SetTagExpiry {
                key: "clawverse:main:latest".into(),
                expires_at_unix: 1_800_000_000,
            },
        ]
    }

    #[test]
    fn kind_bytes_stable_and_round_trip() {
        assert_eq!(Kind::PutRef.as_byte(), 0x01);
        assert_eq!(Kind::PutRefVersioned.as_byte(), 0x02);
        assert_eq!(Kind::PutTag.as_byte(), 0x03);
        assert_eq!(Kind::PutTagVersioned.as_byte(), 0x04);
        assert_eq!(Kind::DeleteTag.as_byte(), 0x05);
        assert_eq!(Kind::SetTagExpiry.as_byte(), 0x06);
        for m in all_variants() {
            assert_eq!(Kind::from_byte(m.kind().as_byte()), Some(m.kind()));
        }
    }

    #[test]
    fn round_trip_every_variant() {
        for m in all_variants() {
            let bytes = m.encode();
            assert_eq!(bytes[0], FRAME_VERSION);
            assert_eq!(bytes[1], m.kind().as_byte());
            let back = WalMutation::decode(&bytes).unwrap();
            assert_eq!(back, m);
        }
    }

    #[test]
    fn decode_rejects_empty_and_short() {
        assert!(matches!(
            WalMutation::decode(&[]),
            Err(WalMutationError::Empty)
        ));
        assert!(matches!(
            WalMutation::decode(&[FRAME_VERSION]),
            Err(WalMutationError::Empty)
        ));
    }

    #[test]
    fn decode_rejects_bad_version() {
        let bad = vec![0xFF, Kind::PutRef as u8];
        assert!(matches!(
            WalMutation::decode(&bad),
            Err(WalMutationError::BadVersion(0xFF, FRAME_VERSION))
        ));
    }

    #[test]
    fn decode_flags_unknown_kind() {
        let bad = vec![FRAME_VERSION, 0xEE];
        assert!(matches!(
            WalMutation::decode(&bad),
            Err(WalMutationError::UnknownKind(0xEE))
        ));
    }

    #[test]
    fn decode_flags_malformed_bodies() {
        // PutRef body is exactly 64 bytes; supply 60.
        let mut short = vec![FRAME_VERSION, Kind::PutRef as u8];
        short.extend(std::iter::repeat(0u8).take(60));
        assert!(matches!(
            WalMutation::decode(&short),
            Err(WalMutationError::Malformed(_))
        ));

        // PutTag with declared key_len larger than actual body.
        let mut bad_tag = vec![FRAME_VERSION, Kind::PutTag as u8];
        bad_tag.extend_from_slice(&100u16.to_le_bytes()); // key_len=100
        bad_tag.extend_from_slice(b"abc"); // only 3 bytes follow
        assert!(matches!(
            WalMutation::decode(&bad_tag),
            Err(WalMutationError::Malformed(_))
        ));

        // DeleteTag with trailing garbage.
        let mut bad_del = vec![FRAME_VERSION, Kind::DeleteTag as u8];
        bad_del.extend_from_slice(&3u16.to_le_bytes());
        bad_del.extend_from_slice(b"abcXX"); // 3-byte key + 2-byte tail
        assert!(matches!(
            WalMutation::decode(&bad_del),
            Err(WalMutationError::Malformed(_))
        ));
    }

    #[test]
    fn decode_flags_non_utf8_key() {
        let mut bytes = vec![FRAME_VERSION, Kind::DeleteTag as u8];
        bytes.extend_from_slice(&3u16.to_le_bytes());
        bytes.extend_from_slice(&[0xFF, 0xFE, 0xFD]); // invalid UTF-8
        assert!(matches!(
            WalMutation::decode(&bytes),
            Err(WalMutationError::Malformed(_))
        ));
    }

    #[tokio::test]
    async fn append_and_replay_round_trip() {
        let tmp = tempfile::TempDir::new().unwrap();
        let mut wal = crate::cluster::wal::WriteAheadLog::open(tmp.path().join("wal"))
            .await
            .unwrap();
        let mutations = all_variants();
        for m in &mutations {
            let seq = append_mutation(&mut wal, m).await.unwrap();
            assert!(seq >= 1);
        }
        let replayed = replay_mutations(&wal, 0).await.unwrap();
        assert_eq!(replayed.len(), mutations.len());
        for ((_, decoded), original) in replayed.iter().zip(mutations.iter()) {
            assert_eq!(decoded.as_ref().unwrap(), original);
        }
    }

    #[tokio::test]
    async fn replay_survives_unknown_kind_records() {
        let tmp = tempfile::TempDir::new().unwrap();
        let mut wal = crate::cluster::wal::WriteAheadLog::open(tmp.path().join("wal"))
            .await
            .unwrap();
        // Real mutation first.
        append_mutation(
            &mut wal,
            &WalMutation::DeleteTag {
                key: "k".into(),
            },
        )
        .await
        .unwrap();
        // Then a future-kind record the current binary doesn't
        // understand. Replay should surface it as Err, not panic.
        wal.append(&[FRAME_VERSION, 0x77, 0x01, 0x02]).await.unwrap();
        // Real mutation last.
        append_mutation(
            &mut wal,
            &WalMutation::DeleteTag {
                key: "k2".into(),
            },
        )
        .await
        .unwrap();

        let replayed = replay_mutations(&wal, 0).await.unwrap();
        assert_eq!(replayed.len(), 3);
        assert!(replayed[0].1.is_ok());
        assert!(matches!(
            replayed[1].1,
            Err(WalMutationError::UnknownKind(0x77))
        ));
        assert!(replayed[2].1.is_ok());
    }
}