use std::time::Duration;
use rainier_support::{BoxFuture, Error, Result};
use serde::de::DeserializeOwned;
use serde::Serialize;
pub trait Cache: Send + Sync + 'static {
fn name(&self) -> &str;
fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Result<Option<Vec<u8>>>>;
fn put<'a>(
&'a self,
key: &'a str,
value: &'a [u8],
ttl: Option<Duration>,
) -> BoxFuture<'a, Result<()>>;
fn forget<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Result<bool>>;
fn flush(&self) -> BoxFuture<'_, Result<()>>;
fn is_shared(&self) -> bool {
true
}
fn supports_atomic_add(&self) -> bool {
true
}
fn has<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Result<bool>> {
Box::pin(async move { Ok(self.get(key).await?.is_some()) })
}
fn increment<'a>(&'a self, key: &'a str, by: i64) -> BoxFuture<'a, Result<i64>>;
fn add<'a>(
&'a self,
key: &'a str,
value: &'a [u8],
ttl: Option<Duration>,
) -> BoxFuture<'a, Result<bool>>;
fn forget_if<'a>(&'a self, key: &'a str, expected: &'a [u8]) -> BoxFuture<'a, Result<bool>>;
}
#[async_trait::async_trait]
pub trait CacheExt: Cache {
async fn get_json<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>> {
match self.get(key).await? {
Some(bytes) => match serde_json::from_slice(&bytes) {
Ok(value) => Ok(Some(value)),
Err(e) => {
tracing::warn!(key, error = %e, "discarding an unreadable cached value");
Ok(None)
}
},
None => Ok(None),
}
}
async fn put_json<T: Serialize + Sync>(
&self,
key: &str,
value: &T,
ttl: Option<Duration>,
) -> Result<()> {
let bytes = serde_json::to_vec(value)?;
self.put(key, &bytes, ttl).await
}
async fn get_string(&self, key: &str) -> Result<Option<String>> {
match self.get(key).await? {
Some(bytes) => Ok(String::from_utf8(bytes).ok()),
None => Ok(None),
}
}
async fn put_string(&self, key: &str, value: &str, ttl: Option<Duration>) -> Result<()> {
self.put(key, value.as_bytes(), ttl).await
}
async fn add_string(&self, key: &str, value: &str, ttl: Option<Duration>) -> Result<bool> {
self.add(key, value.as_bytes(), ttl).await
}
async fn decrement(&self, key: &str, by: i64) -> Result<i64> {
self.increment(key, -by).await
}
async fn remember<T, F, Fut>(&self, key: &str, ttl: Option<Duration>, compute: F) -> Result<T>
where
T: Serialize + DeserializeOwned + Sync + Send,
F: FnOnce() -> Fut + Send,
Fut: std::future::Future<Output = Result<T>> + Send,
{
if let Some(cached) = self.get_json::<T>(key).await? {
return Ok(cached);
}
let value = compute().await?;
if let Err(e) = self.put_json(key, &value, ttl).await {
tracing::warn!(key, error = %e.message(), "could not cache a computed value");
}
Ok(value)
}
async fn remember_forever<T, F, Fut>(&self, key: &str, compute: F) -> Result<T>
where
T: Serialize + DeserializeOwned + Sync + Send,
F: FnOnce() -> Fut + Send,
Fut: std::future::Future<Output = Result<T>> + Send,
{
self.remember(key, None, compute).await
}
}
impl<C: Cache + ?Sized> CacheExt for C {}
pub(crate) fn encode_counter(value: i64) -> Vec<u8> {
value.to_string().into_bytes()
}
pub(crate) fn decode_counter(bytes: &[u8]) -> Result<i64> {
if bytes.is_empty() {
return Ok(0);
}
std::str::from_utf8(bytes)
.ok()
.and_then(|text| text.trim().parse::<i64>().ok())
.ok_or_else(|| Error::internal("the cached value is not a counter"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn counters_round_trip_as_decimal_text() {
for value in [0i64, 1, -1, i64::MAX, i64::MIN] {
assert_eq!(decode_counter(&encode_counter(value)).unwrap(), value);
}
assert_eq!(encode_counter(42), b"42".to_vec(), "readable by redis-cli");
}
#[test]
fn an_empty_value_counts_as_zero() {
assert_eq!(decode_counter(b"").unwrap(), 0);
}
#[test]
fn a_non_counter_is_an_error_rather_than_zero() {
assert!(decode_counter(b"hello").is_err());
}
#[tokio::test]
async fn remember_computes_once_and_then_reads_the_cache() {
use crate::MemoryCache;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
let cache = MemoryCache::new();
let calls = Arc::new(AtomicUsize::new(0));
for _ in 0..3 {
let calls = Arc::clone(&calls);
let value: String = cache
.remember("settings", None, || async move {
calls.fetch_add(1, Ordering::SeqCst);
Ok("computed".to_string())
})
.await
.unwrap();
assert_eq!(value, "computed");
}
assert_eq!(calls.load(Ordering::SeqCst), 1, "it recomputed a value it already had");
}
#[tokio::test]
async fn remember_does_not_cache_a_failure() {
use crate::MemoryCache;
let cache = MemoryCache::new();
let failed: Result<String> = cache
.remember("thing", None, || async { Err(Error::internal("upstream is down")) })
.await;
assert!(failed.is_err());
assert!(cache.get("thing").await.unwrap().is_none(), "the failure was cached");
let recovered: String = cache
.remember("thing", None, || async { Ok("it works now".to_string()) })
.await
.unwrap();
assert_eq!(recovered, "it works now");
}
#[tokio::test]
async fn remember_recomputes_a_value_it_can_no_longer_read() {
use crate::MemoryCache;
let cache = MemoryCache::new();
cache.put("count", b"\"not a number\"", None).await.unwrap();
let count: u32 = cache.remember("count", None, || async { Ok(7) }).await.unwrap();
assert_eq!(count, 7);
}
}