tower_http_cache/
layer.rs

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

tower_http_cache/
layer.rs
