use crate::error::Error;
use crate::tagged::TaggedCache;
use async_trait::async_trait;
use serde::{de::DeserializeOwned, Serialize};
use std::env;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct CacheConfig {
pub default_ttl: Duration,
pub prefix: String,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
default_ttl: Duration::from_secs(3600),
prefix: String::new(),
}
}
}
impl CacheConfig {
pub fn new() -> Self {
Self::default()
}
pub fn from_env() -> Self {
let prefix = env::var("CACHE_PREFIX").unwrap_or_default();
let default_ttl = env::var("CACHE_TTL")
.ok()
.and_then(|v| v.parse().ok())
.map(Duration::from_secs)
.unwrap_or_else(|| Duration::from_secs(3600));
Self {
default_ttl,
prefix,
}
}
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.default_ttl = ttl;
self
}
pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
self.prefix = prefix.into();
self
}
}
#[async_trait]
pub trait CacheStore: Send + Sync {
async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>, Error>;
async fn put_raw(&self, key: &str, value: Vec<u8>, ttl: Duration) -> Result<(), Error>;
async fn has(&self, key: &str) -> Result<bool, Error>;
async fn forget(&self, key: &str) -> Result<bool, Error>;
async fn flush(&self) -> Result<(), Error>;
async fn increment(&self, key: &str, value: i64) -> Result<i64, Error>;
async fn decrement(&self, key: &str, value: i64) -> Result<i64, Error>;
async fn tag_add(&self, tag: &str, key: &str) -> Result<(), Error>;
async fn tag_members(&self, tag: &str) -> Result<Vec<String>, Error>;
async fn tag_flush(&self, tag: &str) -> Result<(), Error>;
}
#[derive(Clone)]
pub struct Cache {
store: Arc<dyn CacheStore>,
config: CacheConfig,
}
impl Cache {
pub fn new(store: Arc<dyn CacheStore>) -> Self {
Self {
store,
config: CacheConfig::default(),
}
}
pub fn with_config(store: Arc<dyn CacheStore>, config: CacheConfig) -> Self {
Self { store, config }
}
#[cfg(feature = "memory")]
pub fn memory() -> Self {
Self::new(Arc::new(crate::stores::MemoryStore::new()))
}
#[cfg(feature = "redis-backend")]
pub async fn redis(url: &str) -> Result<Self, Error> {
let store = crate::stores::RedisStore::new(url).await?;
Ok(Self::new(Arc::new(store)))
}
#[cfg(feature = "memory")]
pub async fn from_env() -> Result<Self, Error> {
let driver = env::var("CACHE_DRIVER").unwrap_or_else(|_| "memory".to_string());
let config = CacheConfig::from_env();
let store: Arc<dyn CacheStore> = match driver.as_str() {
#[cfg(feature = "redis-backend")]
"redis" => {
let url =
env::var("REDIS_URL").unwrap_or_else(|_| "redis://127.0.0.1:6379".to_string());
Arc::new(crate::stores::RedisStore::new(&url).await?)
}
_ => {
let capacity = env::var("CACHE_MEMORY_CAPACITY")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(10_000);
Arc::new(crate::stores::MemoryStore::with_capacity(capacity))
}
};
Ok(Self::with_config(store, config))
}
fn prefixed_key(&self, key: &str) -> String {
if self.config.prefix.is_empty() {
key.to_string()
} else {
format!("{}:{}", self.config.prefix, key)
}
}
pub async fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Error> {
let key = self.prefixed_key(key);
match self.store.get_raw(&key).await? {
Some(bytes) => {
let value = serde_json::from_slice(&bytes)
.map_err(|e| Error::deserialization(e.to_string()))?;
Ok(Some(value))
}
None => Ok(None),
}
}
pub async fn put<T: Serialize>(
&self,
key: &str,
value: &T,
ttl: Duration,
) -> Result<(), Error> {
let key = self.prefixed_key(key);
let bytes = serde_json::to_vec(value).map_err(|e| Error::serialization(e.to_string()))?;
self.store.put_raw(&key, bytes, ttl).await
}
pub async fn put_default<T: Serialize>(&self, key: &str, value: &T) -> Result<(), Error> {
self.put(key, value, self.config.default_ttl).await
}
pub async fn forever<T: Serialize>(&self, key: &str, value: &T) -> Result<(), Error> {
self.put(key, value, Duration::from_secs(315_360_000)).await }
pub async fn has(&self, key: &str) -> Result<bool, Error> {
let key = self.prefixed_key(key);
self.store.has(&key).await
}
pub async fn forget(&self, key: &str) -> Result<bool, Error> {
let key = self.prefixed_key(key);
self.store.forget(&key).await
}
pub async fn flush(&self) -> Result<(), Error> {
self.store.flush().await
}
pub async fn remember<T, F, Fut>(&self, key: &str, ttl: Duration, f: F) -> Result<T, Error>
where
T: Serialize + DeserializeOwned,
F: FnOnce() -> Fut,
Fut: Future<Output = T>,
{
if let Some(value) = self.get(key).await? {
return Ok(value);
}
let value = f().await;
self.put(key, &value, ttl).await?;
Ok(value)
}
pub async fn remember_forever<T, F, Fut>(&self, key: &str, f: F) -> Result<T, Error>
where
T: Serialize + DeserializeOwned,
F: FnOnce() -> Fut,
Fut: Future<Output = T>,
{
self.remember(key, Duration::from_secs(315_360_000), f)
.await
}
pub async fn pull<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Error> {
let value = self.get(key).await?;
if value.is_some() {
self.forget(key).await?;
}
Ok(value)
}
pub async fn increment(&self, key: &str, value: i64) -> Result<i64, Error> {
let key = self.prefixed_key(key);
self.store.increment(&key, value).await
}
pub async fn decrement(&self, key: &str, value: i64) -> Result<i64, Error> {
let key = self.prefixed_key(key);
self.store.decrement(&key, value).await
}
pub fn tags(&self, tags: &[&str]) -> TaggedCache {
TaggedCache::new(
self.store.clone(),
tags.iter().map(|s| s.to_string()).collect(),
self.config.clone(),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_config_builder() {
let config = CacheConfig::new()
.with_ttl(Duration::from_secs(1800))
.with_prefix("myapp");
assert_eq!(config.default_ttl, Duration::from_secs(1800));
assert_eq!(config.prefix, "myapp");
}
#[test]
fn test_prefixed_key() {
let config = CacheConfig::new().with_prefix("test");
let cache = Cache::with_config(Arc::new(crate::stores::MemoryStore::new()), config);
assert_eq!(cache.prefixed_key("key"), "test:key");
}
#[test]
fn test_prefixed_key_no_prefix() {
let cache = Cache::new(Arc::new(crate::stores::MemoryStore::new()));
assert_eq!(cache.prefixed_key("key"), "key");
}
}