1use dashmap::DashMap;
8use std::sync::Arc;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::time::Duration;
11
12use super::{CacheBackend, CacheEntry, CacheRead};
13use crate::error::CacheError;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum PromotionStrategy {
18 HitCount { threshold: u64 },
20
21 HitRate { threshold_per_minute: u64 },
23}
24
25impl Default for PromotionStrategy {
26 fn default() -> Self {
27 Self::HitCount { threshold: 3 }
28 }
29}
30
31#[derive(Debug, Clone, Default)]
33pub struct TierStats {
34 pub l1_hits: u64,
35 pub l2_hits: u64,
36 pub misses: u64,
37 pub promotions: u64,
38}
39
40#[derive(Debug, Clone)]
42pub struct MultiTierConfig {
43 pub promotion_strategy: PromotionStrategy,
45
46 pub write_through: bool,
48
49 pub max_l1_entry_size: Option<usize>,
52}
53
54impl Default for MultiTierConfig {
55 fn default() -> Self {
56 Self {
57 promotion_strategy: PromotionStrategy::default(),
58 write_through: true,
59 max_l1_entry_size: Some(256 * 1024), }
61 }
62}
63
64struct KeyStats {
66 l2_hits: AtomicU64,
67}
68
69impl KeyStats {
70 fn new() -> Self {
71 Self {
72 l2_hits: AtomicU64::new(0),
73 }
74 }
75
76 fn record_hit(&self) -> u64 {
77 self.l2_hits.fetch_add(1, Ordering::Relaxed) + 1
78 }
79
80 fn reset(&self) {
81 self.l2_hits.store(0, Ordering::Relaxed);
82 }
83
84 fn hits(&self) -> u64 {
85 self.l2_hits.load(Ordering::Relaxed)
86 }
87}
88
89#[derive(Clone)]
94pub struct MultiTierBackend<L1, L2> {
95 l1: L1,
96 l2: L2,
97 config: MultiTierConfig,
98 key_stats: Arc<DashMap<String, Arc<KeyStats>>>,
99 tier_stats: Arc<TierStats>,
100}
101
102impl<L1, L2> MultiTierBackend<L1, L2>
103where
104 L1: CacheBackend,
105 L2: CacheBackend,
106{
107 pub fn new(l1: L1, l2: L2) -> Self {
109 Self {
110 l1,
111 l2,
112 config: MultiTierConfig::default(),
113 key_stats: Arc::new(DashMap::new()),
114 tier_stats: Arc::new(TierStats::default()),
115 }
116 }
117
118 pub fn builder() -> MultiTierBuilder<L1, L2> {
120 MultiTierBuilder::new()
121 }
122
123 pub fn l1(&self) -> &L1 {
125 &self.l1
126 }
127
128 pub fn l2(&self) -> &L2 {
130 &self.l2
131 }
132
133 pub fn stats(&self) -> &TierStats {
135 &self.tier_stats
136 }
137
138 fn should_promote(&self, key: &str) -> bool {
140 let stats = self
141 .key_stats
142 .entry(key.to_string())
143 .or_insert_with(|| Arc::new(KeyStats::new()));
144
145 match self.config.promotion_strategy {
146 PromotionStrategy::HitCount { threshold } => stats.hits() >= threshold,
147 PromotionStrategy::HitRate {
148 threshold_per_minute: _,
149 } => {
150 stats.hits() >= 3
153 }
154 }
155 }
156
157 fn record_hit(&self, key: &str) -> u64 {
159 self.key_stats
160 .entry(key.to_string())
161 .or_insert_with(|| Arc::new(KeyStats::new()))
162 .record_hit()
163 }
164
165 #[allow(dead_code)]
167 async fn promote(
168 &self,
169 key: &str,
170 entry: CacheEntry,
171 ttl: Duration,
172 stale_for: Duration,
173 ) -> Result<(), CacheError> {
174 self.l1.set(key.to_string(), entry, ttl, stale_for).await?;
176
177 if let Some(stats) = self.key_stats.get(key) {
179 stats.reset();
180 }
181
182 Ok(())
183 }
184}
185
186impl<L1, L2> CacheBackend for MultiTierBackend<L1, L2>
187where
188 L1: CacheBackend,
189 L2: CacheBackend,
190{
191 async fn get(&self, key: &str) -> Result<Option<CacheRead>, CacheError> {
192 if let Some(entry) = self.l1.get(key).await? {
194 #[cfg(feature = "metrics")]
195 metrics::counter!("tower_http_cache.tier.l1_hit").increment(1);
196 return Ok(Some(entry));
197 }
198
199 if let Some(read) = self.l2.get(key).await? {
201 #[cfg(feature = "metrics")]
202 metrics::counter!("tower_http_cache.tier.l2_hit").increment(1);
203
204 self.record_hit(key);
206
207 if self.should_promote(key) {
208 let entry_size = read.entry.body.len();
209
210 let should_promote_l1 = if let Some(max_size) = self.config.max_l1_entry_size {
212 entry_size <= max_size
213 } else {
214 true
215 };
216
217 if should_promote_l1 {
218 #[cfg(feature = "metrics")]
219 metrics::counter!("tower_http_cache.tier.promoted").increment(1);
220
221 let ttl = if let Some(expires_at) = read.expires_at {
223 expires_at
224 .duration_since(std::time::SystemTime::now())
225 .unwrap_or(Duration::from_secs(60))
226 } else {
227 Duration::from_secs(60)
228 };
229
230 let stale_for = if let (Some(stale_until), Some(expires_at)) =
231 (read.stale_until, read.expires_at)
232 {
233 stale_until.duration_since(expires_at).unwrap_or_default()
234 } else {
235 Duration::ZERO
236 };
237
238 let entry = read.entry.clone();
240 let key = key.to_string();
241 let l1 = self.l1.clone();
242 let key_stats = self.key_stats.clone();
243
244 tokio::spawn(async move {
245 let _ = l1.set(key.clone(), entry, ttl, stale_for).await;
246 if let Some(stats) = key_stats.get(&key) {
247 stats.reset();
248 }
249 });
250 } else {
251 #[cfg(feature = "metrics")]
252 metrics::counter!("tower_http_cache.tier.promotion_skipped_large").increment(1);
253
254 #[cfg(feature = "tracing")]
255 tracing::debug!(
256 key = %key,
257 size = entry_size,
258 max_l1_size = ?self.config.max_l1_entry_size,
259 "skipping promotion for large entry"
260 );
261 }
262 }
263
264 return Ok(Some(read));
265 }
266
267 Ok(None)
268 }
269
270 async fn set(
271 &self,
272 key: String,
273 entry: CacheEntry,
274 ttl: Duration,
275 stale_for: Duration,
276 ) -> Result<(), CacheError> {
277 let entry_size = entry.body.len();
278
279 self.l2
281 .set(key.clone(), entry.clone(), ttl, stale_for)
282 .await?;
283
284 if self.config.write_through {
286 let mut should_write_l1 = true;
291 if let Some(max_size) = self.config.max_l1_entry_size {
292 if entry_size > max_size {
293 #[cfg(feature = "metrics")]
294 metrics::counter!("tower_http_cache.tier.l1_skipped_large").increment(1);
295
296 #[cfg(feature = "tracing")]
297 tracing::debug!(
298 key = %key,
299 size = entry_size,
300 max_l1_size = max_size,
301 "skipping L1 write for large entry"
302 );
303
304 should_write_l1 = false;
305 }
306 }
307
308 if should_write_l1 {
309 let _ = self.l1.set(key.clone(), entry, ttl, stale_for).await;
310 }
311 }
312
313 Ok(())
314 }
315
316 async fn invalidate(&self, key: &str) -> Result<(), CacheError> {
317 let l1_result = self.l1.invalidate(key).await;
319 let l2_result = self.l2.invalidate(key).await;
320
321 self.key_stats.remove(key);
323
324 l1_result.and(l2_result)
326 }
327
328 async fn get_keys_by_tag(&self, tag: &str) -> Result<Vec<String>, CacheError> {
336 let l1 = self.l1.get_keys_by_tag(tag).await;
337 let l2 = self.l2.get_keys_by_tag(tag).await;
338
339 let (mut keys, l2_keys) = match merge_tiers(l1, l2)? {
340 Some(pair) => pair,
341 None => return Ok(Vec::new()),
342 };
343
344 keys.extend(l2_keys);
346 keys.sort();
347 keys.dedup();
348
349 Ok(keys)
350 }
351
352 async fn invalidate_by_tag(&self, tag: &str) -> Result<usize, CacheError> {
355 let l1 = self.l1.invalidate_by_tag(tag).await;
356 let l2 = self.l2.invalidate_by_tag(tag).await;
357
358 match merge_tiers(l1, l2)? {
359 Some((l1_count, l2_count)) => Ok(l1_count + l2_count),
360 None => Ok(0),
361 }
362 }
363
364 async fn list_tags(&self) -> Result<Vec<String>, CacheError> {
367 let l1 = self.l1.list_tags().await;
368 let l2 = self.l2.list_tags().await;
369
370 let (mut tags, l2_tags) = match merge_tiers(l1, l2)? {
371 Some(pair) => pair,
372 None => return Ok(Vec::new()),
373 };
374
375 tags.extend(l2_tags);
376 tags.sort();
377 tags.dedup();
378
379 Ok(tags)
380 }
381}
382
383fn merge_tiers<T: Default>(
390 l1: Result<T, CacheError>,
391 l2: Result<T, CacheError>,
392) -> Result<Option<(T, T)>, CacheError> {
393 match (l1, l2) {
394 (Ok(a), Ok(b)) => Ok(Some((a, b))),
395 (Ok(a), Err(e)) if e.is_unsupported() => Ok(Some((a, T::default()))),
396 (Err(e), Ok(b)) if e.is_unsupported() => Ok(Some((T::default(), b))),
397 (Err(a), Err(b)) if a.is_unsupported() && b.is_unsupported() => Ok(None),
398 (Err(e), _) | (_, Err(e)) => Err(e),
399 }
400}
401
402pub struct MultiTierBuilder<L1, L2> {
404 l1: Option<L1>,
405 l2: Option<L2>,
406 config: MultiTierConfig,
407}
408
409impl<L1, L2> MultiTierBuilder<L1, L2> {
410 pub fn new() -> Self {
412 Self {
413 l1: None,
414 l2: None,
415 config: MultiTierConfig::default(),
416 }
417 }
418
419 pub fn l1(mut self, backend: L1) -> Self {
421 self.l1 = Some(backend);
422 self
423 }
424
425 pub fn l2(mut self, backend: L2) -> Self {
427 self.l2 = Some(backend);
428 self
429 }
430
431 pub fn promotion_strategy(mut self, strategy: PromotionStrategy) -> Self {
433 self.config.promotion_strategy = strategy;
434 self
435 }
436
437 pub fn promotion_threshold(mut self, threshold: u64) -> Self {
439 self.config.promotion_strategy = PromotionStrategy::HitCount { threshold };
440 self
441 }
442
443 pub fn write_through(mut self, enabled: bool) -> Self {
445 self.config.write_through = enabled;
446 self
447 }
448
449 pub fn max_l1_entry_size(mut self, size: Option<usize>) -> Self {
452 self.config.max_l1_entry_size = size;
453 self
454 }
455
456 pub fn build(self) -> MultiTierBackend<L1, L2> {
458 MultiTierBackend {
459 l1: self.l1.expect("L1 backend is required"),
460 l2: self.l2.expect("L2 backend is required"),
461 config: self.config,
462 key_stats: Arc::new(DashMap::new()),
463 tier_stats: Arc::new(TierStats::default()),
464 }
465 }
466}
467
468impl<L1, L2> Default for MultiTierBuilder<L1, L2> {
469 fn default() -> Self {
470 Self::new()
471 }
472}
473
474#[cfg(all(test, feature = "in-memory"))]
475mod tests {
476 use super::*;
477 use crate::backend::memory::InMemoryBackend;
478 use bytes::Bytes;
479 use http::{StatusCode, Version};
480
481 fn test_entry() -> CacheEntry {
482 CacheEntry::new(
483 StatusCode::OK,
484 Version::HTTP_11,
485 Vec::new(),
486 Bytes::from_static(b"test"),
487 )
488 }
489
490 #[tokio::test]
491 async fn multi_tier_l1_hit() {
492 let l1 = InMemoryBackend::new(100);
493 let l2 = InMemoryBackend::new(1000);
494 let backend = MultiTierBackend::new(l1.clone(), l2);
495
496 l1.set(
498 "key".to_string(),
499 test_entry(),
500 Duration::from_secs(60),
501 Duration::ZERO,
502 )
503 .await
504 .unwrap();
505
506 let result = backend.get("key").await.unwrap();
508 assert!(result.is_some());
509 }
510
511 #[tokio::test]
512 async fn multi_tier_l2_hit_and_promote() {
513 let l1 = InMemoryBackend::new(100);
514 let l2 = InMemoryBackend::new(1000);
515
516 let backend = MultiTierBackend::builder()
517 .l1(l1.clone())
518 .l2(l2.clone())
519 .promotion_threshold(3)
520 .build();
521
522 l2.set(
524 "key".to_string(),
525 test_entry(),
526 Duration::from_secs(60),
527 Duration::ZERO,
528 )
529 .await
530 .unwrap();
531
532 for _ in 0..3 {
534 let result = backend.get("key").await.unwrap();
535 assert!(result.is_some());
536 }
537
538 tokio::time::sleep(Duration::from_millis(50)).await;
540
541 let l1_result = l1.get("key").await.unwrap();
543 assert!(l1_result.is_some());
544 }
545
546 #[tokio::test]
547 async fn multi_tier_set_writes_to_both_tiers() {
548 let l1 = InMemoryBackend::new(100);
549 let l2 = InMemoryBackend::new(1000);
550 let backend = MultiTierBackend::builder()
551 .l1(l1.clone())
552 .l2(l2.clone())
553 .write_through(true)
554 .build();
555
556 backend
557 .set(
558 "key".to_string(),
559 test_entry(),
560 Duration::from_secs(60),
561 Duration::ZERO,
562 )
563 .await
564 .unwrap();
565
566 assert!(l1.get("key").await.unwrap().is_some());
568 assert!(l2.get("key").await.unwrap().is_some());
569 }
570
571 #[tokio::test]
572 async fn multi_tier_invalidate_both_tiers() {
573 let l1 = InMemoryBackend::new(100);
574 let l2 = InMemoryBackend::new(1000);
575 let backend = MultiTierBackend::new(l1.clone(), l2.clone());
576
577 l1.set(
579 "key".to_string(),
580 test_entry(),
581 Duration::from_secs(60),
582 Duration::ZERO,
583 )
584 .await
585 .unwrap();
586 l2.set(
587 "key".to_string(),
588 test_entry(),
589 Duration::from_secs(60),
590 Duration::ZERO,
591 )
592 .await
593 .unwrap();
594
595 backend.invalidate("key").await.unwrap();
597
598 assert!(l1.get("key").await.unwrap().is_none());
600 assert!(l2.get("key").await.unwrap().is_none());
601 }
602
603 #[tokio::test]
604 async fn multi_tier_miss() {
605 let l1 = InMemoryBackend::new(100);
606 let l2 = InMemoryBackend::new(1000);
607 let backend = MultiTierBackend::new(l1, l2);
608
609 let result = backend.get("nonexistent").await.unwrap();
610 assert!(result.is_none());
611 }
612
613 #[tokio::test]
614 async fn promotion_strategy_hit_count() {
615 let strategy = PromotionStrategy::HitCount { threshold: 5 };
616 let l1 = InMemoryBackend::new(100);
617 let l2 = InMemoryBackend::new(1000);
618
619 let backend = MultiTierBackend::builder()
620 .l1(l1.clone())
621 .l2(l2.clone())
622 .promotion_strategy(strategy)
623 .build();
624
625 l2.set(
626 "key".to_string(),
627 test_entry(),
628 Duration::from_secs(60),
629 Duration::ZERO,
630 )
631 .await
632 .unwrap();
633
634 for _ in 0..5 {
636 backend.get("key").await.unwrap();
637 }
638
639 tokio::time::sleep(Duration::from_millis(50)).await;
640
641 assert!(l1.get("key").await.unwrap().is_some());
643 }
644}
645
646#[cfg(test)]
647mod merge_tier_tests {
648 use super::merge_tiers;
649 use crate::error::CacheError;
650
651 fn unsupported<T>() -> Result<T, CacheError> {
652 Err(CacheError::Unsupported("no tag index".to_string()))
653 }
654
655 #[test]
656 fn both_tiers_answer() {
657 let merged = merge_tiers(Ok(vec!["a".to_string()]), Ok(vec!["b".to_string()])).unwrap();
658 assert_eq!(merged, Some((vec!["a".to_string()], vec!["b".to_string()])));
659 }
660
661 #[test]
664 fn an_unsupported_tier_contributes_nothing_instead_of_failing() {
665 let merged = merge_tiers(Ok(vec!["a".to_string()]), unsupported()).unwrap();
666 assert_eq!(merged, Some((vec!["a".to_string()], Vec::new())));
667
668 let merged = merge_tiers(unsupported(), Ok(vec!["b".to_string()])).unwrap();
669 assert_eq!(merged, Some((Vec::new(), vec!["b".to_string()])));
670 }
671
672 #[test]
673 fn neither_tier_has_an_index() {
674 assert_eq!(
675 merge_tiers::<Vec<String>>(unsupported(), unsupported()).unwrap(),
676 None
677 );
678 }
679
680 #[test]
681 fn a_real_error_still_propagates() {
682 let err = merge_tiers::<Vec<String>>(
683 Ok(Vec::new()),
684 Err(CacheError::Backend("connection reset".to_string())),
685 )
686 .unwrap_err();
687 assert!(!err.is_unsupported());
688 }
689}