djangors_cache/
lib.rs

1#![deny(missing_docs)]
2//! Raw-byte caches, common backends, and explicitly opted-in response caching.
3
4use async_trait::async_trait;
5use chrono::{DateTime, Utc};
6use http_body_util::Full;
7use hyper::{Request, Response};
8use moka::future::Cache as MokaCache;
9use serde::{de::DeserializeOwned, Serialize};
10use std::{convert::Infallible, future::Future, sync::Arc, time::Duration};
11use thiserror::Error;
12use tower::{Layer, Service};
13
14/// Errors encountered during cache access or value serialization.
15#[derive(Debug, Error)]
16pub enum CacheError {
17    /// An error reported by the underlying cache backend.
18    #[error("cache backend error: {0}")]
19    Backend(String),
20    /// An error serializing or deserializing cached values.
21    #[error("cache serialization error: {0}")]
22    Serialization(String),
23}
24
25/// An object-safe cache of raw byte values.
26#[async_trait]
27pub trait Cache: Send + Sync {
28    /// Retrieves the raw byte value associated with `key` if present and unexpired.
29    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError>;
30    /// Stores `value` under `key` with an optional time-to-live (`ttl`).
31    async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>)
32        -> Result<(), CacheError>;
33    /// Deletes any entry associated with `key`.
34    async fn delete(&self, key: &str) -> Result<(), CacheError>;
35}
36
37/// Extension trait providing convenient `get_or_set` helpers on any [`Cache`].
38#[async_trait]
39pub trait CacheExt: Cache {
40    /// Retrieves `key` if cached; otherwise executes `f`, caches the result under `key`, and returns it.
41    async fn get_or_set<F, Fut>(
42        &self,
43        key: &str,
44        ttl: Option<Duration>,
45        f: F,
46    ) -> Result<Vec<u8>, CacheError>
47    where
48        F: FnOnce() -> Fut + Send,
49        Fut: Future<Output = Result<Vec<u8>, CacheError>> + Send,
50    {
51        if let Some(value) = self.get(key).await? {
52            return Ok(value);
53        }
54        let value = f().await?;
55        self.set(key, value.clone(), ttl).await?;
56        Ok(value)
57    }
58
59    /// Retrieves and deserializes `key` from JSON; otherwise executes `f`, serializes to JSON, caches, and returns `T`.
60    async fn get_or_set_json<T, F, Fut>(
61        &self,
62        key: &str,
63        ttl: Option<Duration>,
64        f: F,
65    ) -> Result<T, CacheError>
66    where
67        T: Serialize + DeserializeOwned + Send,
68        F: FnOnce() -> Fut + Send,
69        Fut: Future<Output = Result<T, CacheError>> + Send,
70    {
71        let bytes = self
72            .get_or_set(key, ttl, || async {
73                let value = f().await?;
74                serde_json::to_vec(&value).map_err(|e| CacheError::Serialization(e.to_string()))
75            })
76            .await?;
77        serde_json::from_slice(&bytes).map_err(|e| CacheError::Serialization(e.to_string()))
78    }
79}
80impl<T: Cache + ?Sized> CacheExt for T {}
81
82/// Helper function to retrieve or compute a template fragment or raw byte cache entry.
83pub async fn get_or_set_fragment<F, Fut>(
84    cache: &dyn Cache,
85    key: &str,
86    ttl: Option<Duration>,
87    f: F,
88) -> Result<Vec<u8>, CacheError>
89where
90    F: FnOnce() -> Fut + Send,
91    Fut: Future<Output = Result<Vec<u8>, CacheError>> + Send,
92{
93    cache.get_or_set(key, ttl, f).await
94}
95
96/// An in-memory cache backend powered by Moka.
97#[derive(Clone)]
98pub struct InMemoryCache {
99    inner: MokaCache<String, MemoryEntry>,
100}
101#[derive(Clone)]
102struct MemoryEntry {
103    value: Vec<u8>,
104    expires_at: Option<std::time::Instant>,
105}
106impl InMemoryCache {
107    /// Creates a new `InMemoryCache` instance with `max_capacity` elements limit.
108    pub fn new(max_capacity: u64) -> Self {
109        Self {
110            inner: MokaCache::builder().max_capacity(max_capacity).build(),
111        }
112    }
113}
114impl Default for InMemoryCache {
115    fn default() -> Self {
116        Self::new(10_000)
117    }
118}
119#[async_trait]
120impl Cache for InMemoryCache {
121    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
122        Ok(self.inner.get(key).await.and_then(|entry| {
123            if entry
124                .expires_at
125                .is_none_or(|e| e > std::time::Instant::now())
126            {
127                Some(entry.value)
128            } else {
129                None
130            }
131        }))
132    }
133    async fn set(
134        &self,
135        key: &str,
136        value: Vec<u8>,
137        ttl: Option<Duration>,
138    ) -> Result<(), CacheError> {
139        self.inner
140            .insert(
141                key.to_owned(),
142                MemoryEntry {
143                    value,
144                    expires_at: ttl.map(|t| std::time::Instant::now() + t),
145                },
146            )
147            .await;
148        Ok(())
149    }
150    async fn delete(&self, key: &str) -> Result<(), CacheError> {
151        self.inner.invalidate(key).await;
152        Ok(())
153    }
154}
155
156/// A database-backed cache backend storing entries in `djangors_cache_entries`.
157#[derive(Clone)]
158pub struct DatabaseCache {
159    db: djangors_db::Database,
160}
161impl DatabaseCache {
162    /// Creates a new `DatabaseCache` backend using `db`.
163    pub fn new(db: djangors_db::Database) -> Self {
164        Self { db }
165    }
166    async fn ensure_table(&self) -> Result<(), CacheError> {
167        sqlx::query("CREATE TABLE IF NOT EXISTS djangors_cache_entries (key TEXT PRIMARY KEY, value BYTEA NOT NULL, expires_at TIMESTAMPTZ)").execute(self.db.pool()).await.map(|_| ()).map_err(|e| CacheError::Backend(e.to_string()))
168    }
169}
170#[async_trait]
171impl Cache for DatabaseCache {
172    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
173        self.ensure_table().await?;
174        let row = sqlx::query_as::<_, (Vec<u8>, Option<DateTime<Utc>>)>(
175            "SELECT value, expires_at FROM djangors_cache_entries WHERE key = $1",
176        )
177        .bind(key)
178        .fetch_optional(self.db.pool())
179        .await
180        .map_err(|e| CacheError::Backend(e.to_string()))?;
181        if let Some((value, expiry)) = row {
182            if expiry.is_none_or(|e| e > Utc::now()) {
183                return Ok(Some(value));
184            }
185            self.delete(key).await?;
186        }
187        Ok(None)
188    }
189    async fn set(
190        &self,
191        key: &str,
192        value: Vec<u8>,
193        ttl: Option<Duration>,
194    ) -> Result<(), CacheError> {
195        self.ensure_table().await?;
196        let expires = ttl.map(|t| Utc::now() + chrono::Duration::from_std(t).unwrap_or_default());
197        sqlx::query("INSERT INTO djangors_cache_entries (key, value, expires_at) VALUES ($1, $2, $3) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, expires_at = EXCLUDED.expires_at").bind(key).bind(value).bind(expires).execute(self.db.pool()).await.map(|_| ()).map_err(|e| CacheError::Backend(e.to_string()))
198    }
199    async fn delete(&self, key: &str) -> Result<(), CacheError> {
200        self.ensure_table().await?;
201        sqlx::query("DELETE FROM djangors_cache_entries WHERE key = $1")
202            .bind(key)
203            .execute(self.db.pool())
204            .await
205            .map(|_| ())
206            .map_err(|e| CacheError::Backend(e.to_string()))
207    }
208}
209
210#[cfg(feature = "redis")]
211/// A Redis-backed cache implementation.
212#[derive(Clone)]
213pub struct RedisCache {
214    client: redis::Client,
215}
216#[cfg(feature = "redis")]
217impl RedisCache {
218    /// Creates a new `RedisCache` connecting to the given Redis URL string.
219    pub fn new(url: &str) -> Result<Self, CacheError> {
220        redis::Client::open(url)
221            .map(|client| Self { client })
222            .map_err(|e| CacheError::Backend(e.to_string()))
223    }
224    async fn connection(&self) -> Result<redis::aio::MultiplexedConnection, CacheError> {
225        self.client
226            .get_multiplexed_async_connection()
227            .await
228            .map_err(|e| CacheError::Backend(e.to_string()))
229    }
230}
231#[cfg(feature = "redis")]
232#[async_trait]
233impl Cache for RedisCache {
234    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
235        let mut c = self.connection().await?;
236        redis::cmd("GET")
237            .arg(key)
238            .query_async(&mut c)
239            .await
240            .map_err(|e| CacheError::Backend(e.to_string()))
241    }
242    async fn set(
243        &self,
244        key: &str,
245        value: Vec<u8>,
246        ttl: Option<Duration>,
247    ) -> Result<(), CacheError> {
248        let mut c = self.connection().await?;
249        let mut cmd = redis::cmd("SET");
250        cmd.arg(key).arg(value);
251        if let Some(t) = ttl {
252            cmd.arg("PX").arg(t.as_millis() as u64);
253        }
254        cmd.query_async::<()>(&mut c)
255            .await
256            .map_err(|e| CacheError::Backend(e.to_string()))
257    }
258    async fn delete(&self, key: &str) -> Result<(), CacheError> {
259        let mut c = self.connection().await?;
260        redis::cmd("DEL")
261            .arg(key)
262            .query_async::<()>(&mut c)
263            .await
264            .map_err(|e| CacheError::Backend(e.to_string()))
265    }
266}
267
268/// Request/response extension marker indicating that a response is cacheable by [`CacheLayer`].
269#[derive(Clone, Copy, Debug)]
270pub struct CacheableResponse;
271
272/// Tower middleware layer for caching responses marked with [`CacheableResponse`].
273#[derive(Clone)]
274pub struct CacheLayer {
275    cache: Arc<dyn Cache>,
276    ttl: Option<Duration>,
277}
278impl CacheLayer {
279    /// Creates a new `CacheLayer` using `cache` and optional default `ttl`.
280    pub fn new(cache: Arc<dyn Cache>, ttl: Option<Duration>) -> Self {
281        Self { cache, ttl }
282    }
283}
284impl<S> Layer<S> for CacheLayer
285where
286    S: Service<
287            Request<Full<bytes::Bytes>>,
288            Response = Response<Full<bytes::Bytes>>,
289            Error = Infallible,
290        > + Clone
291        + Send
292        + 'static,
293    S::Future: Send + 'static,
294{
295    type Service = CacheService<S>;
296    fn layer(&self, inner: S) -> Self::Service {
297        CacheService {
298            inner,
299            cache: self.cache.clone(),
300            ttl: self.ttl,
301        }
302    }
303}
304
305/// Tower service middleware that intercepts GET requests and serves/caches responses.
306#[derive(Clone)]
307pub struct CacheService<S> {
308    inner: S,
309    cache: Arc<dyn Cache>,
310    ttl: Option<Duration>,
311}
312impl<S> Service<Request<Full<bytes::Bytes>>> for CacheService<S>
313where
314    S: Service<
315            Request<Full<bytes::Bytes>>,
316            Response = Response<Full<bytes::Bytes>>,
317            Error = Infallible,
318        > + Clone
319        + Send
320        + 'static,
321    S::Future: Send + 'static,
322{
323    type Response = Response<Full<bytes::Bytes>>;
324    type Error = Infallible;
325    type Future =
326        std::pin::Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
327    fn poll_ready(
328        &mut self,
329        cx: &mut std::task::Context<'_>,
330    ) -> std::task::Poll<Result<(), Self::Error>> {
331        self.inner.poll_ready(cx)
332    }
333    fn call(&mut self, req: Request<Full<bytes::Bytes>>) -> Self::Future {
334        let mut inner = self.inner.clone();
335        let cache = self.cache.clone();
336        let ttl = self.ttl;
337        Box::pin(async move {
338            if req.method() != hyper::Method::GET {
339                return inner.call(req).await;
340            }
341            let key = format!("{}?{}", req.uri().path(), req.uri().query().unwrap_or(""));
342            if let Some(bytes) = cache.get(&key).await.unwrap_or(None) {
343                if let Ok(cached) = serde_json::from_slice::<CachedResponse>(&bytes) {
344                    return Ok(cached.into_response());
345                }
346            }
347            let response = inner.call(req).await?;
348            if response.extensions().get::<CacheableResponse>().is_some()
349                && !response.headers().contains_key(hyper::header::SET_COOKIE)
350            {
351                let cached = CachedResponse::from_response(&response);
352                if let Ok(bytes) = serde_json::to_vec(&cached) {
353                    let _ = cache.set(&key, bytes, ttl).await;
354                }
355            }
356            Ok(response)
357        })
358    }
359}
360#[derive(serde::Serialize, serde::Deserialize)]
361struct CachedResponse {
362    status: u16,
363    headers: Vec<(String, Vec<u8>)>,
364    body: Vec<u8>,
365}
366impl CachedResponse {
367    fn from_response(r: &Response<Full<bytes::Bytes>>) -> Self {
368        Self {
369            status: r.status().as_u16(),
370            headers: r
371                .headers()
372                .iter()
373                .map(|(n, v)| (n.to_string(), v.as_bytes().to_vec()))
374                .collect(),
375            body: r.body().clone().into_inner().unwrap_or_default().to_vec(),
376        }
377    }
378    fn into_response(self) -> Response<Full<bytes::Bytes>> {
379        let mut r = Response::builder()
380            .status(self.status)
381            .body(Full::new(bytes::Bytes::from(self.body)))
382            .unwrap();
383        for (n, v) in self.headers {
384            if let (Ok(n), Ok(v)) = (
385                n.parse::<hyper::header::HeaderName>(),
386                hyper::header::HeaderValue::from_bytes(&v),
387            ) {
388                r.headers_mut().append(n, v);
389            }
390        }
391        r.extensions_mut().insert(CacheableResponse);
392        r
393    }
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399    use bytes::Bytes;
400    use std::{
401        convert::Infallible,
402        sync::{
403            atomic::{AtomicUsize, Ordering},
404            Arc,
405        },
406    };
407    use tower::{service_fn, ServiceExt};
408
409    async fn exercise(cache: &dyn Cache) {
410        assert_eq!(cache.get("missing").await.unwrap(), None);
411        cache
412            .set("key", b"value".to_vec(), Some(Duration::from_secs(5)))
413            .await
414            .unwrap();
415        assert_eq!(cache.get("key").await.unwrap(), Some(b"value".to_vec()));
416        assert_eq!(
417            cache
418                .get_or_set("key", None, || async { Ok(b"new".to_vec()) })
419                .await
420                .unwrap(),
421            b"value"
422        );
423        cache.delete("key").await.unwrap();
424        assert_eq!(
425            cache
426                .get_or_set("key", None, || async { Ok(b"new".to_vec()) })
427                .await
428                .unwrap(),
429            b"new"
430        );
431        cache
432            .set("ttl", b"gone".to_vec(), Some(Duration::from_millis(50)))
433            .await
434            .unwrap();
435        tokio::time::sleep(Duration::from_millis(80)).await;
436        assert_eq!(cache.get("ttl").await.unwrap(), None);
437    }
438    #[tokio::test]
439    async fn in_memory_cache_operations_and_ttl() {
440        exercise(&InMemoryCache::default()).await;
441    }
442    #[tokio::test]
443    async fn cache_layer_requires_response_opt_in() {
444        let count = Arc::new(AtomicUsize::new(0));
445        let c = count.clone();
446        let handler = service_fn(move |_req: Request<Full<Bytes>>| {
447            let c = c.clone();
448            async move {
449                c.fetch_add(1, Ordering::SeqCst);
450                let mut r = Response::new(Full::new(Bytes::from_static(b"cached")));
451                r.extensions_mut().insert(CacheableResponse);
452                Ok::<_, Infallible>(r)
453            }
454        });
455        let mut svc = CacheLayer::new(Arc::new(InMemoryCache::default()), None).layer(handler);
456        let request = || {
457            Request::builder()
458                .method("GET")
459                .uri("/fragment?x=1")
460                .body(Full::new(Bytes::new()))
461                .unwrap()
462        };
463        assert_eq!(
464            svc.ready()
465                .await
466                .unwrap()
467                .call(request())
468                .await
469                .unwrap()
470                .body()
471                .clone()
472                .into_inner(),
473            Some(Bytes::from_static(b"cached"))
474        );
475        svc.ready().await.unwrap().call(request()).await.unwrap();
476        assert_eq!(count.load(Ordering::SeqCst), 1);
477    }
478    #[cfg(feature = "redis")]
479    #[tokio::test]
480    async fn redis_cache_operations_and_ttl() {
481        exercise(&RedisCache::new("redis://127.0.0.1:6379").unwrap()).await;
482    }
483    #[tokio::test]
484    async fn database_cache_operations_and_ttl() {
485        let url = std::env::var("DATABASE_URL")
486            .unwrap_or_else(|_| "postgres://postgres:postgres@localhost/djangors_test".into());
487        let db = djangors_db::Database::connect(&djangors_db::DatabaseConfig::new(url))
488            .await
489            .unwrap();
490        let cache = DatabaseCache::new(db);
491        sqlx::query("DROP TABLE IF EXISTS djangors_cache_entries")
492            .execute(cache.db.pool())
493            .await
494            .unwrap();
495        exercise(&cache).await;
496        sqlx::query("DROP TABLE IF EXISTS djangors_cache_entries")
497            .execute(cache.db.pool())
498            .await
499            .unwrap();
500    }
501}
djangors_cache/lib.rs

djangors_cache/
lib.rs
