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                // Double-checked read. The lookup above and this acquisition
808                // are not one atomic step: another caller can take the lock,
809                // fetch, store and release inside that window, so winning the
810                // lock does not mean the key still needs fetching. Without
811                // this, two requests in flight at the same time each hit the
812                // origin -- the stampede the lock exists to prevent. The
813                // waiting branch above re-reads for the same reason. Only a
814                // fresh entry short-circuits; stale or expired still owes the
815                // origin a request.
816                if primary_guard.is_some() {
817                    if let Ok(Some(hit)) = backend.get(key_ref).await {
818                        if let HitState::Fresh(entry) =
819                            classify_hit(hit, stale_window, refresh_before)
820                        {
821                            #[cfg(feature = "metrics")]
822                            counter!("tower_http_cache.hit_after_lock").increment(1);
823                            return Ok(entry.into_response());
824                        }
825                    }
826                }
827            }
828
829            #[cfg(feature = "metrics")]
830            counter!("tower_http_cache.miss").increment(1);
831
832            #[cfg(feature = "metrics")]
833            let start = std::time::Instant::now();
834            let service = inner;
835            let response = service.oneshot(req).await.map_err(|err| err.into())?;
836            #[cfg(feature = "metrics")]
837            histogram!("tower_http_cache.backend_latency").record(start.elapsed().as_secs_f64());
838
839            let (parts, body) = response.into_parts();
840
841            // NEW: Early streaming decision
842            let content_type = parts
843                .headers
844                .get(http::header::CONTENT_TYPE)
845                .and_then(|v| v.to_str().ok());
846
847            let content_length = parts
848                .headers
849                .get(http::header::CONTENT_LENGTH)
850                .and_then(|v| v.to_str().ok())
851                .and_then(|v| v.parse::<u64>().ok());
852
853            let size_hint = body.size_hint();
854
855            let streaming_decision = should_stream(
856                policy.streaming_policy(),
857                &size_hint,
858                content_type,
859                content_length,
860            );
861
862            // Check for range requests
863            let is_range_request = parse_range_header(&parts.headers).is_some();
864            let is_partial_response = is_partial_content(parts.status);
865            let range_handling = policy.streaming_policy().range_handling;
866
867            // Handle range requests according to policy
868            if (is_range_request || is_partial_response)
869                && range_handling == RangeHandling::PassThrough
870            {
871                #[cfg(feature = "metrics")]
872                counter!("tower_http_cache.range_request_passthrough").increment(1);
873
874                #[cfg(feature = "tracing")]
875                tracing::debug!(
876                    method = %method,
877                    uri = %uri,
878                    is_range = is_range_request,
879                    is_partial = is_partial_response,
880                    "range_request_passthrough"
881                );
882
883                // Stream through without buffering for range requests
884                let boxed_body = SyncBoxBody::new(body.map_err(Into::into).boxed());
885                drop(primary_guard);
886                return Ok(Response::from_parts(parts, boxed_body));
887            }
888
889            // Check if we should skip caching and pass through
890            match streaming_decision {
891                StreamingDecision::SkipCache | StreamingDecision::StreamThrough => {
892                    // TRUE STREAMING: Pass through without buffering!
893                    #[cfg(feature = "metrics")]
894                    counter!("tower_http_cache.streaming_passthrough").increment(1);
895
896                    #[cfg(feature = "tracing")]
897                    tracing::debug!(
898                        method = %method,
899                        uri = %uri,
900                        decision = ?streaming_decision,
901                        content_type = ?content_type,
902                        size = ?extract_size_info(&size_hint, content_length),
903                        "streaming_passthrough"
904                    );
905
906                    // Box the body and stream it through without collecting
907                    let boxed_body = SyncBoxBody::new(body.map_err(Into::into).boxed());
908                    drop(primary_guard);
909                    return Ok(Response::from_parts(parts, boxed_body));
910                }
911                _ => {}
912            }
913
914            let streaming = body.size_hint().upper().is_none();
915            if streaming && !policy.allow_streaming_bodies() {
916                #[cfg(feature = "metrics")]
917                counter!("tower_http_cache.streaming_skip").increment(1);
918            }
919
920            let collected = BodyExt::collect(body).await.map_err(|err| err.into())?;
921            let cache_bytes = collected.to_bytes();
922            let response_bytes = cache_bytes.clone();
923
924            // Populate chunk cache for large files if enabled
925            if let (Some(chunk_cache), Some(key_ref)) = (&chunk_cache, &key) {
926                let streaming_policy = policy.streaming_policy();
927
928                if streaming_policy.enable_chunk_cache
929                    && cache_bytes.len() as u64 >= streaming_policy.min_chunk_file_size
930                    && parts.status.is_success()
931                {
932                    #[cfg(feature = "tracing")]
933                    tracing::debug!(
934                        key = %key_ref,
935                        size = cache_bytes.len(),
936                        chunk_size = streaming_policy.chunk_size,
937                        "populating_chunk_cache"
938                    );
939
940                    // Create chunk metadata
941                    let metadata = ChunkMetadata {
942                        total_size: cache_bytes.len() as u64,
943                        content_type: parts
944                            .headers
945                            .get(http::header::CONTENT_TYPE)
946                            .and_then(|v| v.to_str().ok())
947                            .unwrap_or("application/octet-stream")
948                            .to_string(),
949                        etag: parts
950                            .headers
951                            .get(http::header::ETAG)
952                            .and_then(|v| v.to_str().ok())
953                            .map(|s| s.to_string()),
954                        last_modified: parts
955                            .headers
956                            .get(http::header::LAST_MODIFIED)
957                            .and_then(|v| v.to_str().ok())
958                            .map(|s| s.to_string()),
959                        status: parts.status,
960                        version: parts.version,
961                        headers: policy.headers_to_cache(&parts.headers),
962                    };
963
964                    // Get or create chunked entry
965                    let entry = chunk_cache.get_or_create(key_ref.clone(), metadata);
966
967                    // Split into chunks and store
968                    let chunk_size = streaming_policy.chunk_size;
969                    let mut offset = 0;
970                    let mut chunk_index = 0;
971
972                    while offset < cache_bytes.len() {
973                        let end = std::cmp::min(offset + chunk_size, cache_bytes.len());
974                        let chunk = cache_bytes.slice(offset..end);
975                        entry.add_chunk(chunk_index, chunk);
976
977                        offset = end;
978                        chunk_index += 1;
979                    }
980
981                    #[cfg(feature = "metrics")]
982                    counter!("tower_http_cache.chunk_cache_stored").increment(1);
983
984                    #[cfg(feature = "tracing")]
985                    tracing::debug!(
986                        key = %key_ref,
987                        chunks = chunk_index,
988                        "chunk_cache_populated"
989                    );
990                }
991            }
992
993            let cache_control_block =
994                policy.respect_cache_control() && cache_control_disallows(&parts.headers);
995            let body_too_large = policy
996                .max_body_size()
997                .is_some_and(|max| cache_bytes.len() > max);
998            let body_too_small = policy
999                .min_body_size()
1000                .is_some_and(|min| cache_bytes.len() < min);
1001
1002            let should_store = key.is_some()
1003                && !cache_control_block
1004                && !body_too_large
1005                && !body_too_small
1006                && (policy.allow_streaming_bodies() || !streaming);
1007
1008            let headers_to_cache = if should_store {
1009                Some(policy.headers_to_cache(&parts.headers))
1010            } else {
1011                None
1012            };
1013
1014            let version = parts.version;
1015            let status = parts.status;
1016
1017            if should_store {
1018                if let Some(key_ref) = &key {
1019                    if let Some(ttl) = policy.ttl_for(status) {
1020                        if !ttl.is_zero() {
1021                            let stale_for = policy.stale_while_revalidate();
1022                            let (compressed_bytes, compressed) =
1023                                maybe_compress(cache_bytes.clone(), policy.compression());
1024                            if compressed {
1025                                #[cfg(feature = "metrics")]
1026                                counter!("tower_http_cache.compressed").increment(1);
1027                            }
1028                            let entry = CacheEntry::new(
1029                                status,
1030                                version,
1031                                headers_to_cache.unwrap(),
1032                                compressed_bytes,
1033                            );
1034                            // `extract_tags` returns an empty vector unless
1035                            // `TagPolicy::enabled` is set, which defaults to
1036                            // false -- so this is inert for anyone who has not
1037                            // opted in.
1038                            let tags = policy.extract_tags(&method, &uri);
1039                            let entry = if tags.is_empty() {
1040                                entry
1041                            } else {
1042                                entry.with_tags(tags)
1043                            };
1044                            if backend
1045                                .set(key_ref.clone(), entry, ttl, stale_for)
1046                                .await
1047                                .is_err()
1048                            {
1049                                #[cfg(feature = "metrics")]
1050                                counter!("tower_http_cache.store_error").increment(1);
1051                            } else {
1052                                #[cfg(feature = "metrics")]
1053                                counter!("tower_http_cache.store").increment(1);
1054
1055                                // Store refresh metadata if auto-refresh is enabled
1056                                if let (Some(manager), Some(metadata)) =
1057                                    (&refresh_manager, refresh_metadata)
1058                                {
1059                                    manager.store_metadata(key_ref.clone(), metadata);
1060                                }
1061                            }
1062                        }
1063                    }
1064                }
1065            } else {
1066                #[cfg(feature = "metrics")]
1067                counter!("tower_http_cache.store_skipped").increment(1);
1068            }
1069
1070            drop(primary_guard);
1071
1072            // Box the response body
1073            let full_body = Full::from(response_bytes);
1074            let boxed_body = SyncBoxBody::new(full_body.map_err(Into::into).boxed());
1075            Ok(Response::from_parts(parts, boxed_body))
1076        })
1077    }
1078}
1079
1080struct StampedeGuard {
1081    key: String,
1082    locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
1083    lock: Arc<Mutex<()>>,
1084    _guard: OwnedMutexGuard<()>,
1085}
1086
1087enum StampedeHandle {
1088    Primary(StampedeGuard),
1089    Secondary(Arc<Mutex<()>>),
1090}
1091
1092impl StampedeGuard {
1093    async fn acquire_handle(
1094        locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
1095        key: String,
1096    ) -> StampedeHandle {
1097        // The `let` binding is deliberate: `DashMap::entry` returns a guard
1098        // holding the shard's write lock, and binding the match result pins
1099        // the drop point to the end of this statement -- identically in both
1100        // edition 2021 and 2024. Edition 2024 changed when tail-expression
1101        // temporaries drop, which is why clippy only started flagging this as
1102        // redundant after the edition bump. Returning the `match` directly
1103        // would make lock release depend on those edition rules; in a path
1104        // that acquires a stampede lock and then awaits, that is not a
1105        // trade worth making to save one line.
1106        #[allow(clippy::let_and_return)]
1107        let handle = match locks.entry(key.clone()) {
1108            Entry::Occupied(entry) => StampedeHandle::Secondary(entry.get().clone()),
1109            Entry::Vacant(entry) => {
1110                let lock = Arc::new(Mutex::new(()));
1111                entry.insert(lock.clone());
1112                let guard = lock.clone().lock_owned().await;
1113                let locks_clone = locks.clone();
1114                StampedeHandle::Primary(StampedeGuard {
1115                    key,
1116                    locks: locks_clone,
1117                    lock,
1118                    _guard: guard,
1119                })
1120            }
1121        };
1122        handle
1123    }
1124}
1125
1126impl Drop for StampedeGuard {
1127    fn drop(&mut self) {
1128        if let Some(current) = self.locks.get(&self.key) {
1129            let should_remove = Arc::ptr_eq(&self.lock, current.value());
1130            drop(current);
1131            if should_remove {
1132                self.locks.remove(&self.key);
1133            }
1134        }
1135    }
1136}
1137
1138fn classify_hit(hit: CacheRead, stale_window: Duration, refresh_before: Duration) -> HitState {
1139    let now = SystemTime::now();
1140    let CacheRead {
1141        entry,
1142        expires_at,
1143        stale_until,
1144    } = hit;
1145
1146    if let Some(expires_at) = expires_at {
1147        if expires_at > now {
1148            if refresh_before > Duration::ZERO {
1149                if let Some(threshold) = expires_at.checked_sub(refresh_before) {
1150                    if now >= threshold {
1151                        return HitState::Stale(entry);
1152                    }
1153                } else {
1154                    return HitState::Stale(entry);
1155                }
1156            }
1157            return HitState::Fresh(entry);
1158        }
1159    }
1160
1161    if stale_window > Duration::ZERO {
1162        if let Some(stale_until) = stale_until {
1163            if stale_until > now {
1164                return HitState::Stale(entry);
1165            }
1166        }
1167    }
1168
1169    HitState::Expired
1170}
1171
1172#[derive(Debug)]
1173enum HitState {
1174    Fresh(CacheEntry),
1175    Stale(CacheEntry),
1176    Expired,
1177}
1178
1179fn cache_control_disallows(headers: &HeaderMap) -> bool {
1180    headers
1181        .get_all(CACHE_CONTROL)
1182        .iter()
1183        .filter_map(|value| value.to_str().ok())
1184        .flat_map(|value| value.split(','))
1185        .map(|token| token.trim().to_ascii_lowercase())
1186        .any(|token| matches!(token.as_str(), "no-store" | "no-cache" | "private"))
1187        || headers
1188            .get(PRAGMA)
1189            .and_then(|value| value.to_str().ok())
1190            .map(|value| value.to_ascii_lowercase().contains("no-cache"))
1191            .unwrap_or(false)
1192}
1193
1194#[cfg(feature = "compression")]
1195fn maybe_compress(bytes: Bytes, config: CompressionConfig) -> (Bytes, bool) {
1196    use flate2::{Compression, write::GzEncoder};
1197    use std::io::Write;
1198
1199    match config.strategy {
1200        CompressionStrategy::None => (bytes, false),
1201        CompressionStrategy::Gzip => {
1202            if bytes.len() < config.min_size {
1203                return (bytes, false);
1204            }
1205            let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
1206            if encoder.write_all(&bytes).is_err() {
1207                return (bytes, false);
1208            }
1209            match encoder.finish() {
1210                Ok(data) => (Bytes::from(data), true),
1211                Err(_) => (bytes, false),
1212            }
1213        }
1214    }
1215}
1216
1217#[cfg(not(feature = "compression"))]
1218fn maybe_compress(bytes: Bytes, _config: CompressionConfig) -> (Bytes, bool) {
1219    let _ = _config;
1220    (bytes, false)
1221}
1222
1223#[cfg(all(test, feature = "in-memory"))]
1224mod tests {
1225    use super::*;
1226    use crate::backend::CacheEntry;
1227    use bytes::Bytes;
1228    use http::{HeaderValue, StatusCode, Version};
1229    use tokio::task::yield_now;
1230
1231    fn mock_entry() -> CacheEntry {
1232        CacheEntry::new(
1233            StatusCode::OK,
1234            Version::HTTP_11,
1235            Vec::new(),
1236            Bytes::from_static(b"body"),
1237        )
1238    }
1239
1240    #[test]
1241    fn classify_hit_marks_entry_fresh_when_not_near_expiry() {
1242        let now = SystemTime::now();
1243        let hit = CacheRead {
1244            entry: mock_entry(),
1245            expires_at: Some(now + Duration::from_secs(10)),
1246            stale_until: Some(now + Duration::from_secs(20)),
1247        };
1248
1249        match classify_hit(hit, Duration::from_secs(5), Duration::from_secs(1)) {
1250            HitState::Fresh(_) => {}
1251            other => panic!("expected fresh entry, got {:?}", other),
1252        }
1253    }
1254
1255    #[test]
1256    fn classify_hit_marks_entry_stale_when_within_refresh_window() {
1257        let now = SystemTime::now();
1258        let hit = CacheRead {
1259            entry: mock_entry(),
1260            expires_at: Some(now + Duration::from_secs(2)),
1261            stale_until: Some(now + Duration::from_secs(10)),
1262        };
1263
1264        match classify_hit(hit, Duration::from_secs(5), Duration::from_secs(5)) {
1265            HitState::Stale(_) => {}
1266            other => panic!("expected stale entry, got {:?}", other),
1267        }
1268    }
1269
1270    #[test]
1271    fn classify_hit_marks_entry_stale_when_within_stale_window() {
1272        let now = SystemTime::now();
1273        let hit = CacheRead {
1274            entry: mock_entry(),
1275            expires_at: Some(now - Duration::from_secs(1)),
1276            stale_until: Some(now + Duration::from_secs(1)),
1277        };
1278
1279        match classify_hit(hit, Duration::from_secs(2), Duration::from_secs(0)) {
1280            HitState::Stale(_) => {}
1281            other => panic!("expected stale entry, got {:?}", other),
1282        }
1283    }
1284
1285    #[test]
1286    fn classify_hit_marks_entry_expired_after_stale_window() {
1287        let now = SystemTime::now();
1288        let hit = CacheRead {
1289            entry: mock_entry(),
1290            expires_at: Some(now - Duration::from_secs(5)),
1291            stale_until: Some(now - Duration::from_secs(1)),
1292        };
1293
1294        match classify_hit(hit, Duration::from_secs(5), Duration::from_secs(0)) {
1295            HitState::Expired => {}
1296            other => panic!("expected expired entry, got {:?}", other),
1297        }
1298    }
1299
1300    #[test]
1301    fn cache_control_disallows_detects_no_cache_directives() {
1302        let mut headers = HeaderMap::new();
1303        headers.insert(
1304            CACHE_CONTROL,
1305            HeaderValue::from_static("max-age=0, no-cache"),
1306        );
1307        assert!(cache_control_disallows(&headers));
1308
1309        let mut pragma_only = HeaderMap::new();
1310        pragma_only.insert(PRAGMA, HeaderValue::from_static("no-cache"));
1311        assert!(cache_control_disallows(&pragma_only));
1312    }
1313
1314    #[tokio::test]
1315    async fn stampede_guard_drop_removes_lock_entry() {
1316        let locks = Arc::new(DashMap::new());
1317        let key = "key".to_string();
1318
1319        match StampedeGuard::acquire_handle(locks.clone(), key.clone()).await {
1320            StampedeHandle::Primary(guard) => {
1321                assert!(locks.get(&key).is_some());
1322                drop(guard);
1323                yield_now().await;
1324                assert!(locks.get(&key).is_none());
1325            }
1326            StampedeHandle::Secondary(_) => panic!("expected primary guard"),
1327        }
1328    }
1329
1330    #[test]
1331    fn cache_service_implements_clone() {
1332        use crate::backend::memory::InMemoryBackend;
1333        use tower::service_fn;
1334
1335        // Compile-time check that CacheService implements Clone
1336        fn assert_clone<T: Clone>(_: &T) {}
1337
1338        let backend = InMemoryBackend::new(100);
1339        let layer = CacheLayer::new(backend);
1340        let service = layer.layer(service_fn(|_req: http::Request<()>| async {
1341            Ok::<_, std::convert::Infallible>(http::Response::new(()))
1342        }));
1343
1344        // This will fail to compile if CacheService doesn't implement Clone
1345        assert_clone(&service);
1346    }
1347}
tower_http_cache/layer.rs

tower_http_cache/
layer.rs
