use async_trait::async_trait;
use super::StorageBackend;
use crate::CacheError;
const NARINFO_PREFIX: &str = "sui:narinfo:";
const NAR_PREFIX: &str = "sui:nar:";
#[async_trait]
pub trait RedisConn: Send + Sync {
async fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError>;
async fn set_bytes(&self, key: &str, value: &[u8], ttl_secs: Option<u64>) -> Result<(), CacheError>;
async fn del(&self, key: &str) -> Result<(), CacheError>;
async fn keys_with_prefix(&self, prefix: &str) -> Result<Vec<String>, CacheError>;
}
pub struct RedisBackend<C: RedisConn> {
conn: C,
ttl_secs: Option<u64>,
}
impl<C: RedisConn> RedisBackend<C> {
pub fn new(conn: C) -> Self {
Self { conn, ttl_secs: None }
}
pub fn with_ttl(conn: C, ttl_secs: u64) -> Self {
Self { conn, ttl_secs: Some(ttl_secs) }
}
#[must_use]
pub fn ttl_secs(&self) -> Option<u64> {
self.ttl_secs
}
pub fn conn(&self) -> &C {
&self.conn
}
fn narinfo_key(hash: &str) -> String {
format!("{NARINFO_PREFIX}{hash}")
}
fn nar_key(path: &str) -> String {
format!("{NAR_PREFIX}{path}")
}
}
#[async_trait]
impl<C: RedisConn> StorageBackend for RedisBackend<C> {
async fn get_narinfo(&self, hash: &str) -> Result<Option<String>, CacheError> {
let key = Self::narinfo_key(hash);
match self.conn.get_bytes(&key).await? {
Some(bytes) => {
let text = String::from_utf8(bytes)
.map_err(|e| CacheError::NarInfo(format!("invalid utf-8 in redis narinfo {hash}: {e}")))?;
Ok(Some(text))
}
None => Ok(None),
}
}
async fn put_narinfo(&self, hash: &str, content: &str) -> Result<(), CacheError> {
let key = Self::narinfo_key(hash);
self.conn.set_bytes(&key, content.as_bytes(), self.ttl_secs).await
}
async fn get_nar(&self, path: &str) -> Result<Option<Vec<u8>>, CacheError> {
let key = Self::nar_key(path);
self.conn.get_bytes(&key).await
}
async fn put_nar(&self, path: &str, data: &[u8]) -> Result<(), CacheError> {
let key = Self::nar_key(path);
self.conn.set_bytes(&key, data, self.ttl_secs).await
}
async fn delete(&self, hash: &str) -> Result<(), CacheError> {
self.conn.del(&Self::narinfo_key(hash)).await?;
for ext in ["nar.xz", "nar.zst", "nar"] {
self.conn.del(&Self::nar_key(&format!("nar/{hash}.{ext}"))).await?;
}
Ok(())
}
async fn list_narinfos(&self) -> Result<Vec<String>, CacheError> {
let keys = self.conn.keys_with_prefix(NARINFO_PREFIX).await?;
Ok(keys
.into_iter()
.filter_map(|k| k.strip_prefix(NARINFO_PREFIX).map(str::to_string))
.collect())
}
async fn wipe_all(&self) -> Result<usize, CacheError> {
let narinfos = self.conn.keys_with_prefix(NARINFO_PREFIX).await?;
let n = narinfos.len();
for key in &narinfos {
self.conn.del(key).await?;
}
for key in self.conn.keys_with_prefix(NAR_PREFIX).await? {
self.conn.del(&key).await?;
}
Ok(n)
}
}
#[cfg(feature = "redis-client")]
mod client {
use super::{CacheError, RedisBackend, RedisConn};
use async_trait::async_trait;
fn to_cache_err(e: redis::RedisError) -> CacheError {
CacheError::Io(std::io::Error::other(format!("redis: {e}")))
}
#[derive(Clone)]
pub struct RedisConnectionManager {
mgr: redis::aio::ConnectionManager,
}
impl RedisConnectionManager {
pub async fn connect(url: &str) -> Result<Self, CacheError> {
let client = redis::Client::open(url).map_err(to_cache_err)?;
let mgr = redis::aio::ConnectionManager::new(client)
.await
.map_err(to_cache_err)?;
Ok(Self { mgr })
}
}
#[async_trait]
impl RedisConn for RedisConnectionManager {
async fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
let mut c = self.mgr.clone();
let v: Option<Vec<u8>> = redis::cmd("GET")
.arg(key)
.query_async(&mut c)
.await
.map_err(to_cache_err)?;
Ok(v)
}
async fn set_bytes(&self, key: &str, value: &[u8], ttl_secs: Option<u64>) -> Result<(), CacheError> {
let mut c = self.mgr.clone();
let mut cmd = redis::cmd("SET");
cmd.arg(key).arg(value);
if let Some(secs) = ttl_secs {
cmd.arg("EX").arg(secs);
}
let _: () = cmd.query_async(&mut c).await.map_err(to_cache_err)?;
Ok(())
}
async fn del(&self, key: &str) -> Result<(), CacheError> {
let mut c = self.mgr.clone();
let _: i64 = redis::cmd("DEL")
.arg(key)
.query_async(&mut c)
.await
.map_err(to_cache_err)?;
Ok(())
}
async fn keys_with_prefix(&self, prefix: &str) -> Result<Vec<String>, CacheError> {
let mut c = self.mgr.clone();
let pattern = format!("{prefix}*");
let mut cursor: u64 = 0;
let mut out = Vec::new();
loop {
let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(512)
.query_async(&mut c)
.await
.map_err(to_cache_err)?;
out.extend(batch);
if next == 0 {
break;
}
cursor = next;
}
Ok(out)
}
}
impl RedisBackend<RedisConnectionManager> {
pub async fn connect(url: &str) -> Result<Self, CacheError> {
Ok(Self::new(RedisConnectionManager::connect(url).await?))
}
pub async fn connect_with_ttl(url: &str, ttl_secs: u64) -> Result<Self, CacheError> {
Ok(Self::with_ttl(RedisConnectionManager::connect(url).await?, ttl_secs))
}
}
}
#[cfg(feature = "redis-client")]
pub use client::RedisConnectionManager;
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
use std::sync::Mutex;
#[derive(Default)]
struct MockRedis {
map: Mutex<HashMap<String, (Vec<u8>, Option<u64>)>>,
}
impl MockRedis {
fn ttl_of(&self, key: &str) -> Option<u64> {
self.map.lock().unwrap().get(key).and_then(|(_, t)| *t)
}
fn evict(&self, key: &str) {
self.map.lock().unwrap().remove(key);
}
fn clear(&self) {
self.map.lock().unwrap().clear();
}
fn len(&self) -> usize {
self.map.lock().unwrap().len()
}
}
#[async_trait]
impl RedisConn for MockRedis {
async fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
Ok(self.map.lock().unwrap().get(key).map(|(v, _)| v.clone()))
}
async fn set_bytes(&self, key: &str, value: &[u8], ttl_secs: Option<u64>) -> Result<(), CacheError> {
self.map
.lock()
.unwrap()
.insert(key.to_string(), (value.to_vec(), ttl_secs));
Ok(())
}
async fn del(&self, key: &str) -> Result<(), CacheError> {
self.map.lock().unwrap().remove(key);
Ok(())
}
async fn keys_with_prefix(&self, prefix: &str) -> Result<Vec<String>, CacheError> {
Ok(self
.map
.lock()
.unwrap()
.keys()
.filter(|k| k.starts_with(prefix))
.cloned()
.collect())
}
}
const NARINFO: &str = "StorePath: /nix/store/abc-hello\nURL: nar/abc.nar.xz\nCompression: xz\nNarHash: sha256:bbb\nNarSize: 200\nReferences: \n";
#[tokio::test]
async fn get_missing_narinfo_returns_none() {
let backend = RedisBackend::new(MockRedis::default());
assert!(backend.get_narinfo("nope").await.unwrap().is_none());
}
#[tokio::test]
async fn put_then_get_narinfo_roundtrips() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("abc", NARINFO).await.unwrap();
let got = backend.get_narinfo("abc").await.unwrap().unwrap();
assert_eq!(got, NARINFO);
}
#[tokio::test]
async fn get_missing_nar_returns_none() {
let backend = RedisBackend::new(MockRedis::default());
assert!(backend.get_nar("nar/missing.nar.xz").await.unwrap().is_none());
}
#[tokio::test]
async fn put_then_get_nar_roundtrips() {
let backend = RedisBackend::new(MockRedis::default());
let data = b"\x00\x01\x02 fake nar bytes";
backend.put_nar("nar/abc.nar.xz", data).await.unwrap();
let got = backend.get_nar("nar/abc.nar.xz").await.unwrap().unwrap();
assert_eq!(got, data);
}
#[tokio::test]
async fn narinfo_and_nar_keyspaces_do_not_collide() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("dead", "the-narinfo").await.unwrap();
backend.put_nar("dead", b"the-nar").await.unwrap();
assert_eq!(backend.get_narinfo("dead").await.unwrap().unwrap(), "the-narinfo");
assert_eq!(backend.get_nar("dead").await.unwrap().unwrap(), b"the-nar");
}
#[tokio::test]
async fn no_ttl_by_default() {
let mock = MockRedis::default();
let backend = RedisBackend::new(mock);
assert_eq!(backend.ttl_secs(), None);
backend.put_narinfo("abc", NARINFO).await.unwrap();
assert_eq!(backend.conn().ttl_of("sui:narinfo:abc"), None);
}
#[tokio::test]
async fn with_ttl_stamps_every_write() {
let backend = RedisBackend::with_ttl(MockRedis::default(), 3600);
assert_eq!(backend.ttl_secs(), Some(3600));
backend.put_narinfo("abc", NARINFO).await.unwrap();
backend.put_nar("nar/abc.nar.xz", b"data").await.unwrap();
assert_eq!(backend.conn().ttl_of("sui:narinfo:abc"), Some(3600));
assert_eq!(backend.conn().ttl_of("sui:nar:nar/abc.nar.xz"), Some(3600));
}
#[tokio::test]
async fn eviction_of_a_hot_key_is_a_plain_miss() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("abc", NARINFO).await.unwrap();
assert!(backend.get_narinfo("abc").await.unwrap().is_some());
backend.conn().evict("sui:narinfo:abc");
assert!(backend.get_narinfo("abc").await.unwrap().is_none());
}
#[tokio::test]
async fn pod_roll_clears_the_whole_hot_tier() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("a", NARINFO).await.unwrap();
backend.put_nar("nar/a.nar.xz", b"x").await.unwrap();
backend.conn().clear();
assert!(backend.get_narinfo("a").await.unwrap().is_none());
assert!(backend.get_nar("nar/a.nar.xz").await.unwrap().is_none());
}
#[tokio::test]
async fn delete_removes_narinfo_and_common_nar_patterns() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("xyz", NARINFO).await.unwrap();
backend.put_nar("nar/xyz.nar.xz", b"nar-xz").await.unwrap();
backend.put_nar("nar/xyz.nar.zst", b"nar-zst").await.unwrap();
backend.put_nar("nar/xyz.nar", b"nar-plain").await.unwrap();
backend.delete("xyz").await.unwrap();
assert!(backend.get_narinfo("xyz").await.unwrap().is_none());
assert!(backend.get_nar("nar/xyz.nar.xz").await.unwrap().is_none());
assert!(backend.get_nar("nar/xyz.nar.zst").await.unwrap().is_none());
assert!(backend.get_nar("nar/xyz.nar").await.unwrap().is_none());
}
#[tokio::test]
async fn delete_absent_is_idempotent() {
let backend = RedisBackend::new(MockRedis::default());
backend.delete("ghost").await.unwrap();
assert_eq!(backend.conn().len(), 0);
}
#[tokio::test]
async fn list_narinfos_returns_hot_subset_stripped() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("aaa", "1").await.unwrap();
backend.put_narinfo("bbb", "2").await.unwrap();
backend.put_nar("nar/ccc.nar.xz", b"3").await.unwrap();
let mut hashes = backend.list_narinfos().await.unwrap();
hashes.sort();
assert_eq!(hashes, vec!["aaa".to_string(), "bbb".to_string()]);
}
#[tokio::test]
async fn list_narinfos_empty_when_cold() {
let backend = RedisBackend::new(MockRedis::default());
assert!(backend.list_narinfos().await.unwrap().is_empty());
}
#[tokio::test]
async fn overwrite_narinfo_takes_latest() {
let backend = RedisBackend::new(MockRedis::default());
backend.put_narinfo("h", "v1").await.unwrap();
backend.put_narinfo("h", "v2").await.unwrap();
assert_eq!(backend.get_narinfo("h").await.unwrap().unwrap(), "v2");
}
#[tokio::test]
async fn invalid_utf8_narinfo_surfaces_typed_error() {
let mock = MockRedis::default();
mock.map
.lock()
.unwrap()
.insert("sui:narinfo:bad".to_string(), (vec![0xff, 0xfe, 0xfd], None));
let backend = RedisBackend::new(mock);
let err = backend.get_narinfo("bad").await.unwrap_err();
assert!(matches!(err, CacheError::NarInfo(_)));
}
}