gitdataai/lib/cache/app.rs

use std::time::Duration;

use track::CounterVec;

use crate::{
    cluster::{ClusterCache, ClusterCacheConfig},
    error::{CacheError, CacheResult},
    local::{LocalCacheConfig, MokaCache},
};

// ============================================================================
// Configuration
// ============================================================================

#[derive(Clone, Debug)]
pub struct AppCacheConfig {
    pub local: LocalCacheConfig,
    pub cluster: Option<ClusterCacheConfig>,
    pub default_ttl: Option<Duration>,
    pub cluster_write_through: bool,
}

impl Default for AppCacheConfig {
    fn default() -> Self {
        Self {
            local: LocalCacheConfig::default(),
            cluster: None,
            default_ttl: Some(Duration::from_secs(300)),
            cluster_write_through: true,
        }
    }
}

impl TryFrom<&config::AppConfig> for AppCacheConfig {
    type Error = CacheError;

    fn try_from(config: &config::AppConfig) -> Result<Self, Self::Error> {
        let local = LocalCacheConfig {
            max_capacity: config
                .cache_local_max_capacity()
                .map_err(|error| CacheError::Config(error.to_string()))?,
            time_to_live: config
                .cache_local_ttl()
                .map_err(|error| CacheError::Config(error.to_string()))?,
            time_to_idle: config
                .cache_local_tti()
                .map_err(|error| CacheError::Config(error.to_string()))?,
        };

        let cluster = if config
            .cache_cluster_enabled()
            .map_err(|error| CacheError::Config(error.to_string()))?
        {
            Some(ClusterCacheConfig {
                urls: config
                    .redis_urls()
                    .map_err(|error| CacheError::Config(error.to_string()))?,
                key_prefix: config.cache_cluster_key_prefix(),
                command_timeout: config
                    .cache_cluster_command_timeout()
                    .map_err(|error| CacheError::Config(error.to_string()))?,
            })
        } else {
            None
        };

        Ok(Self {
            local,
            cluster,
            default_ttl: config
                .cache_default_ttl()
                .map_err(|error| CacheError::Config(error.to_string()))?,
            cluster_write_through: config
                .cache_cluster_write_through()
                .map_err(|error| CacheError::Config(error.to_string()))?,
        })
    }
}

// ============================================================================
// AppCache
// ============================================================================

#[derive(Clone)]
pub struct AppCache {
    pub local: MokaCache,
    pub cluster: Option<ClusterCache>,
    default_ttl: Option<Duration>,
    cluster_write_through: bool,
    metrics: Option<track::MetricsRegistry>,
}

impl AppCache {
    #[tracing::instrument(skip(config))]
    pub async fn init(config: AppCacheConfig) -> CacheResult<Self> {
        let local = MokaCache::with_config(config.local);
        let cluster = match config.cluster {
            Some(cluster) => Some(match ClusterCache::connect(cluster).await {
                Ok(cluster) => cluster,
                Err(e) => {
                    tracing::error!(error = %e, "failed to connect to cache cluster");
                    return Err(e);
                }
            }),
            None => None,
        };

        tracing::info!(has_cluster = cluster.is_some(), "cache initialized");
        Ok(Self {
            local,
            cluster,
            default_ttl: config.default_ttl,
            cluster_write_through: config.cluster_write_through,
            metrics: None,
        })
    }

    pub fn local_only(local: MokaCache) -> Self {
        Self {
            local,
            cluster: None,
            default_ttl: None,
            cluster_write_through: false,
            metrics: None,
        }
    }

    /// Attach a metrics registry for recording cache counters.
    pub fn set_metrics(&mut self, registry: track::MetricsRegistry) {
        self.metrics = Some(registry);
    }

    #[tracing::instrument(skip(self), fields(cache.key = %key))]
    pub async fn get<T>(&self, key: &str) -> CacheResult<Option<T>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned,
    {
        if let Some(value) = self.local.get(key).await? {
            tracing::debug!("cache hit (local)");
            self.record_hit("local");
            return Ok(Some(value));
        }

        let Some(cluster) = &self.cluster else {
            tracing::debug!("cache miss");
            self.record_miss();
            return Ok(None);
        };

        let value = cluster.get::<T>(key).await?;
        if let Some(value) = &value {
            self.local.set(key, value).await?;
            tracing::debug!("cache hit (cluster)");
            self.record_hit("cluster");
        } else {
            tracing::debug!("cache miss");
            self.record_miss();
        }
        Ok(value)
    }

    #[tracing::instrument(skip(self, value), fields(cache.key = %key))]
    pub async fn set<T>(&self, key: &str, value: &T) -> CacheResult<()>
    where
        T: serde::Serialize + ?Sized,
    {
        self.local.set(key, value).await?;
        if self.cluster_write_through
            && let Some(cluster) = &self.cluster
        {
            cluster.set(key, value, self.default_ttl).await?;
        }
        self.record_set();
        Ok(())
    }

    pub async fn set_with_ttl<T>(
        &self,
        key: &str,
        value: &T,
        ttl: std::time::Duration,
    ) -> CacheResult<()>
    where
        T: serde::Serialize + ?Sized,
    {
        self.local.set(key, value).await?;
        if self.cluster_write_through
            && let Some(cluster) = &self.cluster
        {
            cluster.set(key, value, Some(ttl)).await?;
        }
        Ok(())
    }

    #[tracing::instrument(skip(self), fields(cache.key = %key))]
    pub async fn remove(&self, key: &str) -> CacheResult<()> {
        self.local.remove(key).await;
        if let Some(cluster) = &self.cluster {
            cluster.remove(key).await?;
        }
        self.record_remove();
        Ok(())
    }

    fn record_hit(&self, tier: &str) {
        if let Some(reg) = &self.metrics {
            cache_hits_vec(reg).with_label_values(&[tier]).inc();
        }
    }

    fn record_miss(&self) {
        if let Some(reg) = &self.metrics {
            cache_misses_vec(reg).with_label_values(&[]).inc();
        }
    }

    fn record_set(&self) {
        if let Some(reg) = &self.metrics {
            cache_sets_vec(reg).with_label_values(&[]).inc();
        }
    }

    fn record_remove(&self) {
        if let Some(reg) = &self.metrics {
            cache_removes_vec(reg).with_label_values(&[]).inc();
        }
    }

    pub async fn delete_pattern(&self, pattern: &str) -> CacheResult<u64> {
        let pattern = pattern.to_string();
        let local_pattern = pattern.clone();
        self.local.invalidate_entries_if(move |key| {
            simple_glob_match(&local_pattern, key)
        });

        let mut removed = 0u64;
        if let Some(cluster) = &self.cluster {
            removed = cluster.delete_pattern(&pattern).await?;
        }
        Ok(removed)
    }

    pub async fn ping_cluster(&self) -> CacheResult<()> {
        if let Some(cluster) = &self.cluster {
            cluster.ping().await?;
        }
        Ok(())
    }

    pub fn conn(&self) -> Option<redis::cluster_async::ClusterConnection> {
        self.cluster.as_ref().map(|c| c.conn())
    }
}

fn cache_hits_vec(registry: &track::MetricsRegistry) -> CounterVec {
    registry
        .register_counter_vec("cache_hits_total", "Total cache hits", &["tier"])
        .expect("failed to register cache_hits_total")
}

fn cache_misses_vec(registry: &track::MetricsRegistry) -> CounterVec {
    registry
        .register_counter_vec("cache_misses_total", "Total cache misses", &[])
        .expect("failed to register cache_misses_total")
}

fn cache_sets_vec(registry: &track::MetricsRegistry) -> CounterVec {
    registry
        .register_counter_vec(
            "cache_sets_total",
            "Total cache set operations",
            &[],
        )
        .expect("failed to register cache_sets_total")
}

fn cache_removes_vec(registry: &track::MetricsRegistry) -> CounterVec {
    registry
        .register_counter_vec(
            "cache_removes_total",
            "Total cache remove operations",
            &[],
        )
        .expect("failed to register cache_removes_total")
}

// ============================================================================
// Helpers
// ============================================================================

fn simple_glob_match(pattern: &str, key: &str) -> bool {
    let p = pattern.as_bytes();
    let k = key.as_bytes();
    let (mut pi, mut ki) = (0usize, 0usize);

    let mut backtrack_p: Option<usize> = None;
    let mut backtrack_k: usize = 0;

    loop {
        if pi < p.len() && ki < k.len() && (p[pi] == b'?' || p[pi] == k[ki]) {
            pi += 1;
            ki += 1;
        } else if pi < p.len() && p[pi] == b'*' {
            backtrack_p = Some(pi);
            backtrack_k = ki;
            pi += 1;
        } else if let Some(saved_pi) = backtrack_p {
            backtrack_k += 1;
            ki = backtrack_k;
            pi = saved_pi + 1;
        } else {
            return pi == p.len() && ki == k.len();
        }

        if pi == p.len() && ki == k.len() {
            return true;
        }
    }
}