1#![deny(missing_docs)]
2use async_trait::async_trait;
15use chrono::Utc;
16use http_body_util::Full;
17use hyper::{Request, Response};
18use moka::future::Cache as MokaCache;
19use serde::{de::DeserializeOwned, Serialize};
20use std::{convert::Infallible, future::Future, sync::Arc, time::Duration};
21use thiserror::Error;
22use tower::{Layer, Service};
23
24#[derive(Debug, Error)]
26pub enum CacheError {
27 #[error("cache backend error: {0}")]
29 Backend(String),
30 #[error("cache serialization error: {0}")]
32 Serialization(String),
33}
34
35#[async_trait]
37pub trait Cache: Send + Sync {
38 async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError>;
40 async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>)
42 -> Result<(), CacheError>;
43 async fn delete(&self, key: &str) -> Result<(), CacheError>;
45}
46
47#[async_trait]
49pub trait CacheExt: Cache {
50 async fn get_or_set<F, Fut>(
52 &self,
53 key: &str,
54 ttl: Option<Duration>,
55 f: F,
56 ) -> Result<Vec<u8>, CacheError>
57 where
58 F: FnOnce() -> Fut + Send,
59 Fut: Future<Output = Result<Vec<u8>, CacheError>> + Send,
60 {
61 if let Some(value) = self.get(key).await? {
62 return Ok(value);
63 }
64 let value = f().await?;
65 self.set(key, value.clone(), ttl).await?;
66 Ok(value)
67 }
68
69 async fn get_or_set_json<T, F, Fut>(
71 &self,
72 key: &str,
73 ttl: Option<Duration>,
74 f: F,
75 ) -> Result<T, CacheError>
76 where
77 T: Serialize + DeserializeOwned + Send,
78 F: FnOnce() -> Fut + Send,
79 Fut: Future<Output = Result<T, CacheError>> + Send,
80 {
81 let bytes = self
82 .get_or_set(key, ttl, || async {
83 let value = f().await?;
84 serde_json::to_vec(&value).map_err(|e| CacheError::Serialization(e.to_string()))
85 })
86 .await?;
87 serde_json::from_slice(&bytes).map_err(|e| CacheError::Serialization(e.to_string()))
88 }
89}
90impl<T: Cache + ?Sized> CacheExt for T {}
91
92pub async fn get_or_set_fragment<F, Fut>(
94 cache: &dyn Cache,
95 key: &str,
96 ttl: Option<Duration>,
97 f: F,
98) -> Result<Vec<u8>, CacheError>
99where
100 F: FnOnce() -> Fut + Send,
101 Fut: Future<Output = Result<Vec<u8>, CacheError>> + Send,
102{
103 cache.get_or_set(key, ttl, f).await
104}
105
106#[derive(Clone)]
108pub struct InMemoryCache {
109 inner: MokaCache<String, MemoryEntry>,
110}
111#[derive(Clone)]
113struct MemoryEntry {
114 value: Vec<u8>,
116 expires_at: Option<std::time::Instant>,
118}
119impl InMemoryCache {
120 pub fn new(max_capacity: u64) -> Self {
122 Self {
123 inner: MokaCache::builder().max_capacity(max_capacity).build(),
124 }
125 }
126}
127impl Default for InMemoryCache {
128 fn default() -> Self {
129 Self::new(10_000)
130 }
131}
132#[async_trait]
133impl Cache for InMemoryCache {
134 async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
135 Ok(self.inner.get(key).await.and_then(|entry| {
136 if entry
137 .expires_at
138 .is_none_or(|e| e > std::time::Instant::now())
139 {
140 Some(entry.value)
141 } else {
142 None
143 }
144 }))
145 }
146 async fn set(
147 &self,
148 key: &str,
149 value: Vec<u8>,
150 ttl: Option<Duration>,
151 ) -> Result<(), CacheError> {
152 self.inner
153 .insert(
154 key.to_owned(),
155 MemoryEntry {
156 value,
157 expires_at: ttl.map(|t| std::time::Instant::now() + t),
158 },
159 )
160 .await;
161 Ok(())
162 }
163 async fn delete(&self, key: &str) -> Result<(), CacheError> {
164 self.inner.invalidate(key).await;
165 Ok(())
166 }
167}
168
169#[derive(Clone)]
171pub struct DatabaseCache {
172 db: djangors_db::Database,
173}
174impl DatabaseCache {
175 pub fn new(db: djangors_db::Database) -> Self {
177 Self { db }
178 }
179 async fn ensure_table(&self) -> Result<(), CacheError> {
182 let mut conn = self.db.conn();
183 let dialect = conn.dialect();
184 let bytea = dialect.bytea_type();
185 let sql = format!(
186 "CREATE TABLE IF NOT EXISTS djangors_cache_entries (key TEXT PRIMARY KEY, value {bytea} NOT NULL, expires_at TIMESTAMPTZ)"
187 );
188 conn.execute(&sql, &[])
189 .await
190 .map(|_| ())
191 .map_err(|e| CacheError::Backend(e.to_string()))
192 }
193}
194#[async_trait]
195impl Cache for DatabaseCache {
196 async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
197 self.ensure_table().await?;
198 let mut conn = self.db.conn();
199 let p1 = conn.dialect().placeholder(1);
200 let sql = format!("SELECT value, expires_at FROM djangors_cache_entries WHERE key = {p1}");
201 let params = vec![djangors_db::BindValue::Text(key.to_string())];
202 let row_opt = conn
203 .fetch_optional(&sql, ¶ms)
204 .await
205 .map_err(|e| CacheError::Backend(e.to_string()))?;
206 if let Some(row) = row_opt {
207 let value = row
208 .try_bytes(0)
209 .map_err(|e| CacheError::Backend(e.to_string()))?
210 .unwrap_or_default();
211 let expiry = row
212 .try_datetime(1)
213 .map_err(|e| CacheError::Backend(e.to_string()))?;
214 if expiry.is_none_or(|e| e > Utc::now()) {
215 return Ok(Some(value));
216 }
217 self.delete(key).await?;
218 }
219 Ok(None)
220 }
221 async fn set(
222 &self,
223 key: &str,
224 value: Vec<u8>,
225 ttl: Option<Duration>,
226 ) -> Result<(), CacheError> {
227 self.ensure_table().await?;
228 let expires = ttl.map(|t| Utc::now() + chrono::Duration::from_std(t).unwrap_or_default());
229 let mut conn = self.db.conn();
230 let dialect = conn.dialect();
231 let p1 = dialect.placeholder(1);
232 let p2 = dialect.placeholder(2);
233 let p3 = dialect.placeholder(3);
234 let sql = format!(
235 "INSERT INTO djangors_cache_entries (key, value, expires_at) VALUES ({p1}, {p2}, {p3}) \
236 ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, expires_at = EXCLUDED.expires_at"
237 );
238 let expires_param = match expires {
239 Some(dt) => djangors_db::BindValue::DateTime(dt),
240 None => djangors_db::BindValue::Null(djangors_db::NullKind::DateTime),
241 };
242 let params = vec![
243 djangors_db::BindValue::Text(key.to_string()),
244 djangors_db::BindValue::Bytes(value),
245 expires_param,
246 ];
247 conn.execute(&sql, ¶ms)
248 .await
249 .map(|_| ())
250 .map_err(|e| CacheError::Backend(e.to_string()))
251 }
252 async fn delete(&self, key: &str) -> Result<(), CacheError> {
253 self.ensure_table().await?;
254 let mut conn = self.db.conn();
255 let p1 = conn.dialect().placeholder(1);
256 let sql = format!("DELETE FROM djangors_cache_entries WHERE key = {p1}");
257 let params = vec![djangors_db::BindValue::Text(key.to_string())];
258 conn.execute(&sql, ¶ms)
259 .await
260 .map(|_| ())
261 .map_err(|e| CacheError::Backend(e.to_string()))
262 }
263}
264
265#[cfg(feature = "redis")]
266#[derive(Clone)]
268pub struct RedisCache {
269 client: redis::Client,
270}
271#[cfg(feature = "redis")]
272impl RedisCache {
273 pub fn new(url: &str) -> Result<Self, CacheError> {
275 redis::Client::open(url)
276 .map(|client| Self { client })
277 .map_err(|e| CacheError::Backend(e.to_string()))
278 }
279 async fn connection(&self) -> Result<redis::aio::MultiplexedConnection, CacheError> {
280 self.client
281 .get_multiplexed_async_connection()
282 .await
283 .map_err(|e| CacheError::Backend(e.to_string()))
284 }
285}
286#[cfg(feature = "redis")]
287#[async_trait]
288impl Cache for RedisCache {
289 async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
290 let mut c = self.connection().await?;
291 redis::cmd("GET")
292 .arg(key)
293 .query_async(&mut c)
294 .await
295 .map_err(|e| CacheError::Backend(e.to_string()))
296 }
297 async fn set(
298 &self,
299 key: &str,
300 value: Vec<u8>,
301 ttl: Option<Duration>,
302 ) -> Result<(), CacheError> {
303 let mut c = self.connection().await?;
304 let mut cmd = redis::cmd("SET");
305 cmd.arg(key).arg(value);
306 if let Some(t) = ttl {
307 cmd.arg("PX").arg(t.as_millis() as u64);
308 }
309 cmd.query_async::<()>(&mut c)
310 .await
311 .map_err(|e| CacheError::Backend(e.to_string()))
312 }
313 async fn delete(&self, key: &str) -> Result<(), CacheError> {
314 let mut c = self.connection().await?;
315 redis::cmd("DEL")
316 .arg(key)
317 .query_async::<()>(&mut c)
318 .await
319 .map_err(|e| CacheError::Backend(e.to_string()))
320 }
321}
322
323#[derive(Clone, Copy, Debug)]
325pub struct CacheableResponse;
326
327#[derive(Clone)]
329pub struct CacheLayer {
330 cache: Arc<dyn Cache>,
332 ttl: Option<Duration>,
334}
335impl CacheLayer {
336 pub fn new(cache: Arc<dyn Cache>, ttl: Option<Duration>) -> Self {
338 Self { cache, ttl }
339 }
340}
341impl<S> Layer<S> for CacheLayer
342where
343 S: Service<
344 Request<Full<bytes::Bytes>>,
345 Response = Response<Full<bytes::Bytes>>,
346 Error = Infallible,
347 > + Clone
348 + Send
349 + 'static,
350 S::Future: Send + 'static,
351{
352 type Service = CacheService<S>;
353 fn layer(&self, inner: S) -> Self::Service {
354 CacheService {
355 inner,
356 cache: self.cache.clone(),
357 ttl: self.ttl,
358 }
359 }
360}
361
362#[derive(Clone)]
364pub struct CacheService<S> {
365 inner: S,
367 cache: Arc<dyn Cache>,
369 ttl: Option<Duration>,
371}
372impl<S> Service<Request<Full<bytes::Bytes>>> for CacheService<S>
373where
374 S: Service<
375 Request<Full<bytes::Bytes>>,
376 Response = Response<Full<bytes::Bytes>>,
377 Error = Infallible,
378 > + Clone
379 + Send
380 + 'static,
381 S::Future: Send + 'static,
382{
383 type Response = Response<Full<bytes::Bytes>>;
384 type Error = Infallible;
385 type Future =
386 std::pin::Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
387 fn poll_ready(
388 &mut self,
389 cx: &mut std::task::Context<'_>,
390 ) -> std::task::Poll<Result<(), Self::Error>> {
391 self.inner.poll_ready(cx)
392 }
393 fn call(&mut self, req: Request<Full<bytes::Bytes>>) -> Self::Future {
394 let mut inner = self.inner.clone();
395 let cache = self.cache.clone();
396 let ttl = self.ttl;
397 Box::pin(async move {
398 if req.method() != hyper::Method::GET {
399 return inner.call(req).await;
400 }
401 let key = format!("{}?{}", req.uri().path(), req.uri().query().unwrap_or(""));
402 if let Some(bytes) = cache.get(&key).await.unwrap_or(None) {
403 if let Ok(cached) = serde_json::from_slice::<CachedResponse>(&bytes) {
404 return Ok(cached.into_response());
405 }
406 }
407 let response = inner.call(req).await?;
408 if response.extensions().get::<CacheableResponse>().is_some()
409 && !response.headers().contains_key(hyper::header::SET_COOKIE)
410 {
411 let cached = CachedResponse::from_response(&response);
412 if let Ok(bytes) = serde_json::to_vec(&cached) {
413 let _ = cache.set(&key, bytes, ttl).await;
414 }
415 }
416 Ok(response)
417 })
418 }
419}
420#[derive(serde::Serialize, serde::Deserialize)]
423struct CachedResponse {
424 status: u16,
426 headers: Vec<(String, Vec<u8>)>,
428 body: Vec<u8>,
430}
431impl CachedResponse {
432 fn from_response(r: &Response<Full<bytes::Bytes>>) -> Self {
433 Self {
434 status: r.status().as_u16(),
435 headers: r
436 .headers()
437 .iter()
438 .map(|(n, v)| (n.to_string(), v.as_bytes().to_vec()))
439 .collect(),
440 body: r.body().clone().into_inner().unwrap_or_default().to_vec(),
441 }
442 }
443 fn into_response(self) -> Response<Full<bytes::Bytes>> {
444 let mut r = Response::builder()
445 .status(self.status)
446 .body(Full::new(bytes::Bytes::from(self.body)))
447 .unwrap();
448 for (n, v) in self.headers {
449 if let (Ok(n), Ok(v)) = (
450 n.parse::<hyper::header::HeaderName>(),
451 hyper::header::HeaderValue::from_bytes(&v),
452 ) {
453 r.headers_mut().append(n, v);
454 }
455 }
456 r.extensions_mut().insert(CacheableResponse);
457 r
458 }
459}
460
461#[cfg(test)]
462mod tests {
463 use super::*;
464 use bytes::Bytes;
465 use std::{
466 convert::Infallible,
467 sync::{
468 atomic::{AtomicUsize, Ordering},
469 Arc,
470 },
471 };
472 use tower::{service_fn, ServiceExt};
473
474 async fn exercise(cache: &dyn Cache) {
475 assert_eq!(cache.get("missing").await.unwrap(), None);
476 cache
477 .set("key", b"value".to_vec(), Some(Duration::from_secs(5)))
478 .await
479 .unwrap();
480 assert_eq!(cache.get("key").await.unwrap(), Some(b"value".to_vec()));
481 assert_eq!(
482 cache
483 .get_or_set("key", None, || async { Ok(b"new".to_vec()) })
484 .await
485 .unwrap(),
486 b"value"
487 );
488 cache.delete("key").await.unwrap();
489 assert_eq!(
490 cache
491 .get_or_set("key", None, || async { Ok(b"new".to_vec()) })
492 .await
493 .unwrap(),
494 b"new"
495 );
496 cache
497 .set("ttl", b"gone".to_vec(), Some(Duration::from_millis(50)))
498 .await
499 .unwrap();
500 tokio::time::sleep(Duration::from_millis(80)).await;
501 assert_eq!(cache.get("ttl").await.unwrap(), None);
502 }
503 #[tokio::test]
504 async fn in_memory_cache_operations_and_ttl() {
505 exercise(&InMemoryCache::default()).await;
506 }
507 #[tokio::test]
508 async fn cache_layer_requires_response_opt_in() {
509 let count = Arc::new(AtomicUsize::new(0));
510 let c = count.clone();
511 let handler = service_fn(move |_req: Request<Full<Bytes>>| {
512 let c = c.clone();
513 async move {
514 c.fetch_add(1, Ordering::SeqCst);
515 let mut r = Response::new(Full::new(Bytes::from_static(b"cached")));
516 r.extensions_mut().insert(CacheableResponse);
517 Ok::<_, Infallible>(r)
518 }
519 });
520 let mut svc = CacheLayer::new(Arc::new(InMemoryCache::default()), None).layer(handler);
521 let request = || {
522 Request::builder()
523 .method("GET")
524 .uri("/fragment?x=1")
525 .body(Full::new(Bytes::new()))
526 .unwrap()
527 };
528 assert_eq!(
529 svc.ready()
530 .await
531 .unwrap()
532 .call(request())
533 .await
534 .unwrap()
535 .body()
536 .clone()
537 .into_inner(),
538 Some(Bytes::from_static(b"cached"))
539 );
540 svc.ready().await.unwrap().call(request()).await.unwrap();
541 assert_eq!(count.load(Ordering::SeqCst), 1);
542 }
543 #[cfg(feature = "redis")]
544 #[tokio::test]
545 async fn redis_cache_operations_and_ttl() {
546 exercise(&RedisCache::new("redis://127.0.0.1:6379").unwrap()).await;
547 }
548 #[tokio::test]
549 async fn database_cache_operations_and_ttl() {
550 let url = std::env::var("DATABASE_URL")
551 .unwrap_or_else(|_| "postgres://postgres:postgres@localhost/djangors_test".into());
552 let db = djangors_db::Database::connect(&djangors_db::DatabaseConfig::new(url))
553 .await
554 .unwrap();
555 let cache = DatabaseCache::new(db);
556 sqlx::query("DROP TABLE IF EXISTS djangors_cache_entries")
557 .execute(cache.db.pool())
558 .await
559 .unwrap();
560 exercise(&cache).await;
561 sqlx::query("DROP TABLE IF EXISTS djangors_cache_entries")
562 .execute(cache.db.pool())
563 .await
564 .unwrap();
565 }
566}