mx_cache/

local.rs

//! Local in-memory cache implementation using Moka.
//!
//! This implementation supports per-entry TTL to ensure behavioral consistency
//! with [`RedisCache`](crate::RedisCache). Each `set()` and `set_nx_px()` call
//! respects the provided TTL parameter.
//!
//! # Performance Optimizations
//!
//! This implementation uses `Bytes` for both keys and values to minimize
//! allocations on the hot path:
//! - Keys use `Bytes` which is reference-counted and cheap to clone
//! - Values use `Bytes` avoiding per-lookup allocations
//! - Lookups can use borrowed slices via the `equivalent` trait

use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::time::{Duration, Instant};

use bytes::Bytes;
use moka::Expiry;
use moka::sync::Cache;

use crate::Cache as CacheTrait;

/// A wrapper around `Bytes` that supports borrowed lookups via the `equivalent` trait.
///
/// This allows Moka to look up keys using `&[u8]` without allocating a new `Bytes`.
#[derive(Clone, Debug, Eq)]
struct CacheKey(Bytes);

impl CacheKey {
    #[inline]
    fn new(key: &[u8]) -> Self {
        Self(Bytes::copy_from_slice(key))
    }

    #[inline]
    fn as_slice(&self) -> &[u8] {
        &self.0
    }
}

impl PartialEq for CacheKey {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

impl Hash for CacheKey {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

/// Wrapper for borrowed key lookups without allocation.
///
/// This implements the `equivalent` trait pattern used by Moka to allow
/// looking up entries using `&[u8]` without creating a `CacheKey`.
#[derive(PartialEq, Eq)]
struct BorrowedKey<'a>(&'a [u8]);

impl Hash for BorrowedKey<'_> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl equivalent::Equivalent<CacheKey> for BorrowedKey<'_> {
    #[inline]
    fn equivalent(&self, key: &CacheKey) -> bool {
        self.0 == key.as_slice()
    }
}

/// Cache entry containing value and its expiration time.
///
/// Uses `Bytes` for the value to avoid cloning the data on reads.
#[derive(Clone, Debug)]
struct CacheEntry {
    value: Bytes,
    expires_at: Instant,
}

/// Custom expiry policy that uses per-entry expiration times.
///
/// This allows each entry to have its own TTL, matching Redis behavior
/// where each SET operation can specify a different expiration.
struct PerEntryExpiry;

impl Expiry<CacheKey, CacheEntry> for PerEntryExpiry {
    /// Returns the TTL for a newly created entry.
    fn expire_after_create(
        &self,
        _key: &CacheKey,
        value: &CacheEntry,
        _current_time: Instant,
    ) -> Option<Duration> {
        let now = Instant::now();
        if value.expires_at > now {
            Some(value.expires_at.duration_since(now))
        } else {
            // Already expired, expire immediately
            Some(Duration::ZERO)
        }
    }

    /// Returns the TTL after a read operation (no change).
    fn expire_after_read(
        &self,
        _key: &CacheKey,
        _value: &CacheEntry,
        _current_time: Instant,
        current_duration: Option<Duration>,
        _last_modified_at: Instant,
    ) -> Option<Duration> {
        // Don't extend TTL on read - maintain original expiration
        current_duration
    }

    /// Returns the TTL after an update operation.
    fn expire_after_update(
        &self,
        _key: &CacheKey,
        value: &CacheEntry,
        _current_time: Instant,
        _current_duration: Option<Duration>,
    ) -> Option<Duration> {
        let now = Instant::now();
        if value.expires_at > now {
            Some(value.expires_at.duration_since(now))
        } else {
            Some(Duration::ZERO)
        }
    }
}

/// Local in-memory cache with per-entry TTL support.
///
/// Uses Moka's sync cache with a custom expiry policy to support
/// per-operation TTL values, ensuring behavioral consistency with
/// [`RedisCache`](crate::RedisCache).
///
/// # Performance
///
/// This implementation is optimized for hot-path usage:
/// - Uses `Bytes` for keys and values (reference-counted, cheap clones)
/// - Supports borrowed lookups via the `equivalent` trait for `contains_sync`
/// - Minimizes allocations on repeated operations with the same key
///
/// # Example
///
/// ```
/// use std::time::Duration;
/// use mx_cache::{LocalCache, Cache};
///
/// # tokio_test::block_on(async {
/// let cache = LocalCache::new(10_000, Duration::from_secs(300));
///
/// // Each operation can specify its own TTL
/// cache.set(b"short_lived", b"value1", Duration::from_secs(10)).await.unwrap();
/// cache.set(b"long_lived", b"value2", Duration::from_secs(3600)).await.unwrap();
/// # });
/// ```
#[derive(Clone, Debug)]
pub struct LocalCache {
    inner: Arc<Cache<CacheKey, CacheEntry>>,
}

impl LocalCache {
    /// Creates a new local cache with the specified capacity and default TTL.
    ///
    /// # Arguments
    ///
    /// * `capacity` - Maximum number of entries in the cache
    /// * `default_ttl` - Default TTL used as an upper bound for entries.
    ///   Individual operations can specify shorter TTLs, but entries will
    ///   never live longer than this default.
    ///
    /// # Note
    ///
    /// The `default_ttl` parameter sets a maximum lifetime for entries.
    /// Per-operation TTLs (passed to `set()` and `set_nx_px()`) take
    /// precedence and are fully respected, as long as they don't exceed
    /// this default.
    pub fn new(capacity: u64, default_ttl: Duration) -> Self {
        let cache = Cache::builder()
            .max_capacity(capacity.max(1))
            // Set a maximum TTL as a safety bound
            .time_to_live(default_ttl)
            // Use custom expiry for per-entry TTL support
            .expire_after(PerEntryExpiry)
            .build();
        Self {
            inner: Arc::new(cache),
        }
    }

    /// Synchronous check for local cache presence - fast path for deduplication.
    ///
    /// Returns `true` if the key exists in the cache and has not expired.
    ///
    /// # Performance
    ///
    /// This is a synchronous operation optimized for hot-path deduplication
    /// checks. Uses borrowed key lookup via the `equivalent` trait to avoid
    /// allocating a new `Bytes` on every call.
    #[inline]
    pub fn contains_sync(&self, key: &[u8]) -> bool {
        // Use borrowed key lookup to avoid allocation
        self.inner.contains_key(&BorrowedKey(key))
    }

    /// Synchronous get for local cache - fast path for hot data.
    ///
    /// Returns the value if the key exists and has not expired.
    ///
    /// # Performance
    ///
    /// This is a synchronous operation that returns a `Bytes` reference
    /// to avoid copying the value data.
    #[inline]
    pub fn get_sync(&self, key: &[u8]) -> Option<Bytes> {
        self.inner.get(&BorrowedKey(key)).map(|entry| entry.value)
    }
}

impl CacheTrait for LocalCache {
    /// Sets a value only if the key does not exist (NX = Not eXists).
    ///
    /// # Arguments
    ///
    /// * `key` - The cache key
    /// * `value` - The value to store
    /// * `ttl` - Time-to-live for this entry
    ///
    /// # Returns
    ///
    /// * `Ok(true)` - The key was newly inserted
    /// * `Ok(false)` - The key already existed, no change made
    fn set_nx_px(
        &self,
        key: &[u8],
        value: &[u8],
        ttl: Duration,
    ) -> impl Future<Output = anyhow::Result<bool>> + Send {
        // Single allocation for key, single allocation for value
        let cache_key = CacheKey::new(key);
        let entry = CacheEntry {
            value: Bytes::copy_from_slice(value),
            expires_at: Instant::now() + ttl,
        };
        let inner = Arc::clone(&self.inner);

        async move {
            // Use Moka's entry API for atomic insert-if-not-exists operation.
            // This eliminates the TOCTOU race condition between contains_key and insert.
            let result = inner.entry(cache_key).or_insert(entry);

            // is_fresh() returns true if the value was just inserted (key didn't exist),
            // false if the key already existed.
            Ok(result.is_fresh())
        }
    }

    /// Sets a value with the specified TTL.
    ///
    /// # Arguments
    ///
    /// * `key` - The cache key
    /// * `value` - The value to store
    /// * `ttl` - Time-to-live for this entry
    fn set(
        &self,
        key: &[u8],
        value: &[u8],
        ttl: Duration,
    ) -> impl Future<Output = anyhow::Result<()>> + Send {
        // Single allocation for key, single allocation for value
        let cache_key = CacheKey::new(key);
        let entry = CacheEntry {
            value: Bytes::copy_from_slice(value),
            expires_at: Instant::now() + ttl,
        };
        let inner = Arc::clone(&self.inner);

        async move {
            inner.insert(cache_key, entry);
            Ok(())
        }
    }

    /// Gets a value from the cache.
    ///
    /// # Returns
    ///
    /// * `Ok(Some(value))` - The value exists and has not expired
    /// * `Ok(None)` - The key does not exist or has expired
    fn get(&self, key: &[u8]) -> impl Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send {
        // Use borrowed key lookup to avoid allocation on lookup
        let result = self.inner.get(&BorrowedKey(key)).map(|entry| entry.value);
        async move { Ok(result.map(|bytes| bytes.to_vec())) }
    }

    /// Deletes a key from the cache.
    fn del(&self, key: &[u8]) -> impl Future<Output = anyhow::Result<()>> + Send {
        // Use borrowed key for invalidation
        self.inner.invalidate(&BorrowedKey(key));
        async move { Ok(()) }
    }
}

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

    #[tokio::test]
    async fn test_set_and_get() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        cache
            .set(b"key1", b"value1", Duration::from_secs(60))
            .await
            .unwrap();

        let result = cache.get(b"key1").await.unwrap();
        assert_eq!(result, Some(b"value1".to_vec()));
    }

    #[tokio::test]
    async fn test_get_nonexistent() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        let result = cache.get(b"nonexistent").await.unwrap();
        assert_eq!(result, None);
    }

    #[tokio::test]
    async fn test_set_nx_px_new_key() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        let was_set = cache
            .set_nx_px(b"key1", b"value1", Duration::from_secs(60))
            .await
            .unwrap();
        assert!(was_set, "Expected key to be set (new key)");

        let result = cache.get(b"key1").await.unwrap();
        assert_eq!(result, Some(b"value1".to_vec()));
    }

    #[tokio::test]
    async fn test_set_nx_px_existing_key() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        // First set should succeed
        let was_set1 = cache
            .set_nx_px(b"key1", b"value1", Duration::from_secs(60))
            .await
            .unwrap();
        assert!(was_set1);

        // Second set should fail (key exists)
        let was_set2 = cache
            .set_nx_px(b"key1", b"value2", Duration::from_secs(60))
            .await
            .unwrap();
        assert!(!was_set2, "Expected key NOT to be set (key exists)");

        // Original value should remain
        let result = cache.get(b"key1").await.unwrap();
        assert_eq!(result, Some(b"value1".to_vec()));
    }

    #[tokio::test]
    async fn test_del() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        cache
            .set(b"key1", b"value1", Duration::from_secs(60))
            .await
            .unwrap();
        cache.del(b"key1").await.unwrap();

        let result = cache.get(b"key1").await.unwrap();
        assert_eq!(result, None);
    }

    #[tokio::test]
    async fn test_contains_sync() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        assert!(!cache.contains_sync(b"key1"));

        cache
            .set(b"key1", b"value1", Duration::from_secs(60))
            .await
            .unwrap();

        assert!(cache.contains_sync(b"key1"));
    }

    #[tokio::test]
    async fn test_get_sync() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        assert!(cache.get_sync(b"key1").is_none());

        cache
            .set(b"key1", b"value1", Duration::from_secs(60))
            .await
            .unwrap();

        let result = cache.get_sync(b"key1");
        assert_eq!(result, Some(Bytes::from_static(b"value1")));
    }

    #[tokio::test]
    async fn test_per_entry_ttl_respected() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        // Set with very short TTL
        cache
            .set(b"short_ttl", b"value", Duration::from_millis(50))
            .await
            .unwrap();

        // Should exist immediately
        let result = cache.get(b"short_ttl").await.unwrap();
        assert_eq!(result, Some(b"value".to_vec()));

        // Wait for expiration
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Should be expired now
        let result = cache.get(b"short_ttl").await.unwrap();
        assert_eq!(result, None, "Entry should have expired after TTL");
    }

    #[tokio::test]
    async fn test_different_ttls_for_different_keys() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        // Set entries with different TTLs
        cache
            .set(b"short", b"value1", Duration::from_millis(50))
            .await
            .unwrap();
        cache
            .set(b"long", b"value2", Duration::from_secs(10))
            .await
            .unwrap();

        // Both should exist initially
        assert!(cache.get(b"short").await.unwrap().is_some());
        assert!(cache.get(b"long").await.unwrap().is_some());

        // Wait for short TTL to expire
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Short TTL entry should be gone, long TTL should remain
        assert!(
            cache.get(b"short").await.unwrap().is_none(),
            "Short TTL entry should have expired"
        );
        assert!(
            cache.get(b"long").await.unwrap().is_some(),
            "Long TTL entry should still exist"
        );
    }

    #[tokio::test]
    async fn test_set_nx_px_ttl_respected() {
        let cache = LocalCache::new(100, Duration::from_secs(60));

        // Set with very short TTL using set_nx_px
        let was_set = cache
            .set_nx_px(b"key", b"value", Duration::from_millis(50))
            .await
            .unwrap();
        assert!(was_set);

        // Wait for expiration
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Should be expired now
        let result = cache.get(b"key").await.unwrap();
        assert_eq!(result, None, "Entry should have expired after TTL");

        // Should be able to set again since key expired
        let was_set_again = cache
            .set_nx_px(b"key", b"new_value", Duration::from_secs(60))
            .await
            .unwrap();
        assert!(
            was_set_again,
            "Should be able to set after previous entry expired"
        );

        let result = cache.get(b"key").await.unwrap();
        assert_eq!(result, Some(b"new_value".to_vec()));
    }
}