wedb_embed 0.1.0

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

/// 布隆过滤器链元数据(对标 Apache Kvrocks BloomChainMetadata 46字节)
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct BloomChainMeta {
    pub base: KeyMeta,
    pub n_filters: u16,
    pub expansion: u16,
    pub base_capacity: u32,
    pub error_rate: f64,
    pub bloom_bytes: u32,
}

impl BloomChainMeta {
    pub const ENCODED_SIZE: usize = KeyMeta::ENCODED_SIZE + 2 + 2 + 4 + 8 + 4; // 26 + 20 = 46

    #[inline]
    pub fn new(
        base_capacity: u32,
        error_rate: f64,
        expansion: u16,
        version: u64,
        expire_at: u64,
        bloom_bytes: u32,
    ) -> Self {
        let ver = if version == 0 {
            crate::meta::generate_version()
        } else {
            version
        };
        Self {
            base: KeyMeta::new(RedisType::Bloom, expire_at, ver, 0),
            n_filters: 1,
            expansion,
            base_capacity,
            error_rate,
            bloom_bytes,
        }
    }

    #[inline]
    pub fn is_scaling(&self) -> bool {
        self.expansion != 0
    }

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

    #[inline]
    pub fn get_capacity(&self) -> u32 {
        if self.expansion == 0 {
            return self.base_capacity;
        }
        if self.expansion == 1 {
            return self.base_capacity.saturating_mul(self.n_filters as u32);
        }
        let mut total = 0u64;
        let mut cur = self.base_capacity as u64;
        for _ in 0..self.n_filters {
            total = total.saturating_add(cur);
            cur = cur.saturating_mul(self.expansion as u64);
        }
        total.min(u32::MAX as u64) as u32
    }

    #[inline]
    pub fn sub_filter_capacity(&self, filter_index: u16) -> u32 {
        if self.expansion == 0 {
            self.base_capacity
        } else {
            let mut cap = self.base_capacity as u64;
            for _ in 0..filter_index {
                cap = cap.saturating_mul(self.expansion as u64);
            }
            cap.min(u32::MAX as u64) as u32
        }
    }

    #[inline]
    pub fn encode(&self) -> [u8; Self::ENCODED_SIZE] {
        let mut buf = [0u8; Self::ENCODED_SIZE];
        buf[..KeyMeta::ENCODED_SIZE].copy_from_slice(&self.base.encode());
        let mut offset = KeyMeta::ENCODED_SIZE;
        buf[offset..offset + 2].copy_from_slice(&self.n_filters.to_be_bytes());
        offset += 2;
        buf[offset..offset + 2].copy_from_slice(&self.expansion.to_be_bytes());
        offset += 2;
        buf[offset..offset + 4].copy_from_slice(&self.base_capacity.to_be_bytes());
        offset += 4;
        buf[offset..offset + 8].copy_from_slice(&self.error_rate.to_be_bytes());
        offset += 8;
        buf[offset..offset + 4].copy_from_slice(&self.bloom_bytes.to_be_bytes());
        buf
    }

    #[inline]
    pub fn decode(bytes: &[u8]) -> Option<Self> {
        if bytes.len() < Self::ENCODED_SIZE {
            return None;
        }
        let base = KeyMeta::decode(&bytes[..KeyMeta::ENCODED_SIZE])?;
        let mut offset = KeyMeta::ENCODED_SIZE;

        let mut b2 = [0u8; 2];
        b2.copy_from_slice(&bytes[offset..offset + 2]);
        let n_filters = u16::from_be_bytes(b2);
        offset += 2;

        b2.copy_from_slice(&bytes[offset..offset + 2]);
        let expansion = u16::from_be_bytes(b2);
        offset += 2;

        let mut b4 = [0u8; 4];
        b4.copy_from_slice(&bytes[offset..offset + 4]);
        let base_capacity = u32::from_be_bytes(b4);
        offset += 4;

        let mut b8 = [0u8; 8];
        b8.copy_from_slice(&bytes[offset..offset + 8]);
        let error_rate = f64::from_be_bytes(b8);
        offset += 8;

        b4.copy_from_slice(&bytes[offset..offset + 4]);
        let bloom_bytes = u32::from_be_bytes(b4);

        Some(Self {
            base,
            n_filters,
            expansion,
            base_capacity,
            error_rate,
            bloom_bytes,
        })
    }
}

/// 布谷鸟过滤器链元数据(对标 Apache Kvrocks CuckooChainMetadata 53字节)
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct CuckooChainMeta {
    pub base: KeyMeta,
    pub n_filters: u16,
    pub expansion: u16,
    pub base_capacity: u64,
    pub bucket_size: u8,
    pub max_iterations: u16,
    pub num_deleted_items: u64,
    pub page_size: u32,
}

impl CuckooChainMeta {
    pub const ENCODED_SIZE: usize = KeyMeta::ENCODED_SIZE + 2 + 2 + 8 + 1 + 2 + 8 + 4; // 26 + 27 = 53

    #[inline]
    pub fn new(
        base_capacity: u64,
        bucket_size: u8,
        max_iterations: u16,
        expansion: u16,
        page_size: u32,
        version: u64,
        expire_at: u64,
    ) -> Self {
        let ver = if version == 0 {
            crate::meta::generate_version()
        } else {
            version
        };
        Self {
            base: KeyMeta::new(RedisType::CuckooFilter, expire_at, ver, 0),
            n_filters: 1,
            expansion,
            base_capacity,
            bucket_size,
            max_iterations,
            num_deleted_items: 0,
            page_size,
        }
    }

    #[inline]
    pub fn is_scaling(&self) -> bool {
        self.expansion > 0
    }

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

    #[inline]
    pub fn get_total_capacity(&self) -> u64 {
        if self.expansion == 0 || self.n_filters == 1 {
            return self.base_capacity;
        }
        let mut total = 0u64;
        let mut filter_cap = self.base_capacity;
        for _ in 0..self.n_filters {
            total = total.saturating_add(filter_cap);
            filter_cap = filter_cap.saturating_mul(self.expansion as u64);
        }
        total
    }

    #[inline]
    pub fn sub_filter_capacity(&self, filter_index: u16) -> Option<u64> {
        let mut cap = self.base_capacity;
        for _ in 0..filter_index {
            if self.expansion != 0 && cap > u64::MAX / (self.expansion as u64) {
                return None;
            }
            cap *= self.expansion as u64;
        }
        Some(cap)
    }

    #[inline]
    pub fn sub_filter_num_buckets(&self, filter_index: u16) -> Result<u32> {
        let cap = self
            .sub_filter_capacity(filter_index)
            .ok_or_else(|| Error::invalid_data("filter capacity overflow"))?;
        super::CuckooFilterHelper::calculate_required_buckets(cap, self.bucket_size)
    }

    #[inline]
    pub fn encode(&self) -> [u8; Self::ENCODED_SIZE] {
        let mut buf = [0u8; Self::ENCODED_SIZE];
        buf[..KeyMeta::ENCODED_SIZE].copy_from_slice(&self.base.encode());
        let mut offset = KeyMeta::ENCODED_SIZE;

        buf[offset..offset + 2].copy_from_slice(&self.n_filters.to_be_bytes());
        offset += 2;
        buf[offset..offset + 2].copy_from_slice(&self.expansion.to_be_bytes());
        offset += 2;
        buf[offset..offset + 8].copy_from_slice(&self.base_capacity.to_be_bytes());
        offset += 8;
        buf[offset] = self.bucket_size;
        offset += 1;
        buf[offset..offset + 2].copy_from_slice(&self.max_iterations.to_be_bytes());
        offset += 2;
        buf[offset..offset + 8].copy_from_slice(&self.num_deleted_items.to_be_bytes());
        offset += 8;
        buf[offset..offset + 4].copy_from_slice(&self.page_size.to_be_bytes());
        buf
    }

    #[inline]
    pub fn decode(bytes: &[u8]) -> Option<Self> {
        if bytes.len() < Self::ENCODED_SIZE {
            return None;
        }
        let base = KeyMeta::decode(&bytes[..KeyMeta::ENCODED_SIZE])?;
        let mut offset = KeyMeta::ENCODED_SIZE;

        let mut b2 = [0u8; 2];
        b2.copy_from_slice(&bytes[offset..offset + 2]);
        let n_filters = u16::from_be_bytes(b2);
        offset += 2;

        b2.copy_from_slice(&bytes[offset..offset + 2]);
        let expansion = u16::from_be_bytes(b2);
        offset += 2;

        let mut b8 = [0u8; 8];
        b8.copy_from_slice(&bytes[offset..offset + 8]);
        let base_capacity = u64::from_be_bytes(b8);
        offset += 8;

        let bucket_size = bytes[offset];
        offset += 1;

        b2.copy_from_slice(&bytes[offset..offset + 2]);
        let max_iterations = u16::from_be_bytes(b2);
        offset += 2;

        b8.copy_from_slice(&bytes[offset..offset + 8]);
        let num_deleted_items = u64::from_be_bytes(b8);
        offset += 8;

        let mut b4 = [0u8; 4];
        b4.copy_from_slice(&bytes[offset..offset + 4]);
        let page_size = u32::from_be_bytes(b4);

        Some(Self {
            base,
            n_filters,
            expansion,
            base_capacity,
            bucket_size,
            max_iterations,
            num_deleted_items,
            page_size,
        })
    }
}