rust_cache_server/

store.rs

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use dashmap::DashMap;
use tokio::time;

#[derive(Clone, Debug)]
pub struct Value {
    pub data: Vec<u8>,
    pub expires_at: Option<Instant>,
}

impl Value {
    pub fn new(data: Vec<u8>, ttl: Option<Duration>) -> Self {
        let expires_at = ttl.map(|d| Instant::now() + d);
        Self { data, expires_at }
    }

    pub fn is_expired(&self) -> bool {
        if let Some(exp) = self.expires_at {
            Instant::now() >= exp
        } else {
            false
        }
    }
}

#[derive(Default)]
pub struct Store {
    inner: DashMap<String, Value>,
    max_memory_bytes: Option<u64>,
    used_bytes: AtomicU64,
}

impl Store {
    pub fn new(max_memory_bytes: Option<u64>) -> Self {
        Self {
            inner: DashMap::new(),
            max_memory_bytes,
            used_bytes: AtomicU64::new(0),
        }
    }

    fn value_size(value: &Value) -> u64 {
        value.data.len() as u64
    }

    pub fn get(&self, key: &str) -> Option<Vec<u8>> {
        if let Some(entry) = self.inner.get(key) {
            if entry.is_expired() {
                drop(entry);
                if let Some((_, old)) = self.inner.remove(key) {
                    let size = Self::value_size(&old);
                    self.used_bytes
                        .fetch_sub(size, Ordering::Relaxed);
                }
                None
            } else {
                Some(entry.data.clone())
            }
        } else {
            None
        }
    }

    pub fn set(
        &self,
        key: String,
        value: Vec<u8>,
        ttl: Option<Duration>,
    ) -> Result<(), ()> {
        let new_value = Value::new(value, ttl);
        let new_size = Self::value_size(&new_value);

        let mut old_size = 0;
        if let Some(old) = self.inner.get(&key) {
            old_size = Self::value_size(&old);
        }

        // Adjust used bytes and check against limit.
        let prev = self
            .used_bytes
            .fetch_add(new_size.saturating_sub(old_size), Ordering::Relaxed);
        let after = prev + new_size.saturating_sub(old_size);

        if let Some(limit) = self.max_memory_bytes {
            if after > limit {
                // Revert and reject.
                self.used_bytes
                    .fetch_sub(new_size.saturating_sub(old_size), Ordering::Relaxed);
                return Err(());
            }
        }

        self.inner.insert(key, new_value);
        Ok(())
    }

    pub fn del(&self, keys: &[String]) -> usize {
        let mut removed = 0;
        for k in keys {
            if let Some((_, v)) = self.inner.remove(k) {
                let size = Self::value_size(&v);
                self.used_bytes
                    .fetch_sub(size, Ordering::Relaxed);
                removed += 1;
            }
        }
        removed
    }

    pub fn expire(&self, key: &str, ttl: Duration) -> bool {
        if let Some(mut entry) = self.inner.get_mut(key) {
            if entry.is_expired() {
                drop(entry);
                if let Some((_, v)) = self.inner.remove(key) {
                    let size = Self::value_size(&v);
                    self.used_bytes
                        .fetch_sub(size, Ordering::Relaxed);
                }
                false
            } else {
                entry.expires_at = Some(Instant::now() + ttl);
                true
            }
        } else {
            false
        }
    }

    pub fn ttl(&self, key: &str) -> Option<Duration> {
        if let Some(entry) = self.inner.get(key) {
            if let Some(exp) = entry.expires_at {
                if Instant::now() >= exp {
                    drop(entry);
                    if let Some((_, v)) = self.inner.remove(key) {
                        let size = Self::value_size(&v);
                        self.used_bytes
                            .fetch_sub(size, Ordering::Relaxed);
                    }
                    None
                } else {
                    Some(exp.saturating_duration_since(Instant::now()))
                }
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Return all non-expired keys. Used by the KEYS command.
    pub fn keys(&self) -> Vec<String> {
        let mut result = Vec::new();
        for entry in self.inner.iter() {
            if entry.is_expired() {
                continue;
            }
            result.push(entry.key().clone());
        }
        result
    }

    /// Delete all keys and reset memory usage.
    pub fn clear_all(&self) {
        self.inner.clear();
        self.used_bytes.store(0, Ordering::Relaxed);
    }

    /// Increment/decrement a numeric value atomically.
    /// Returns the new value, or None if the key doesn't exist and create_if_missing is false,
    /// or if the value is not a valid integer.
    pub fn incr_by(&self, key: &str, delta: i64, create_if_missing: bool) -> Option<i64> {
        // Try to get existing value
        let current = if let Some(entry) = self.inner.get(key) {
            if entry.is_expired() {
                drop(entry);
                if let Some((_, old)) = self.inner.remove(key) {
                    let size = Self::value_size(&old);
                    self.used_bytes.fetch_sub(size, Ordering::Relaxed);
                }
                None
            } else {
                // Try to parse as integer
                match std::str::from_utf8(&entry.data) {
                    Ok(s) => s.trim().parse::<i64>().ok(),
                    Err(_) => None,
                }
            }
        } else {
            None
        };

        let new_value = if let Some(val) = current {
            val.checked_add(delta)?
        } else if create_if_missing {
            delta
        } else {
            return None;
        };

        // Update or insert the new value
        let new_data = new_value.to_string().into_bytes();
        let new_value_obj = Value::new(new_data.clone(), None);
        let new_size = Self::value_size(&new_value_obj);

        let old_size = if let Some(old) = self.inner.get(key) {
            Self::value_size(&old)
        } else {
            0
        };

        // Check memory limit
        if let Some(limit) = self.max_memory_bytes {
            let prev = self.used_bytes.load(Ordering::Relaxed);
            let after = prev + new_size.saturating_sub(old_size);
            if after > limit {
                return None;
            }
        }

        // Update memory counter
        self.used_bytes
            .fetch_add(new_size.saturating_sub(old_size), Ordering::Relaxed);

        self.inner.insert(key.to_string(), new_value_obj);
        Some(new_value)
    }
}

pub fn spawn_cleanup(store: Arc<Store>, interval: Duration) {
    tokio::spawn(async move {
        let mut ticker = time::interval(interval);
        loop {
            ticker.tick().await;

            // Collect expired keys first to avoid holding locks too long.
            let mut expired_keys = Vec::new();
            for entry in store.inner.iter() {
                if entry.is_expired() {
                    expired_keys.push(entry.key().clone());
                }
            }

            for key in expired_keys {
                if let Some((_, v)) = store.inner.remove(&key) {
                    let size = Store::value_size(&v);
                    store
                        .used_bytes
                        .fetch_sub(size, Ordering::Relaxed);
                }
            }
        }
    });
}