use std::collections::HashMap; use std::{ io::{Error, ErrorKind}, path::Path, sync::Arc, }; use async_trait::async_trait; use parking_lot::RwLock; use serde::{Deserialize, Serialize}; use yoke::{Yoke, Yokeable}; use crate::config::{CacheStore, Config}; pub const BINCODE_CONFIG: bincode::config::Configuration = bincode::config::standard(); /// Marker trait for values that can be stored within the cache pub trait Cacheable: Serialize + Send + for<'a> Yokeable<'a> { /// The key used to uniquely identify a cache item. type Key<'a>: Send + 'a; /// A unique kind for the `Cacheable` used to prefix the `Key`. const KIND: CacheKind; /// Builds the key to store in the cache by prefixing `Self::format_key` with `Self::KIND`. fn build_key(k: Self::Key<'_>) -> Vec { let mut key = Vec::new(); key.push(Self::KIND as u8); Self::format_key(&mut key, k); key } /// Serializes `k` to `out`. fn format_key(out: &mut Vec, k: Self::Key<'_>); } /// A unique prefix for each type stored within the cache to prevent conflicts. #[repr(u8)] pub enum CacheKind { Eligibility = 1, CrateMetadata = 2, } /// A generic-erased `Cache`. #[derive(Clone)] pub enum ConcreteCache { RocksDb(RocksDb), InMemory(InMemory), } impl ConcreteCache { /// Instantiates a new `Cache`. pub fn new(config: &Config) -> Result { Ok(match &config.cache { CacheStore::RocksDb { path } => Self::RocksDb(RocksDb::new(path)?), CacheStore::InMemory => Self::InMemory(InMemory::default()), }) } } #[async_trait] impl Cache for ConcreteCache { async fn put(&self, key: C::Key<'_>, value: &C) -> Result<(), Error> { match self { Self::RocksDb(r) => r.put(key, value).await, Self::InMemory(i) => i.put(key, value).await, } } async fn get( &self, key: C::Key<'_>, ) -> Result>>, Error> where for<'a> >::Output: Deserialize<'a>, { match self { Self::RocksDb(r) => r.get(key).await, Self::InMemory(i) => i.get(key).await, } } async fn remove(&self, key: C::Key<'_>) -> Result<(), Error> { match self { Self::RocksDb(r) => r.remove::(key).await, Self::InMemory(i) => i.remove::(key).await, } } } #[async_trait] pub trait Cache { /// Inserts a value into the cache. async fn put(&self, key: C::Key<'_>, value: &C) -> Result<(), Error>; /// Retrieves a value from the cache. async fn get( &self, key: C::Key<'_>, ) -> Result>>, Error> where for<'a> >::Output: Deserialize<'a>; /// Removes a value from the cache. async fn remove(&self, key: C::Key<'_>) -> Result<(), Error>; } #[derive(Clone, Default)] #[allow(clippy::type_complexity)] pub struct InMemory { db: Arc, Box<[u8]>>>>, } #[async_trait] impl Cache for InMemory { async fn put(&self, key: C::Key<'_>, value: &C) -> Result<(), Error> { let serialized = bincode::serde::encode_to_vec(value, BINCODE_CONFIG) .map_err(|e| Error::new(ErrorKind::Other, e))?; let key = C::build_key(key); self.db .write() .insert(key.into_boxed_slice(), serialized.into_boxed_slice()); Ok(()) } async fn get( &self, key: C::Key<'_>, ) -> Result>>, Error> where for<'a> >::Output: Deserialize<'a>, { let key = C::build_key(key); let Some(value) = self.db.read().get(key.as_slice()).map(|v| v.to_vec()) else { return Ok(None); }; Yoke::try_attach_to_cart(value, |v| { bincode::serde::decode_borrowed_from_slice(v, BINCODE_CONFIG) }) .map(Some) .map_err(|e| Error::new(ErrorKind::Other, e)) } async fn remove(&self, key: C::Key<'_>) -> Result<(), Error> { self.db.write().remove(C::build_key(key).as_slice()); Ok(()) } } #[derive(Clone)] pub struct RocksDb { rocks: Arc, } impl RocksDb { pub fn new(path: &Path) -> Result { let rocks = rocksdb::DB::open_default(path).map_err(|e| Error::new(ErrorKind::Other, e))?; Ok(Self { rocks: Arc::new(rocks), }) } } #[async_trait] impl Cache for RocksDb { async fn put(&self, key: C::Key<'_>, value: &C) -> Result<(), Error> { let serialized = bincode::serde::encode_to_vec(value, BINCODE_CONFIG) .map_err(|e| Error::new(ErrorKind::Other, e))?; let rocks = self.rocks.clone(); let key = C::build_key(key); tokio::task::spawn_blocking(move || { rocks .put(key, serialized) .map_err(|e| Error::new(ErrorKind::Other, e)) }) .await .map_err(|e| Error::new(ErrorKind::Other, e))? } async fn get( &self, key: C::Key<'_>, ) -> Result>>, Error> where for<'a> >::Output: Deserialize<'a>, { let rocks = self.rocks.clone(); let key = C::build_key(key); tokio::task::spawn_blocking(move || { rocks .get(key) .map_err(|e| Error::new(ErrorKind::Other, e))? .map(|v| { Yoke::try_attach_to_cart(v, |v| { bincode::serde::decode_borrowed_from_slice(v, BINCODE_CONFIG) }) }) .transpose() .map_err(|e| Error::new(ErrorKind::Other, e)) }) .await .map_err(|e| Error::new(ErrorKind::Other, e))? } async fn remove(&self, key: C::Key<'_>) -> Result<(), Error> { let rocks = self.rocks.clone(); let key = C::build_key(key); tokio::task::spawn_blocking(move || { rocks .delete(key) .map_err(|e| Error::new(ErrorKind::Other, e)) }) .await .map_err(|e| Error::new(ErrorKind::Other, e))? } } pub type Yoked = Yoke>; #[cfg(test)] mod test { use crate::cache::{Cache, InMemory, RocksDb}; use crate::providers::{EligibilityCacheKey, Release}; use std::borrow::Cow; use tempfile::tempdir; async fn test_suite(cache: T) { let out = cache .get::>>(EligibilityCacheKey::new( "my-project", "my-crate", "my-crate-version", )) .await .unwrap(); assert!(out.is_none()); cache .put( EligibilityCacheKey::new("my-project", "my-crate", "my-crate-version"), &None, ) .await .unwrap(); let out = cache .get::>>(EligibilityCacheKey::new( "my-project", "my-crate", "my-crate-version", )) .await .unwrap(); assert!(out.unwrap().get().is_none()); cache .put( EligibilityCacheKey::new("my-project", "my-crate", "my-crate-version"), &Some(Release { name: Cow::Borrowed("helloworld"), version: Cow::Borrowed("1.0.0"), checksum: Cow::Borrowed("123456"), project: Cow::Borrowed("test"), yanked: false, }), ) .await .unwrap(); let out = cache .get::>>(EligibilityCacheKey::new( "my-project", "my-crate", "my-crate-version", )) .await .unwrap(); assert_eq!( out.unwrap().get().as_ref().unwrap().name.as_ref(), "helloworld" ); let out = cache .get::>>(EligibilityCacheKey::new( "my-project", "my-crate", "my-crate-version-2", )) .await .unwrap(); assert!(out.is_none()); } #[tokio::test] async fn rocksdb() { let temp_dir = tempdir().unwrap(); let cache = RocksDb::new(temp_dir.path()).unwrap(); test_suite(cache).await; } #[tokio::test] async fn in_memory() { let cache = InMemory::default(); test_suite(cache).await; } }