memory_table.rs

← Back to explorer
        Rust
        150 lines · UTF-8
      
use crate::types::{BranchId, Key, Result, Timestamp, Value};
use crossbeam_skiplist::SkipMap;
use parking_lot::RwLock;
use std::sync::atomic::{AtomicUsize, Ordering};

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
struct MemoryTableKey {
    key: Vec<u8>,
    branch_id: BranchId,
    timestamp: std::cmp::Reverse<Timestamp>,
}

pub struct MemoryTable {
    data: SkipMap<MemoryTableKey, Option<Value>>,
    size: AtomicUsize,
    max_size: usize,
    is_frozen: RwLock<bool>,
}

impl MemoryTable {
    pub fn new_table(max_table_size: usize) -> Self {
        Self {
            data: SkipMap::new(),
            size: AtomicUsize::new(0),
            max_size: max_table_size,
            is_frozen: RwLock::new(false),
        }
    }

    pub fn put_data(&self, table_key: &Key, value: Value, timestamp: Timestamp, branch_id: BranchId) -> Result<()> {
        if *self.is_frozen.read() {
            return Ok(());
        }

        let memory_table_key = MemoryTableKey {
            key: table_key.0.clone(),
            branch_id,
            timestamp: std::cmp::Reverse(timestamp),
        };

        let data_entry_size = table_key.0.len() + value.length() + 32;
        self.size.fetch_add(data_entry_size, Ordering::Relaxed);

        self.data.insert(memory_table_key, Some(value));
        Ok(())
    }

    pub fn delete_data(&self, table_key: &Key, timestamp: Timestamp, branch_id: BranchId) -> Result<()> {
        if *self.is_frozen.read() {
            return Ok(());
        }

        let memory_table_key = MemoryTableKey {
            key: table_key.0.clone(),
            branch_id,
            timestamp: std::cmp::Reverse(timestamp),
        };

        let data_entry_size = table_key.0.len() + 32;
        self.size.fetch_add(data_entry_size, Ordering::Relaxed);

        self.data.insert(memory_table_key, None);

        Ok(())
    }

    pub fn get_data(&self, table_key: &Key, branch_id: BranchId, time: Option<Timestamp>) -> Option<Option<Value>> {
        let start_range = MemoryTableKey {
            key: table_key.0.clone(),
            branch_id,
            timestamp: std::cmp::Reverse(time.unwrap_or(u64::MAX)),
        };

        for data_entry in self.data.range(start_range..) {
            let _key = data_entry.key();

            if _key.key != table_key.0 || _key.branch_id != branch_id {
                break;
            }
            return Some(data_entry.value().clone());
        }

        None
    }

    pub fn get_history(&self, table_key: &Key, branch_id: BranchId, get_limit: usize) -> Vec<(Timestamp, Option<Value>)> {
        let start_range = MemoryTableKey {
            key: table_key.0.clone(),
            branch_id,
            timestamp: std::cmp::Reverse(u64::MAX),
        };

        let mut history = Vec::new();

        for data_entry in self.data.range(start_range..) {
            let _key = data_entry.key();

            if _key.key != table_key.0 || _key.branch_id != branch_id {
                break;
            }

            history.push((_key.timestamp.0, data_entry.value().clone()));

            if history.len() >= get_limit {
                break;
            }
        }

        history
    }

    pub fn scan_table(&self, start_range: &Key, end_range: &Key, branch_id: BranchId, time: Option<Timestamp>,) -> Vec<(Key, Value)> {
        let mut scan_results = Vec::new();
        let mut last_key: Option<Vec<u8>> = None;
        let scan_start = MemoryTableKey {
            key: start_range.0.clone(),
            branch_id,
            timestamp: std::cmp::Reverse(u64::MAX),
        };
        let scan_limit = time.unwrap_or(u64::MAX);

        for data_entry in self.data.range(scan_start..) {
            let _key = data_entry.key();

            if _key.key.as_slice() >= end_range.0.as_slice() {
                break;
            }

            if _key.branch_id != branch_id {
                continue;
            }

            if _key.timestamp.0 > scan_limit {
                continue;
            }

            if last_key.as_ref() == Some(&_key.key) {
                continue;
            }

            last_key = Some(_key.key.clone());

            if let Some(ref value) = *data_entry.value() {
                scan_results.push((Key(_key.key.clone()), value.clone()));
            }
        }

        scan_results
    }
}