wedb_embed 0.1.0

Embedded Kvrocks-compatible storage engine for WeDb
Documentation
use crate::meta::{KeyMeta, RedisType};
use serde::{Deserialize, Serialize};

/// JSON 存储格式(对标 Apache Kvrocks JsonStorageFormat)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[repr(u8)]
pub enum JsonStorageFormat {
    #[default]
    Json = 0,
    Cbor = 1,
}

/// JSON 结构元数据(对标 Apache Kvrocks JsonMetadata 27 字节/26 字节)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct JsonMeta {
    pub base: KeyMeta,
    pub format: JsonStorageFormat,
}

impl JsonMeta {
    pub const ENCODED_SIZE: usize = KeyMeta::ENCODED_SIZE + 1; // 26 + 1 = 27
    pub const KVROCKS_ENCODED_SIZE: usize = KeyMeta::KVROCKS_COMPLEX_ENCODED_SIZE + 1; // 25 + 1 = 26

    #[inline]
    pub fn new(expire_at: u64, version: u64, size: u64) -> Self {
        Self {
            base: KeyMeta::new(RedisType::Json, expire_at, version, size),
            format: JsonStorageFormat::Json,
        }
    }

    #[inline]
    pub fn new_with_version(expire_at: u64, size: u64) -> Self {
        Self {
            base: KeyMeta::new_with_version(RedisType::Json, expire_at, size),
            format: JsonStorageFormat::Json,
        }
    }

    #[inline]
    pub fn with_format(format: JsonStorageFormat, expire_at: u64, version: u64, size: u64) -> Self {
        Self {
            base: KeyMeta::new(RedisType::Json, expire_at, version, size),
            format,
        }
    }

    #[inline]
    pub fn is_expired(&self, now_ms: u64) -> bool {
        self.base.is_expired(now_ms)
    }

    /// 编码为标准 27 字节 wedb 元数据头(栈上定长数组,零堆分配)
    #[inline]
    pub fn encode(&self) -> [u8; Self::ENCODED_SIZE] {
        let mut buf = [0u8; Self::ENCODED_SIZE];
        let base_enc = self.base.encode();
        buf[..KeyMeta::ENCODED_SIZE].copy_from_slice(&base_enc);
        buf[KeyMeta::ENCODED_SIZE] = self.format as u8;
        buf
    }

    /// 编码为 Kvrocks 26 字节紧凑元数据头(栈上定长数组,零堆分配)
    #[inline]
    pub fn encode_kvrocks(&self) -> [u8; Self::KVROCKS_ENCODED_SIZE] {
        let mut buf = [0u8; Self::KVROCKS_ENCODED_SIZE];
        let flags = KeyMeta::META_64BIT_ENCODING_MASK | (self.base.rtype as u8 & KeyMeta::META_TYPE_MASK);
        buf[0] = flags;
        buf[1..9].copy_from_slice(&self.base.expire_at.to_be_bytes());
        buf[9..17].copy_from_slice(&self.base.version.to_be_bytes());
        buf[17..25].copy_from_slice(&self.base.size.to_be_bytes());
        buf[25] = self.format as u8;
        buf
    }

    /// 解码元数据头与载荷切片(自适应支持 27 字节 wedb 格式、26 字节 Kvrocks 格式及裸 JSON)
    #[inline]
    pub fn decode(bytes: &[u8]) -> Option<(Self, &[u8])> {
        if bytes.is_empty() {
            return None;
        }

        // 1. 标准 27 字节 wedb JsonMeta 头部 (rtype == RedisType::Json)
        if bytes.len() >= Self::ENCODED_SIZE
            && bytes[0] == RedisType::Json as u8
            && let Some(base) = KeyMeta::decode(&bytes[..KeyMeta::ENCODED_SIZE])
        {
            let format = match bytes[KeyMeta::ENCODED_SIZE] {
                1 => JsonStorageFormat::Cbor,
                _ => JsonStorageFormat::Json,
            };
            let payload = &bytes[Self::ENCODED_SIZE..];
            return Some((Self { base, format }, payload));
        }

        // 2. Kvrocks 26 字节 JsonMetadata 头部 (flags & 0x0F == 10, flags & 0x80 != 0)
        if bytes.len() >= Self::KVROCKS_ENCODED_SIZE
            && (bytes[0] & KeyMeta::META_TYPE_MASK == RedisType::Json as u8)
            && (bytes[0] & KeyMeta::META_64BIT_ENCODING_MASK != 0)
            && let Some(base) = KeyMeta::decode(&bytes[..KeyMeta::KVROCKS_COMPLEX_ENCODED_SIZE])
        {
            let format = match bytes[KeyMeta::KVROCKS_COMPLEX_ENCODED_SIZE] {
                1 => JsonStorageFormat::Cbor,
                _ => JsonStorageFormat::Json,
            };
            let payload = &bytes[Self::KVROCKS_ENCODED_SIZE..];
            return Some((Self { base, format }, payload));
        }

        // 3. 裸 JSON 文本兼容(跳过前导空白符)
        if let Some(&first_non_ws) = bytes.iter().find(|&&b| !b.is_ascii_whitespace())
            && matches!(
                first_non_ws,
                b'{' | b'[' | b'"' | b't' | b'f' | b'n' | b'0'..=b'9' | b'-'
            )
        {
            return Some((Self::new(0, 0, bytes.len() as u64), bytes));
        }

        None
    }
}

/// 编码完整 JSON 存储值 (元数据头 + 载荷)
#[inline]
pub fn encode_json_value(meta: &JsonMeta, payload: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(JsonMeta::ENCODED_SIZE + payload.len());
    out.extend_from_slice(&meta.encode());
    out.extend_from_slice(payload);
    out
}

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

    #[test]
    fn test_json_meta_roundtrip() {
        let meta = JsonMeta::with_format(JsonStorageFormat::Json, 1000, 42, 128);
        assert_eq!(JsonMeta::ENCODED_SIZE, 27);
        let enc = meta.encode();
        assert_eq!(enc.len(), 27);

        let (dec, payload) = JsonMeta::decode(&enc).expect("decode failed");
        assert_eq!(dec.format, JsonStorageFormat::Json);
        assert_eq!(dec.base.expire_at, 1000);
        assert_eq!(dec.base.version, 42);
        assert_eq!(dec.base.size, 128);
        assert!(payload.is_empty());
    }

    #[test]
    fn test_json_meta_kvrocks_roundtrip() {
        let meta = JsonMeta::with_format(JsonStorageFormat::Cbor, 2000, 99, 256);
        let enc_kv = meta.encode_kvrocks();
        assert_eq!(enc_kv.len(), 26);

        let (dec, payload) = JsonMeta::decode(&enc_kv).expect("decode kvrocks failed");
        assert_eq!(dec.format, JsonStorageFormat::Cbor);
        assert_eq!(dec.base.expire_at, 2000);
        assert_eq!(dec.base.version, 99);
        assert_eq!(dec.base.size, 256);
        assert!(payload.is_empty());
    }

    #[test]
    fn test_json_meta_raw_json() {
        let raw = br#"{"a":1,"b":"hello"}"#;
        let (dec, payload) = JsonMeta::decode(raw).expect("decode raw json failed");
        assert_eq!(dec.format, JsonStorageFormat::Json);
        assert_eq!(dec.base.expire_at, 0);
        assert_eq!(payload, raw);
    }
}