gitdataai/lib/cache/lib.rs

pub mod cluster;
pub mod error;
pub mod local;

use std::time::Duration;

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

#[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,
        }
    }
}

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

impl AppCache {
    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) => {
                    println!("cache:init:error with: {}", e);
                    return Err(e);
                }
            }),
            None => None,
        };

        Ok(Self {
            local,
            cluster,
            default_ttl: config.default_ttl,
            cluster_write_through: config.cluster_write_through,
        })
    }

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

    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? {
            return Ok(Some(value));
        }

        let Some(cluster) = &self.cluster else {
            return Ok(None);
        };

        let value = cluster.get::<T>(key).await?;
        if let Some(value) = &value {
            self.local.set(key, value).await?;
        }
        Ok(value)
    }

    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?;
        }
        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(())
    }

    pub async fn remove(&self, key: &str) -> CacheResult<()> {
        self.local.remove(key).await;
        if let Some(cluster) = &self.cluster {
            cluster.remove(key).await?;
        }
        Ok(())
    }
    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())
    }
}

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()))?,
        })
    }
}
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;
        }
    }
}