RustProxy/src/cache/lookup/hot.rs

impl Cache {
    fn get_hot(&self, ticket: &Ticket, count_miss: bool) -> Option<Hit> {
        let mut store = self.store.lock().unwrap_or_else(|e| e.into_inner());
        if store.revision != ticket.revision
            || store.generation != ticket.generation
            || store
                .scope_generations
                .get(&ticket.scope)
                .copied()
                .unwrap_or(0)
                != ticket.scope_generation
            || store.key_generations.get(&ticket.key).copied().unwrap_or(0) != ticket.key_generation
        {
            return None;
        }
        if let Some(storage_key) = matching_hot_key(&store, ticket) {
            let banned = store
                .entries
                .get(&storage_key)
                .is_some_and(|entry| self.hot_entry_banned(&storage_key, entry));
            if banned {
                store.remove(&storage_key);
                if count_miss {
                    self.counters.misses.fetch_add(1, Ordering::Relaxed);
                }
                drop(store);
                release_allocator_memory();
                return None;
            }
            let entry = store.entries.get(&storage_key)?;
            let resident_age_seconds = hot_resident_age(entry).as_secs();
            let age = entry
                .initial_age
                .saturating_add(entry.received.elapsed())
                .as_secs();
            let state = cache_state(
                age,
                entry.ttl,
                entry.stale_while_revalidate,
                entry.must_revalidate,
                entry.proxy_revalidate,
                entry.requires_revalidation,
                &ticket.request_headers,
            );
            if state.servable() {
                let sequence = store.next_sequence();
                let (
                    old_sequence,
                    headers,
                    body,
                    hits,
                    transformed,
                    skip_compression,
                    remaining_ttl,
                    body_bytes,
                ) = {
                    let entry = store.entries.get_mut(&storage_key)?;
                    let old_sequence = entry.sequence;
                    entry.hits = entry.hits.saturating_add(1);
                    entry.last_access = Instant::now();
                    entry.sequence = sequence;
                    (
                        old_sequence,
                        Arc::clone(&entry.headers),
                        entry.body.clone(),
                        entry.hits,
                        entry.transformed,
                        entry.skip_compression,
                        entry.ttl.saturating_sub(age),
                        entry.body.len(),
                    )
                };
                store.lru.remove(&old_sequence);
                store.lru.insert(sequence, storage_key);
                self.counters.hits.fetch_add(1, Ordering::Relaxed);
                drop(store);

                let mut response = Response::new(Body::from(body.clone()));
                *response.headers_mut() = headers.as_ref().clone();
                if transformed {
                    response
                        .extensions_mut()
                        .insert(crate::transform::CachedTransformed);
                }
                if skip_compression {
                    response.extensions_mut().insert(crate::compression::Skip);
                }
                response.headers_mut().insert(
                    "age",
                    HeaderValue::from_str(&age.to_string()).expect("integer header"),
                );
                self.domain_stats
                    .saved_body_hit(&ticket.scope, &ticket.domain, body_bytes);
                return Some(Hit {
                    response,
                    body,
                    hits,
                    transformed,
                    remaining_ttl,
                    resident_age_seconds,
                    stale: state.stale(),
                    background_revalidate: state.background_revalidate(),
                });
            }
        }
        if count_miss {
            self.counters.misses.fetch_add(1, Ordering::Relaxed);
        }
        None
    }
    pub fn get(&self, ticket: &Ticket) -> Option<Hit> {
        self.get_hot(ticket, true)
    }
    async fn get_cold(&self, ticket: &Ticket, settings: &Settings) -> Option<Hit> {
        if !self.ticket_current(ticket) {
            return None;
        }
        let banned = {
            let mut cold = self.cold.lock().unwrap_or_else(|e| e.into_inner());
            matching_cold_key(&cold, ticket).and_then(|storage_key| {
                let banned = cold
                    .entries
                    .get(&storage_key)
                    .is_some_and(|entry| self.cold_entry_banned(&storage_key, entry));
                banned.then(|| cold.remove(&storage_key)).flatten()
            })
        };
        if let Some(entry) = banned {
            self.delete_cold_entry(entry);
            return None;
        }
        let (storage_key, entry) = {
            let mut cold = self.cold.lock().unwrap_or_else(|e| e.into_inner());
            let storage_key = matching_cold_key(&cold, ticket)?;
            let old_sequence = {
                let entry = cold.entries.get(&storage_key)?;
                if entry.revision != ticket.revision || entry.scope != ticket.scope {
                    return None;
                }
                let state = cache_state(
                    cold_age(entry),
                    entry.ttl,
                    entry.stale_while_revalidate,
                    entry.must_revalidate,
                    entry.proxy_revalidate,
                    entry.requires_revalidation,
                    &ticket.request_headers,
                );
                if !state.servable() {
                    return None;
                }
                entry.sequence
            };
            let sequence = cold.next_sequence();
            cold.lru.remove(&old_sequence);
            cold.lru.insert(sequence, storage_key);
            let entry = cold.entries.get_mut(&storage_key)?;
            entry.sequence = sequence;
            entry.last_access = SystemTime::now();
            entry.pinned_reads = entry.pinned_reads.saturating_add(1);
            (storage_key, entry.clone())
        };
        let _pin = ColdReadGuard {
            cache: self,
            key: storage_key,
            generation: entry.generation,
        };
        let Some(metadata) = read_cold_metadata(&storage_key, &entry).await else {
            self.counters.cold_errors.fetch_add(1, Ordering::Relaxed);
            if let Some(entry) = self.take_cold(&storage_key) {
                self.delete_cold_entry(entry);
            }
            return None;
        };
        if entry.body_bytes > settings.max_object_bytes {
            self.counters
                .memory_pressure_rejections
                .fetch_add(1, Ordering::Relaxed);
            return None;
        }
        let Some(_read_lease) = self.try_reserve_transient(entry.body_bytes, settings) else {
            return None;
        };
        let _cold_read = ByteGaugeGuard::new(&self.cold_read_bytes_active, entry.body_bytes as u64);
        let body_vec = match tokio::fs::read(&entry.body_path).await {
            Ok(body) if cold_body_matches(&entry, &body) => body,
            _ => {
                self.counters.cold_errors.fetch_add(1, Ordering::Relaxed);
                if let Some(entry) = self.take_cold(&storage_key) {
                    self.delete_cold_entry(entry);
                }
                return None;
            }
        };
        let body_allocated = body_vec.capacity().max(body_vec.len());
        let extra = body_allocated.saturating_sub(entry.body_bytes);
        let extra_lease = if extra > 0 {
            self.try_reserve_transient(extra, settings)
        } else {
            None
        };
        if extra > 0 && extra_lease.is_none() {
            return None;
        }
        let _cold_read_extra = extra_lease
            .as_ref()
            .map(|_| ByteGaugeGuard::new(&self.cold_read_bytes_active, extra as u64));
        let body = Bytes::from(body_vec);
        if !self.ticket_current(ticket) {
            return None;
        }
        let age = cold_age(&entry);
        let state = cache_state(
            age,
            entry.ttl,
            entry.stale_while_revalidate,
            entry.must_revalidate,
            entry.proxy_revalidate,
            entry.requires_revalidation,
            &ticket.request_headers,
        );
        if !state.servable() {
            return None;
        }
        let hits = entry.hits.saturating_add(1);
        let cost = hot_entry_cost(
            body_allocated,
            &metadata.headers,
            &entry.vary,
            &entry.scope,
            &entry.scheme,
            &entry.authority,
            &entry.domain,
            &entry.path,
            &entry.target,
            &entry.tags,
            metadata.origin_etag.as_ref(),
            metadata.origin_last_modified.as_ref(),
        );
        let resident_age_seconds = cold_resident_age(&entry);
        let hot_entry = Entry {
            base_key: entry.base_key,
            vary: entry.vary.clone(),
            scope: entry.scope.clone(),
            scheme: entry.scheme.clone(),
            authority: entry.authority.clone(),
            domain: entry.domain.clone(),
            path: entry.path.clone(),
            target: entry.target.clone(),
            tags: entry.tags.clone(),
            created_at: entry.created_at,
            headers: Arc::new(metadata.headers.clone()),
            body: body.clone(),
            body_allocated,
            received: Instant::now(),
            last_access: Instant::now(),
            resident_since: Instant::now(),
            resident_age_base: Duration::from_secs(resident_age_seconds),
            initial_age: Duration::from_secs(age),
            ttl: entry.ttl,
            stale_while_revalidate: entry.stale_while_revalidate,
            stale_if_error: entry.stale_if_error,
            must_revalidate: entry.must_revalidate,
            proxy_revalidate: entry.proxy_revalidate,
            origin_etag: metadata.origin_etag.clone(),
            origin_last_modified: metadata.origin_last_modified.clone(),
            cost,
            sequence: 0,
            hits,
            transformed: entry.transformed,
            skip_compression: entry.skip_compression,
            requires_revalidation: entry.requires_revalidation,
        };
        if !self.insert_inner(ticket.clone(), hot_entry, settings, false) {
            return None;
        }
        let mut response = Response::new(Body::from(body.clone()));
        *response.headers_mut() = metadata.headers;
        if entry.transformed {
            response
                .extensions_mut()
                .insert(crate::transform::CachedTransformed);
        }
        if entry.skip_compression {
            response.extensions_mut().insert(crate::compression::Skip);
        }
        response.headers_mut().insert(
            "age",
            HeaderValue::from_str(&age.to_string()).expect("integer header"),
        );
        self.domain_stats
            .saved_body_hit(&ticket.scope, &ticket.domain, entry.body_bytes);
        self.counters.hits.fetch_add(1, Ordering::Relaxed);
        self.counters.cold_hits.fetch_add(1, Ordering::Relaxed);
        self.counters
            .cold_promotions
            .fetch_add(1, Ordering::Relaxed);
        Some(Hit {
            response,
            body,
            hits,
            transformed: entry.transformed,
            remaining_ttl: entry.ttl.saturating_sub(age),
            resident_age_seconds,
            stale: state.stale(),
            background_revalidate: state.background_revalidate(),
        })
    }
}