use moka::future::Cache as MokaCache;
use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheConfig {
pub max_capacity: u64,
pub ttl_secs: u64,
pub tti_secs: Option<u64>,
pub initial_capacity: Option<u64>,
pub enable_stats: bool,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
max_capacity: 10_000,
ttl_secs: 300, tti_secs: None,
initial_capacity: Some(1_000),
enable_stats: false,
}
}
}
impl CacheConfig {
pub fn high_performance() -> Self {
Self {
max_capacity: 100_000,
ttl_secs: 600, tti_secs: Some(300), initial_capacity: Some(10_000),
enable_stats: true,
}
}
pub fn low_memory() -> Self {
Self {
max_capacity: 1_000,
ttl_secs: 60, tti_secs: Some(30),
initial_capacity: Some(100),
enable_stats: false,
}
}
}
pub struct CacheManager {
cache: MokaCache<String, Vec<u8>>,
config: CacheConfig,
}
impl CacheManager {
pub fn new(config: CacheConfig) -> Self {
let mut builder = MokaCache::builder()
.max_capacity(config.max_capacity)
.time_to_live(Duration::from_secs(config.ttl_secs));
if let Some(tti_secs) = config.tti_secs {
builder = builder.time_to_idle(Duration::from_secs(tti_secs));
}
if let Some(initial_capacity) = config.initial_capacity {
builder = builder.initial_capacity(initial_capacity as usize);
}
let cache = builder.build();
Self { cache, config }
}
pub async fn get(&self, key: &str) -> Option<Vec<u8>> {
self.cache.get(key).await
}
pub async fn set(&self, key: String, value: Vec<u8>) {
self.cache.insert(key, value).await;
}
pub async fn remove(&self, key: &str) {
self.cache.invalidate(key).await;
}
pub async fn get_batch(&self, keys: &[String]) -> Vec<Option<Vec<u8>>> {
let mut results = Vec::with_capacity(keys.len());
for key in keys {
results.push(self.cache.get(key).await);
}
results
}
pub async fn set_batch(&self, entries: Vec<(String, Vec<u8>)>) {
for (key, value) in entries {
self.cache.insert(key, value).await;
}
}
pub async fn clear(&self) {
self.cache.invalidate_all();
}
pub fn entry_count(&self) -> u64 {
self.cache.entry_count()
}
pub fn config(&self) -> &CacheConfig {
&self.config
}
pub async fn contains(&self, key: &str) -> bool {
self.cache.get(key).await.is_some()
}
pub async fn get_or_insert<F>(&self, key: String, default: F) -> Vec<u8>
where
F: FnOnce() -> Vec<u8> + Send,
{
match self.cache.get(&key).await {
Some(value) => value,
None => {
let value = default();
self.cache.insert(key, value.clone()).await;
value
}
}
}
}
impl Default for CacheManager {
fn default() -> Self {
Self::new(CacheConfig::default())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_basic_operations() {
let cache = CacheManager::default();
cache.set("key1".to_string(), b"value1".to_vec()).await;
let value = cache.get("key1").await;
assert_eq!(value, Some(b"value1".to_vec()));
cache.remove("key1").await;
let value = cache.get("key1").await;
assert!(value.is_none());
}
#[tokio::test]
async fn test_ttl() {
let config = CacheConfig {
ttl_secs: 1,
..Default::default()
};
let cache = CacheManager::new(config);
cache.set("key1".to_string(), b"value1".to_vec()).await;
assert!(cache.get("key1").await.is_some());
tokio::time::sleep(Duration::from_secs(2)).await;
assert!(cache.get("key1").await.is_none());
}
#[tokio::test]
async fn test_batch_operations() {
let cache = CacheManager::default();
let entries = vec![
("key1".to_string(), b"value1".to_vec()),
("key2".to_string(), b"value2".to_vec()),
("key3".to_string(), b"value3".to_vec()),
];
cache.set_batch(entries).await;
let values = cache
.get_batch(&["key1".to_string(), "key2".to_string(), "key3".to_string()])
.await;
assert_eq!(values.len(), 3);
assert!(values[0].is_some());
assert!(values[1].is_some());
assert!(values[2].is_some());
}
#[tokio::test]
async fn test_clear() {
let cache = CacheManager::default();
cache.set("key1".to_string(), b"value1".to_vec()).await;
cache.set("key2".to_string(), b"value2".to_vec()).await;
tokio::time::sleep(Duration::from_millis(10)).await;
cache.clear().await;
assert!(cache.get("key1").await.is_none());
assert!(cache.get("key2").await.is_none());
}
#[tokio::test]
async fn test_contains() {
let cache = CacheManager::default();
cache.set("key1".to_string(), b"value1".to_vec()).await;
assert!(cache.contains("key1").await);
assert!(!cache.contains("key2").await);
}
#[tokio::test]
async fn test_get_or_insert() {
let cache = CacheManager::default();
let value = cache
.get_or_insert("key1".to_string(), || b"default".to_vec())
.await;
assert_eq!(value, b"default".to_vec());
let value = cache
.get_or_insert("key1".to_string(), || b"new_default".to_vec())
.await;
assert_eq!(value, b"default".to_vec());
}
#[test]
fn test_config_presets() {
let hp_config = CacheConfig::high_performance();
assert_eq!(hp_config.max_capacity, 100_000);
assert!(hp_config.tti_secs.is_some());
let lm_config = CacheConfig::low_memory();
assert_eq!(lm_config.max_capacity, 1_000);
assert!(lm_config.tti_secs.is_some());
}
#[tokio::test]
async fn test_max_capacity() {
let config = CacheConfig {
max_capacity: 5,
..Default::default()
};
let cache = CacheManager::new(config);
for i in 0..10 {
cache
.set(format!("key{}", i), format!("value{}", i).into_bytes())
.await;
}
assert!(cache.entry_count() <= 10);
}
}