tower_http_cache/
layer.rs

1use std::error::Error as StdError;
2use std::sync::Arc;
3use std::task::{Context, Poll};
4use std::time::{Duration, SystemTime};
5
6use bytes::Bytes;
7use dashmap::mapref::entry::Entry;
8use dashmap::DashMap;
9use futures_util::future::BoxFuture;
10use http::header::{CACHE_CONTROL, PRAGMA};
11use http::{HeaderMap, Method, Request, Response, Uri};
12use http_body::Body;
13use http_body_util::combinators::BoxBody;
14use http_body_util::{BodyExt, Full};
15use tokio::sync::{Mutex, OwnedMutexGuard};
16use tower::{Layer, Service, ServiceExt};
17
18#[cfg(feature = "metrics")]
19use metrics::{counter, histogram};
20
21#[cfg(feature = "in-memory")]
22use crate::backend::memory::InMemoryBackend;
23use crate::backend::{CacheBackend, CacheEntry, CacheRead};
24use crate::chunks::{ChunkCache, ChunkMetadata};
25#[cfg(feature = "compression")]
26use crate::policy::CompressionStrategy;
27use crate::policy::{CachePolicy, CompressionConfig};
28use crate::range::{is_partial_content, parse_range_header, RangeHandling};
29use crate::refresh::{AutoRefreshConfig, RefreshCallback, RefreshManager, RefreshMetadata};
30#[cfg(feature = "tracing")]
31use crate::streaming::extract_size_info;
32use crate::streaming::{should_stream, StreamingDecision};
33
34pub type BoxError = Box<dyn StdError + Send + Sync>;
35
36pin_project_lite::pin_project! {
37    /// Response body type that implements Sync for Axum compatibility.
38    ///
39    /// This wraps BoxBody and manually implements Sync using the same pattern as Axum.
40    /// See the `unsafe impl Sync` below for the safety justification.
41    pub struct SyncBoxBody {
42        #[pin]
43        inner: BoxBody<Bytes, BoxError>,
44    }
45}
46
47impl SyncBoxBody {
48    /// Creates a new SyncBoxBody by wrapping a BoxBody.
49    pub fn new(inner: BoxBody<Bytes, BoxError>) -> Self {
50        Self { inner }
51    }
52}
53
54// SAFETY: The inner BoxBody is Send + 'static. Access to the body is exclusively through
55// poll_frame which requires Pin<&mut Self>, preventing concurrent access without external
56// synchronization. This is the same pattern used by Axum's own Body type.
57unsafe impl Sync for SyncBoxBody {}
58
59impl Body for SyncBoxBody {
60    type Data = Bytes;
61    type Error = BoxError;
62
63    fn poll_frame(
64        self: std::pin::Pin<&mut Self>,
65        cx: &mut std::task::Context<'_>,
66    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
67        self.project().inner.poll_frame(cx)
68    }
69
70    fn is_end_stream(&self) -> bool {
71        self.inner.is_end_stream()
72    }
73
74    fn size_hint(&self) -> http_body::SizeHint {
75        self.inner.size_hint()
76    }
77}
78
79/// Type alias for the key extractor function
80type KeyExtractorFn = Arc<dyn Fn(&Method, &Uri) -> Option<String> + Send + Sync>;
81
82/// Configurable caching layer for Tower services.
83///
84/// The layer wraps an inner service and caches HTTP responses based on the
85/// configured [`CachePolicy`]. Create instances via [`CacheLayer::builder`]
86/// or [`CacheLayer::new`] for a sensible default policy.
87///
88/// Cloning a `CacheLayer` is cheap and shares the underlying backend and
89/// in-flight stampede locks.
90#[derive(Clone)]
91pub struct CacheLayer<B> {
92    backend: B,
93    policy: CachePolicy,
94    key_extractor: KeyExtractor,
95    locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
96    refresh_manager: Option<Arc<RefreshManager>>,
97    chunk_cache: Option<Arc<ChunkCache>>,
98}
99
100/// Strategy used to turn requests into cache keys.
101///
102/// The layer ships with helpers for common patterns such as
103/// [`KeyExtractor::path_and_query`] and [`KeyExtractor::path`].
104/// You can also provide your own extractor with [`KeyExtractor::custom`].
105#[derive(Clone)]
106pub struct KeyExtractor {
107    inner: KeyExtractorFn,
108}
109
110impl KeyExtractor {
111    /// Builds an extractor that uses `method + path + query` for GET/HEAD requests.
112    pub fn path_and_query() -> Self {
113        Self {
114            inner: Arc::new(|method: &Method, uri: &Uri| {
115                if matches!(method, &Method::GET | &Method::HEAD) {
116                    let mut key = uri.path().to_owned();
117                    if let Some(query) = uri.query() {
118                        key.push('?');
119                        key.push_str(query);
120                    }
121                    Some(key)
122                } else {
123                    None
124                }
125            }),
126        }
127    }
128
129    pub fn path() -> Self {
130        Self {
131            inner: Arc::new(|method: &Method, uri: &Uri| {
132                if matches!(method, &Method::GET | &Method::HEAD) {
133                    Some(uri.path().to_owned())
134                } else {
135                    None
136                }
137            }),
138        }
139    }
140
141    pub fn custom<F>(func: F) -> Self
142    where
143        F: Fn(&Method, &Uri) -> Option<String> + Send + Sync + 'static,
144    {
145        Self {
146            inner: Arc::new(func),
147        }
148    }
149
150    /// Extracts a cache key from the provided request parts.
151    ///
152    /// Returns `None` when the request should be skipped.
153    pub fn extract(&self, method: &Method, uri: &Uri) -> Option<String> {
154        (self.inner)(method, uri)
155    }
156}
157
158impl Default for KeyExtractor {
159    fn default() -> Self {
160        Self::path_and_query()
161    }
162}
163
164/// Builder for configuring [`CacheLayer`] instances.
165pub struct CacheLayerBuilder<B> {
166    backend: B,
167    policy: CachePolicy,
168    key_extractor: KeyExtractor,
169    auto_refresh_config: Option<AutoRefreshConfig>,
170}
171
172impl<B> CacheLayerBuilder<B>
173where
174    B: CacheBackend,
175{
176    pub fn new(backend: B) -> Self {
177        Self {
178            backend,
179            policy: CachePolicy::default(),
180            key_extractor: KeyExtractor::default(),
181            auto_refresh_config: None,
182        }
183    }
184
185    /// Replaces the cache policy with a pre-built value.
186    pub fn policy(mut self, policy: CachePolicy) -> Self {
187        self.policy = policy;
188        self
189    }
190
191    /// Sets the positive cache TTL for successful responses.
192    pub fn ttl(mut self, ttl: Duration) -> Self {
193        self.policy = self.policy.with_ttl(ttl);
194        self
195    }
196
197    /// Sets the cache TTL for negative (4xx) responses.
198    pub fn negative_ttl(mut self, ttl: Duration) -> Self {
199        self.policy = self.policy.with_negative_ttl(ttl);
200        self
201    }
202
203    pub fn stale_while_revalidate(mut self, duration: Duration) -> Self {
204        self.policy = self.policy.with_stale_while_revalidate(duration);
205        self
206    }
207
208    pub fn refresh_before(mut self, duration: Duration) -> Self {
209        self.policy = self.policy.with_refresh_before(duration);
210        self
211    }
212
213    pub fn max_body_size(mut self, size: Option<usize>) -> Self {
214        self.policy = self.policy.with_max_body_size(size);
215        self
216    }
217
218    pub fn min_body_size(mut self, size: Option<usize>) -> Self {
219        self.policy = self.policy.with_min_body_size(size);
220        self
221    }
222
223    pub fn allow_streaming_bodies(mut self, allow: bool) -> Self {
224        self.policy = self.policy.with_allow_streaming_bodies(allow);
225        self
226    }
227
228    pub fn compression(mut self, config: CompressionConfig) -> Self {
229        self.policy = self.policy.with_compression(config);
230        self
231    }
232
233    pub fn respect_cache_control(mut self, enabled: bool) -> Self {
234        self.policy = self.policy.with_respect_cache_control(enabled);
235        self
236    }
237
238    pub fn statuses(mut self, statuses: impl IntoIterator<Item = u16>) -> Self {
239        self.policy = self.policy.with_statuses(statuses);
240        self
241    }
242
243    pub fn method_predicate<F>(mut self, predicate: F) -> Self
244    where
245        F: Fn(&Method) -> bool + Send + Sync + 'static,
246    {
247        self.policy = self.policy.with_method_predicate(predicate);
248        self
249    }
250
251    pub fn header_allowlist<I, S>(mut self, headers: I) -> Self
252    where
253        I: IntoIterator<Item = S>,
254        S: Into<String>,
255    {
256        self.policy = self.policy.with_header_allowlist(headers);
257        self
258    }
259
260    pub fn key_extractor(mut self, extractor: KeyExtractor) -> Self {
261        self.key_extractor = extractor;
262        self
263    }
264
265    /// Enables auto-refresh functionality with the provided configuration.
266    ///
267    /// When enabled, frequently accessed cache entries will be proactively
268    /// refreshed before they expire, reducing cache misses and latency.
269    pub fn auto_refresh(mut self, config: AutoRefreshConfig) -> Self {
270        self.auto_refresh_config = Some(config);
271        self
272    }
273
274    pub fn build(self) -> CacheLayer<B> {
275        let refresh_manager = self
276            .auto_refresh_config
277            .filter(|cfg| cfg.enabled)
278            .map(|cfg| Arc::new(RefreshManager::new(cfg)));
279
280        // Create chunk cache if enabled in streaming policy
281        let chunk_cache = if self.policy.streaming_policy().enable_chunk_cache {
282            Some(Arc::new(ChunkCache::new(
283                self.policy.streaming_policy().chunk_size,
284            )))
285        } else {
286            None
287        };
288
289        CacheLayer {
290            backend: self.backend,
291            policy: self.policy,
292            key_extractor: self.key_extractor,
293            locks: Arc::new(DashMap::new()),
294            refresh_manager,
295            chunk_cache,
296        }
297    }
298}
299
300#[cfg(feature = "in-memory")]
301impl CacheLayer<InMemoryBackend> {
302    /// Creates a cache layer backed by an in-memory [`InMemoryBackend`].
303    pub fn new_in_memory(max_capacity: u64) -> Self {
304        CacheLayerBuilder::new(InMemoryBackend::new(max_capacity)).build()
305    }
306}
307
308impl<B> CacheLayer<B>
309where
310    B: CacheBackend,
311{
312    /// Builds a cache layer with the default [`CachePolicy`].
313    pub fn new(backend: B) -> Self {
314        CacheLayerBuilder::new(backend).build()
315    }
316
317    /// Returns a builder for fine-grained control over the cache policy.
318    pub fn builder(backend: B) -> CacheLayerBuilder<B> {
319        CacheLayerBuilder::new(backend)
320    }
321
322    pub fn with_policy(mut self, policy: CachePolicy) -> Self {
323        self.policy = policy;
324        self
325    }
326
327    pub fn with_ttl(mut self, ttl: Duration) -> Self {
328        self.policy = self.policy.clone().with_ttl(ttl);
329        self
330    }
331
332    pub fn with_negative_ttl(mut self, ttl: Duration) -> Self {
333        self.policy = self.policy.clone().with_negative_ttl(ttl);
334        self
335    }
336
337    pub fn with_stale_while_revalidate(mut self, duration: Duration) -> Self {
338        self.policy = self.policy.clone().with_stale_while_revalidate(duration);
339        self
340    }
341
342    pub fn with_refresh_before(mut self, duration: Duration) -> Self {
343        self.policy = self.policy.clone().with_refresh_before(duration);
344        self
345    }
346
347    pub fn with_max_body_size(mut self, size: Option<usize>) -> Self {
348        self.policy = self.policy.clone().with_max_body_size(size);
349        self
350    }
351
352    pub fn with_min_body_size(mut self, size: Option<usize>) -> Self {
353        self.policy = self.policy.clone().with_min_body_size(size);
354        self
355    }
356
357    pub fn with_allow_streaming_bodies(mut self, allow: bool) -> Self {
358        self.policy = self.policy.clone().with_allow_streaming_bodies(allow);
359        self
360    }
361
362    pub fn with_compression(mut self, config: CompressionConfig) -> Self {
363        self.policy = self.policy.clone().with_compression(config);
364        self
365    }
366
367    pub fn with_respect_cache_control(mut self, enabled: bool) -> Self {
368        self.policy = self.policy.clone().with_respect_cache_control(enabled);
369        self
370    }
371
372    pub fn with_cache_statuses(mut self, statuses: impl IntoIterator<Item = u16>) -> Self {
373        self.policy = self.policy.clone().with_statuses(statuses);
374        self
375    }
376
377    pub fn with_method_predicate<F>(mut self, predicate: F) -> Self
378    where
379        F: Fn(&Method) -> bool + Send + Sync + 'static,
380    {
381        self.policy = self.policy.clone().with_method_predicate(predicate);
382        self
383    }
384
385    pub fn with_header_allowlist<I, S>(mut self, headers: I) -> Self
386    where
387        I: IntoIterator<Item = S>,
388        S: Into<String>,
389    {
390        self.policy = self.policy.clone().with_header_allowlist(headers);
391        self
392    }
393
394    pub fn with_key_extractor(mut self, extractor: KeyExtractor) -> Self {
395        self.key_extractor = extractor;
396        self
397    }
398
399    /// Manually initialize the auto-refresh manager with a service instance.
400    ///
401    /// This should be called after constructing the service to start the background
402    /// refresh task. This is only necessary if auto-refresh is enabled.
403    ///
404    /// # Example
405    ///
406    /// ```ignore
407    /// let layer = CacheLayer::builder(backend)
408    ///     .auto_refresh(config)
409    ///     .build();
410    ///
411    /// layer.init_auto_refresh(my_service.clone()).await?;
412    /// ```
413    pub async fn init_auto_refresh<S, ResBody>(&self, service: S) -> Result<(), String>
414    where
415        S: Service<Request<()>, Response = Response<ResBody>> + Clone + Send + Sync + 'static,
416        S::Future: Send + 'static,
417        S::Error: Into<BoxError> + Send,
418        ResBody: Body<Data = Bytes> + Send + 'static,
419        ResBody::Error: Into<BoxError> + Send,
420        B: Clone,
421    {
422        if let Some(ref manager) = self.refresh_manager {
423            let callback = Arc::new(CacheRefreshCallback::new(
424                service,
425                self.backend.clone(),
426                self.policy.clone(),
427                self.key_extractor.clone(),
428            ));
429            manager.start(callback).await
430        } else {
431            Ok(())
432        }
433    }
434}
435
436impl<S, B> Layer<S> for CacheLayer<B>
437where
438    B: CacheBackend,
439{
440    type Service = CacheService<S, B>;
441
442    fn layer(&self, inner: S) -> Self::Service {
443        CacheService {
444            inner,
445            backend: self.backend.clone(),
446            policy: self.policy.clone(),
447            key_extractor: self.key_extractor.clone(),
448            locks: self.locks.clone(),
449            refresh_manager: self.refresh_manager.clone(),
450            chunk_cache: self.chunk_cache.clone(),
451        }
452    }
453}
454
455impl<B> Drop for CacheLayer<B> {
456    fn drop(&mut self) {
457        // Trigger graceful shutdown of refresh manager
458        // We use tokio::spawn to avoid blocking in Drop
459        if let Some(manager) = &self.refresh_manager {
460            let manager = manager.clone();
461            // Best-effort shutdown - spawn detached task
462            // Note: We cannot guarantee execution in Drop, but we try our best
463            if let Ok(handle) = tokio::runtime::Handle::try_current() {
464                handle.spawn(async move {
465                    manager.shutdown().await;
466                });
467            }
468        }
469    }
470}
471
472#[derive(Clone)]
473pub struct CacheService<S, B> {
474    inner: S,
475    backend: B,
476    policy: CachePolicy,
477    key_extractor: KeyExtractor,
478    locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
479    refresh_manager: Option<Arc<RefreshManager>>,
480    chunk_cache: Option<Arc<ChunkCache>>,
481}
482
483/// Implementation of RefreshCallback for CacheService.
484struct CacheRefreshCallback<S, B> {
485    inner: S,
486    backend: B,
487    policy: CachePolicy,
488}
489
490impl<S, B> CacheRefreshCallback<S, B> {
491    fn new(inner: S, backend: B, policy: CachePolicy, _key_extractor: KeyExtractor) -> Self {
492        Self {
493            inner,
494            backend,
495            policy,
496        }
497    }
498}
499
500impl<S, B, ResBody> RefreshCallback for CacheRefreshCallback<S, B>
501where
502    S: Service<Request<()>, Response = Response<ResBody>> + Clone + Send + Sync + 'static,
503    S::Future: Send + 'static,
504    S::Error: Into<BoxError> + Send,
505    ResBody: Body<Data = Bytes> + Send + 'static,
506    ResBody::Error: Into<BoxError> + Send,
507    B: CacheBackend,
508{
509    fn refresh(&self, key: String, metadata: RefreshMetadata) -> crate::refresh::RefreshFuture {
510        let backend = self.backend.clone();
511        let policy = self.policy.clone();
512        let inner = self.inner.clone();
513
514        Box::pin(async move {
515            #[cfg(feature = "tracing")]
516            tracing::debug!(key = %key, uri = %metadata.uri, "Auto-refresh triggered");
517
518            // Reconstruct the request
519            let request = match metadata.try_into_request() {
520                Some(req) => req,
521                None => {
522                    #[cfg(feature = "tracing")]
523                    tracing::warn!(key = %key, "Failed to reconstruct request for auto-refresh");
524                    return Err("Failed to reconstruct request".into());
525                }
526            };
527
528            // Call the inner service
529            let service = inner;
530            let response = match service.oneshot(request).await {
531                Ok(resp) => resp,
532                Err(_err) => {
533                    #[cfg(feature = "tracing")]
534                    tracing::error!(key = %key, "Service error during auto-refresh");
535                    return Err("Service error during refresh".into());
536                }
537            };
538
539            let (parts, body) = response.into_parts();
540
541            // Collect the body
542            let collected = match BodyExt::collect(body).await {
543                Ok(c) => c,
544                Err(_err) => {
545                    #[cfg(feature = "tracing")]
546                    tracing::error!(key = %key, "Body collection error during auto-refresh");
547                    return Err("Body collection error".into());
548                }
549            };
550
551            let cache_bytes = collected.to_bytes();
552
553            // Check if we should cache this response
554            let body_too_large = policy
555                .max_body_size()
556                .is_some_and(|max| cache_bytes.len() > max);
557            let body_too_small = policy
558                .min_body_size()
559                .is_some_and(|min| cache_bytes.len() < min);
560
561            if body_too_large || body_too_small {
562                return Ok(()); // Successfully refreshed but not stored
563            }
564
565            // Store the refreshed entry
566            if let Some(ttl) = policy.ttl_for(parts.status) {
567                if !ttl.is_zero() {
568                    let stale_for = policy.stale_while_revalidate();
569                    let headers_to_cache = policy.headers_to_cache(&parts.headers);
570                    let (compressed_bytes, _compressed) =
571                        maybe_compress(cache_bytes, policy.compression());
572
573                    let entry = CacheEntry::new(
574                        parts.status,
575                        parts.version,
576                        headers_to_cache,
577                        compressed_bytes,
578                    );
579
580                    if let Err(_err) = backend.set(key.clone(), entry, ttl, stale_for).await {
581                        #[cfg(feature = "tracing")]
582                        tracing::error!(key = %key, "Failed to store refreshed entry");
583                        return Err("Failed to store entry".into());
584                    }
585                }
586            }
587
588            Ok(())
589        })
590    }
591}
592
593// Note: We cannot easily initialize the refresh manager from within the service
594// because the service may have different type parameters than required by the callback.
595// Instead, users who want to use auto-refresh should ensure the service is called at least once,
596// or manually initialize the refresh functionality if needed.
597
598impl<S, B, ReqBody, ResBody> Service<Request<ReqBody>> for CacheService<S, B>
599where
600    S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
601    S::Future: Send + 'static,
602    S::Error: Into<BoxError> + Send,
603    ReqBody: Send + 'static,
604    ResBody: Body<Data = Bytes> + Send + Sync + 'static,
605    ResBody::Error: Into<BoxError> + Send,
606    B: CacheBackend,
607{
608    type Response = Response<SyncBoxBody>;
609    type Error = BoxError;
610    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
611
612    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
613        self.inner.poll_ready(cx).map_err(Into::into)
614    }
615
616    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
617        let method = req.method().clone();
618        let uri = req.uri().clone();
619        let should_cache_method = self.policy.should_cache_method(&method);
620        let request_bypass =
621            self.policy.respect_cache_control() && cache_control_disallows(req.headers());
622        let key = if should_cache_method && !request_bypass {
623            self.key_extractor.extract(&method, &uri)
624        } else {
625            None
626        };
627
628        let backend = self.backend.clone();
629        let policy = self.policy.clone();
630        let locks = self.locks.clone();
631        let inner = self.inner.clone();
632        let stale_window = policy.stale_while_revalidate();
633        let refresh_before = policy.refresh_before();
634        let refresh_manager = self.refresh_manager.clone();
635        let chunk_cache = self.chunk_cache.clone();
636
637        // Check for range request early
638        let range_request = parse_range_header(req.headers());
639
640        // Prepare refresh metadata if auto-refresh is enabled
641        let refresh_metadata = if refresh_manager.is_some() && key.is_some() {
642            Some(RefreshMetadata::from_request(&req))
643        } else {
644            None
645        };
646
647        Box::pin(async move {
648            #[cfg(feature = "tracing")]
649            tracing::debug!(method = %method, uri = %uri, "cache_call");
650
651            // Try to serve from chunk cache if this is a range request
652            if let (Some(range_req), Some(ref chunk_cache), Some(ref key_ref)) =
653                (range_request.as_ref(), &chunk_cache, &key)
654            {
655                if let Some(entry) = chunk_cache.get(key_ref) {
656                    // Check if range is satisfiable and normalize it
657                    if let Some(normalized) = range_req.normalize(entry.metadata.total_size) {
658                        let end = normalized.end.unwrap_or(entry.metadata.total_size - 1);
659
660                        // Try to get range from chunks
661                        if let Some(range_data) = entry.get_range(normalized.start, end) {
662                            #[cfg(feature = "metrics")]
663                            counter!("tower_http_cache.chunk_cache_hit").increment(1);
664
665                            #[cfg(feature = "tracing")]
666                            tracing::debug!(
667                                key = %key_ref,
668                                start = normalized.start,
669                                end = end,
670                                "chunk_cache_hit"
671                            );
672
673                            // Build 206 Partial Content response
674                            let mut response = Response::builder()
675                                .status(http::StatusCode::PARTIAL_CONTENT)
676                                .body(SyncBoxBody::new(
677                                    Full::from(range_data).map_err(Into::into).boxed(),
678                                ))
679                                .unwrap();
680
681                            // Copy headers from metadata
682                            for (name, value) in &entry.metadata.headers {
683                                if let (Ok(header_name), Ok(header_value)) = (
684                                    http::header::HeaderName::from_bytes(name.as_bytes()),
685                                    http::header::HeaderValue::from_bytes(value),
686                                ) {
687                                    response.headers_mut().insert(header_name, header_value);
688                                }
689                            }
690
691                            // Add Content-Range header
692                            let content_range = format!(
693                                "bytes {}-{}/{}",
694                                normalized.start, end, entry.metadata.total_size
695                            );
696                            response.headers_mut().insert(
697                                http::header::CONTENT_RANGE,
698                                http::header::HeaderValue::from_str(&content_range).unwrap(),
699                            );
700
701                            // Add Content-Length
702                            let content_length = (end - normalized.start + 1).to_string();
703                            response.headers_mut().insert(
704                                http::header::CONTENT_LENGTH,
705                                http::header::HeaderValue::from_str(&content_length).unwrap(),
706                            );
707
708                            return Ok(response);
709                        }
710                    }
711                }
712
713                #[cfg(feature = "metrics")]
714                counter!("tower_http_cache.chunk_cache_miss").increment(1);
715            }
716
717            let mut stale_entry: Option<CacheEntry> = None;
718            if let Some(ref key_ref) = key {
719                if let Ok(Some(hit)) = backend.get(key_ref).await {
720                    match classify_hit(hit, stale_window, refresh_before) {
721                        HitState::Fresh(entry) => {
722                            #[cfg(feature = "metrics")]
723                            counter!("tower_http_cache.hit").increment(1);
724
725                            // Record hit for auto-refresh tracking
726                            if let Some(ref manager) = refresh_manager {
727                                manager.tracker().record_hit(key_ref);
728                            }
729
730                            return Ok(entry.into_response());
731                        }
732                        HitState::Stale(entry) => {
733                            #[cfg(feature = "metrics")]
734                            counter!("tower_http_cache.stale_hit").increment(1);
735
736                            // Record hit for auto-refresh tracking
737                            if let Some(ref manager) = refresh_manager {
738                                manager.tracker().record_hit(key_ref);
739                            }
740
741                            stale_entry = Some(entry);
742                        }
743                        HitState::Expired => {}
744                    }
745                }
746            }
747
748            let mut primary_guard: Option<StampedeGuard> = None;
749            if let Some(ref key_ref) = key {
750                match StampedeGuard::acquire_handle(locks.clone(), key_ref.clone()).await {
751                    StampedeHandle::Primary(guard) => {
752                        primary_guard = Some(guard);
753                    }
754                    StampedeHandle::Secondary(lock) => {
755                        if let Some(entry) = stale_entry.clone() {
756                            #[cfg(feature = "metrics")]
757                            counter!("tower_http_cache.stale_served").increment(1);
758                            return Ok(entry.into_response());
759                        }
760
761                        let secondary_guard = lock.lock_owned().await;
762                        drop(secondary_guard);
763
764                        if let Ok(Some(hit)) = backend.get(key_ref).await {
765                            match classify_hit(hit, stale_window, refresh_before) {
766                                HitState::Fresh(entry) => {
767                                    #[cfg(feature = "metrics")]
768                                    counter!("tower_http_cache.hit_after_wait").increment(1);
769                                    return Ok(entry.into_response());
770                                }
771                                HitState::Stale(entry) => {
772                                    #[cfg(feature = "metrics")]
773                                    counter!("tower_http_cache.stale_served").increment(1);
774                                    return Ok(entry.into_response());
775                                }
776                                HitState::Expired => {}
777                            }
778                        }
779
780                        if let StampedeHandle::Primary(guard) =
781                            StampedeGuard::acquire_handle(locks.clone(), key_ref.clone()).await
782                        {
783                            primary_guard = Some(guard);
784                        }
785                    }
786                }
787            }
788
789            #[cfg(feature = "metrics")]
790            counter!("tower_http_cache.miss").increment(1);
791
792            #[cfg(feature = "metrics")]
793            let start = std::time::Instant::now();
794            let service = inner;
795            let response = service.oneshot(req).await.map_err(|err| err.into())?;
796            #[cfg(feature = "metrics")]
797            histogram!("tower_http_cache.backend_latency").record(start.elapsed().as_secs_f64());
798
799            let (parts, body) = response.into_parts();
800
801            // NEW: Early streaming decision
802            let content_type = parts
803                .headers
804                .get(http::header::CONTENT_TYPE)
805                .and_then(|v| v.to_str().ok());
806
807            let content_length = parts
808                .headers
809                .get(http::header::CONTENT_LENGTH)
810                .and_then(|v| v.to_str().ok())
811                .and_then(|v| v.parse::<u64>().ok());
812
813            let size_hint = body.size_hint();
814
815            let streaming_decision = should_stream(
816                policy.streaming_policy(),
817                &size_hint,
818                content_type,
819                content_length,
820            );
821
822            // Check for range requests
823            let is_range_request = parse_range_header(&parts.headers).is_some();
824            let is_partial_response = is_partial_content(parts.status);
825            let range_handling = policy.streaming_policy().range_handling;
826
827            // Handle range requests according to policy
828            if (is_range_request || is_partial_response)
829                && range_handling == RangeHandling::PassThrough
830            {
831                #[cfg(feature = "metrics")]
832                counter!("tower_http_cache.range_request_passthrough").increment(1);
833
834                #[cfg(feature = "tracing")]
835                tracing::debug!(
836                    method = %method,
837                    uri = %uri,
838                    is_range = is_range_request,
839                    is_partial = is_partial_response,
840                    "range_request_passthrough"
841                );
842
843                // Stream through without buffering for range requests
844                let boxed_body = SyncBoxBody::new(body.map_err(Into::into).boxed());
845                drop(primary_guard);
846                return Ok(Response::from_parts(parts, boxed_body));
847            }
848
849            // Check if we should skip caching and pass through
850            match streaming_decision {
851                StreamingDecision::SkipCache | StreamingDecision::StreamThrough => {
852                    // TRUE STREAMING: Pass through without buffering!
853                    #[cfg(feature = "metrics")]
854                    counter!("tower_http_cache.streaming_passthrough").increment(1);
855
856                    #[cfg(feature = "tracing")]
857                    tracing::debug!(
858                        method = %method,
859                        uri = %uri,
860                        decision = ?streaming_decision,
861                        content_type = ?content_type,
862                        size = ?extract_size_info(&size_hint, content_length),
863                        "streaming_passthrough"
864                    );
865
866                    // Box the body and stream it through without collecting
867                    let boxed_body = SyncBoxBody::new(body.map_err(Into::into).boxed());
868                    drop(primary_guard);
869                    return Ok(Response::from_parts(parts, boxed_body));
870                }
871                _ => {}
872            }
873
874            let streaming = body.size_hint().upper().is_none();
875            if streaming && !policy.allow_streaming_bodies() {
876                #[cfg(feature = "metrics")]
877                counter!("tower_http_cache.streaming_skip").increment(1);
878            }
879
880            let collected = BodyExt::collect(body).await.map_err(|err| err.into())?;
881            let cache_bytes = collected.to_bytes();
882            let response_bytes = cache_bytes.clone();
883
884            // Populate chunk cache for large files if enabled
885            if let (Some(ref chunk_cache), Some(ref key_ref)) = (&chunk_cache, &key) {
886                let streaming_policy = policy.streaming_policy();
887
888                if streaming_policy.enable_chunk_cache
889                    && cache_bytes.len() as u64 >= streaming_policy.min_chunk_file_size
890                    && parts.status.is_success()
891                {
892                    #[cfg(feature = "tracing")]
893                    tracing::debug!(
894                        key = %key_ref,
895                        size = cache_bytes.len(),
896                        chunk_size = streaming_policy.chunk_size,
897                        "populating_chunk_cache"
898                    );
899
900                    // Create chunk metadata
901                    let metadata = ChunkMetadata {
902                        total_size: cache_bytes.len() as u64,
903                        content_type: parts
904                            .headers
905                            .get(http::header::CONTENT_TYPE)
906                            .and_then(|v| v.to_str().ok())
907                            .unwrap_or("application/octet-stream")
908                            .to_string(),
909                        etag: parts
910                            .headers
911                            .get(http::header::ETAG)
912                            .and_then(|v| v.to_str().ok())
913                            .map(|s| s.to_string()),
914                        last_modified: parts
915                            .headers
916                            .get(http::header::LAST_MODIFIED)
917                            .and_then(|v| v.to_str().ok())
918                            .map(|s| s.to_string()),
919                        status: parts.status,
920                        version: parts.version,
921                        headers: policy.headers_to_cache(&parts.headers),
922                    };
923
924                    // Get or create chunked entry
925                    let entry = chunk_cache.get_or_create(key_ref.clone(), metadata);
926
927                    // Split into chunks and store
928                    let chunk_size = streaming_policy.chunk_size;
929                    let mut offset = 0;
930                    let mut chunk_index = 0;
931
932                    while offset < cache_bytes.len() {
933                        let end = std::cmp::min(offset + chunk_size, cache_bytes.len());
934                        let chunk = cache_bytes.slice(offset..end);
935                        entry.add_chunk(chunk_index, chunk);
936
937                        offset = end;
938                        chunk_index += 1;
939                    }
940
941                    #[cfg(feature = "metrics")]
942                    counter!("tower_http_cache.chunk_cache_stored").increment(1);
943
944                    #[cfg(feature = "tracing")]
945                    tracing::debug!(
946                        key = %key_ref,
947                        chunks = chunk_index,
948                        "chunk_cache_populated"
949                    );
950                }
951            }
952
953            let cache_control_block =
954                policy.respect_cache_control() && cache_control_disallows(&parts.headers);
955            let body_too_large = policy
956                .max_body_size()
957                .is_some_and(|max| cache_bytes.len() > max);
958            let body_too_small = policy
959                .min_body_size()
960                .is_some_and(|min| cache_bytes.len() < min);
961
962            let should_store = key.is_some()
963                && !cache_control_block
964                && !body_too_large
965                && !body_too_small
966                && (policy.allow_streaming_bodies() || !streaming);
967
968            let headers_to_cache = if should_store {
969                Some(policy.headers_to_cache(&parts.headers))
970            } else {
971                None
972            };
973
974            let version = parts.version;
975            let status = parts.status;
976
977            if should_store {
978                if let Some(key_ref) = &key {
979                    if let Some(ttl) = policy.ttl_for(status) {
980                        if !ttl.is_zero() {
981                            let stale_for = policy.stale_while_revalidate();
982                            let (compressed_bytes, compressed) =
983                                maybe_compress(cache_bytes.clone(), policy.compression());
984                            if compressed {
985                                #[cfg(feature = "metrics")]
986                                counter!("tower_http_cache.compressed").increment(1);
987                            }
988                            let entry = CacheEntry::new(
989                                status,
990                                version,
991                                headers_to_cache.unwrap(),
992                                compressed_bytes,
993                            );
994                            if backend
995                                .set(key_ref.clone(), entry, ttl, stale_for)
996                                .await
997                                .is_err()
998                            {
999                                #[cfg(feature = "metrics")]
1000                                counter!("tower_http_cache.store_error").increment(1);
1001                            } else {
1002                                #[cfg(feature = "metrics")]
1003                                counter!("tower_http_cache.store").increment(1);
1004
1005                                // Store refresh metadata if auto-refresh is enabled
1006                                if let (Some(ref manager), Some(metadata)) =
1007                                    (&refresh_manager, refresh_metadata)
1008                                {
1009                                    manager.store_metadata(key_ref.clone(), metadata);
1010                                }
1011                            }
1012                        }
1013                    }
1014                }
1015            } else {
1016                #[cfg(feature = "metrics")]
1017                counter!("tower_http_cache.store_skipped").increment(1);
1018            }
1019
1020            drop(primary_guard);
1021
1022            // Box the response body
1023            let full_body = Full::from(response_bytes);
1024            let boxed_body = SyncBoxBody::new(full_body.map_err(Into::into).boxed());
1025            Ok(Response::from_parts(parts, boxed_body))
1026        })
1027    }
1028}
1029
1030struct StampedeGuard {
1031    key: String,
1032    locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
1033    lock: Arc<Mutex<()>>,
1034    _guard: OwnedMutexGuard<()>,
1035}
1036
1037enum StampedeHandle {
1038    Primary(StampedeGuard),
1039    Secondary(Arc<Mutex<()>>),
1040}
1041
1042impl StampedeGuard {
1043    async fn acquire_handle(
1044        locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
1045        key: String,
1046    ) -> StampedeHandle {
1047        let handle = match locks.entry(key.clone()) {
1048            Entry::Occupied(entry) => StampedeHandle::Secondary(entry.get().clone()),
1049            Entry::Vacant(entry) => {
1050                let lock = Arc::new(Mutex::new(()));
1051                entry.insert(lock.clone());
1052                let guard = lock.clone().lock_owned().await;
1053                let locks_clone = locks.clone();
1054                StampedeHandle::Primary(StampedeGuard {
1055                    key,
1056                    locks: locks_clone,
1057                    lock,
1058                    _guard: guard,
1059                })
1060            }
1061        };
1062        handle
1063    }
1064}
1065
1066impl Drop for StampedeGuard {
1067    fn drop(&mut self) {
1068        if let Some(current) = self.locks.get(&self.key) {
1069            let should_remove = Arc::ptr_eq(&self.lock, current.value());
1070            drop(current);
1071            if should_remove {
1072                self.locks.remove(&self.key);
1073            }
1074        }
1075    }
1076}
1077
1078fn classify_hit(hit: CacheRead, stale_window: Duration, refresh_before: Duration) -> HitState {
1079    let now = SystemTime::now();
1080    let CacheRead {
1081        entry,
1082        expires_at,
1083        stale_until,
1084    } = hit;
1085
1086    if let Some(expires_at) = expires_at {
1087        if expires_at > now {
1088            if refresh_before > Duration::ZERO {
1089                if let Some(threshold) = expires_at.checked_sub(refresh_before) {
1090                    if now >= threshold {
1091                        return HitState::Stale(entry);
1092                    }
1093                } else {
1094                    return HitState::Stale(entry);
1095                }
1096            }
1097            return HitState::Fresh(entry);
1098        }
1099    }
1100
1101    if stale_window > Duration::ZERO {
1102        if let Some(stale_until) = stale_until {
1103            if stale_until > now {
1104                return HitState::Stale(entry);
1105            }
1106        }
1107    }
1108
1109    HitState::Expired
1110}
1111
1112#[derive(Debug)]
1113enum HitState {
1114    Fresh(CacheEntry),
1115    Stale(CacheEntry),
1116    Expired,
1117}
1118
1119fn cache_control_disallows(headers: &HeaderMap) -> bool {
1120    headers
1121        .get_all(CACHE_CONTROL)
1122        .iter()
1123        .filter_map(|value| value.to_str().ok())
1124        .flat_map(|value| value.split(','))
1125        .map(|token| token.trim().to_ascii_lowercase())
1126        .any(|token| matches!(token.as_str(), "no-store" | "no-cache" | "private"))
1127        || headers
1128            .get(PRAGMA)
1129            .and_then(|value| value.to_str().ok())
1130            .map(|value| value.to_ascii_lowercase().contains("no-cache"))
1131            .unwrap_or(false)
1132}
1133
1134#[cfg(feature = "compression")]
1135fn maybe_compress(bytes: Bytes, config: CompressionConfig) -> (Bytes, bool) {
1136    use flate2::{write::GzEncoder, Compression};
1137    use std::io::Write;
1138
1139    match config.strategy {
1140        CompressionStrategy::None => (bytes, false),
1141        CompressionStrategy::Gzip => {
1142            if bytes.len() < config.min_size {
1143                return (bytes, false);
1144            }
1145            let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
1146            if encoder.write_all(&bytes).is_err() {
1147                return (bytes, false);
1148            }
1149            match encoder.finish() {
1150                Ok(data) => (Bytes::from(data), true),
1151                Err(_) => (bytes, false),
1152            }
1153        }
1154    }
1155}
1156
1157#[cfg(not(feature = "compression"))]
1158fn maybe_compress(bytes: Bytes, _config: CompressionConfig) -> (Bytes, bool) {
1159    let _ = _config;
1160    (bytes, false)
1161}
1162
1163#[cfg(all(test, feature = "in-memory"))]
1164mod tests {
1165    use super::*;
1166    use crate::backend::CacheEntry;
1167    use bytes::Bytes;
1168    use http::{HeaderValue, StatusCode, Version};
1169    use tokio::task::yield_now;
1170
1171    fn mock_entry() -> CacheEntry {
1172        CacheEntry::new(
1173            StatusCode::OK,
1174            Version::HTTP_11,
1175            Vec::new(),
1176            Bytes::from_static(b"body"),
1177        )
1178    }
1179
1180    #[test]
1181    fn classify_hit_marks_entry_fresh_when_not_near_expiry() {
1182        let now = SystemTime::now();
1183        let hit = CacheRead {
1184            entry: mock_entry(),
1185            expires_at: Some(now + Duration::from_secs(10)),
1186            stale_until: Some(now + Duration::from_secs(20)),
1187        };
1188
1189        match classify_hit(hit, Duration::from_secs(5), Duration::from_secs(1)) {
1190            HitState::Fresh(_) => {}
1191            other => panic!("expected fresh entry, got {:?}", other),
1192        }
1193    }
1194
1195    #[test]
1196    fn classify_hit_marks_entry_stale_when_within_refresh_window() {
1197        let now = SystemTime::now();
1198        let hit = CacheRead {
1199            entry: mock_entry(),
1200            expires_at: Some(now + Duration::from_secs(2)),
1201            stale_until: Some(now + Duration::from_secs(10)),
1202        };
1203
1204        match classify_hit(hit, Duration::from_secs(5), Duration::from_secs(5)) {
1205            HitState::Stale(_) => {}
1206            other => panic!("expected stale entry, got {:?}", other),
1207        }
1208    }
1209
1210    #[test]
1211    fn classify_hit_marks_entry_stale_when_within_stale_window() {
1212        let now = SystemTime::now();
1213        let hit = CacheRead {
1214            entry: mock_entry(),
1215            expires_at: Some(now - Duration::from_secs(1)),
1216            stale_until: Some(now + Duration::from_secs(1)),
1217        };
1218
1219        match classify_hit(hit, Duration::from_secs(2), Duration::from_secs(0)) {
1220            HitState::Stale(_) => {}
1221            other => panic!("expected stale entry, got {:?}", other),
1222        }
1223    }
1224
1225    #[test]
1226    fn classify_hit_marks_entry_expired_after_stale_window() {
1227        let now = SystemTime::now();
1228        let hit = CacheRead {
1229            entry: mock_entry(),
1230            expires_at: Some(now - Duration::from_secs(5)),
1231            stale_until: Some(now - Duration::from_secs(1)),
1232        };
1233
1234        match classify_hit(hit, Duration::from_secs(5), Duration::from_secs(0)) {
1235            HitState::Expired => {}
1236            other => panic!("expected expired entry, got {:?}", other),
1237        }
1238    }
1239
1240    #[test]
1241    fn cache_control_disallows_detects_no_cache_directives() {
1242        let mut headers = HeaderMap::new();
1243        headers.insert(
1244            CACHE_CONTROL,
1245            HeaderValue::from_static("max-age=0, no-cache"),
1246        );
1247        assert!(cache_control_disallows(&headers));
1248
1249        let mut pragma_only = HeaderMap::new();
1250        pragma_only.insert(PRAGMA, HeaderValue::from_static("no-cache"));
1251        assert!(cache_control_disallows(&pragma_only));
1252    }
1253
1254    #[tokio::test]
1255    async fn stampede_guard_drop_removes_lock_entry() {
1256        let locks = Arc::new(DashMap::new());
1257        let key = "key".to_string();
1258
1259        match StampedeGuard::acquire_handle(locks.clone(), key.clone()).await {
1260            StampedeHandle::Primary(guard) => {
1261                assert!(locks.get(&key).is_some());
1262                drop(guard);
1263                yield_now().await;
1264                assert!(locks.get(&key).is_none());
1265            }
1266            StampedeHandle::Secondary(_) => panic!("expected primary guard"),
1267        }
1268    }
1269
1270    #[test]
1271    fn cache_service_implements_clone() {
1272        use crate::backend::memory::InMemoryBackend;
1273        use tower::service_fn;
1274
1275        // Compile-time check that CacheService implements Clone
1276        fn assert_clone<T: Clone>(_: &T) {}
1277
1278        let backend = InMemoryBackend::new(100);
1279        let layer = CacheLayer::new(backend);
1280        let service = layer.layer(service_fn(|_req: http::Request<()>| async {
1281            Ok::<_, std::convert::Infallible>(http::Response::new(()))
1282        }));
1283
1284        // This will fail to compile if CacheService doesn't implement Clone
1285        assert_clone(&service);
1286    }
1287}
tower_http_cache/layer.rs

tower_http_cache/
layer.rs
