RustProxy/src/cache/lookup/tiered.rs

impl Cache {
    pub async fn get_tiered(&self, ticket: &Ticket, settings: &Settings) -> Option<Hit> {
        if let Some(hit) = self.get_hot(ticket, false) {
            return Some(hit);
        }
        if let Some(hit) = self.get_cold(ticket, settings).await {
            return Some(hit);
        }
        self.counters.misses.fetch_add(1, Ordering::Relaxed);
        None
    }
    pub fn publish_revalidated(
        &self,
        ticket: Ticket,
        cached: &Revalidation,
        mut headers: HeaderMap,
        validation_headers: &HeaderMap,
        settings: &Settings,
        response_delay: Duration,
    ) -> Option<Hit> {
        let plan = plan(StatusCode::OK, &headers, settings, response_delay)?;
        if cached.body.len() > settings.max_object_bytes {
            return None;
        }
        if !headers.contains_key(::http::header::DATE) {
            if let Ok(value) = httpdate::fmt_http_date(SystemTime::now()).parse() {
                headers.insert(::http::header::DATE, value);
            }
        }
        headers.remove("x-request-id");
        let hits = cached.hits.saturating_add(1);
        let age = plan.initial_age.as_secs();
        let remaining_ttl = plan.ttl.saturating_sub(age);
        let vary = plan.vary.clone();
        let tags = cache_tags(&headers);
        let origin_etag = validation_headers
            .get(::http::header::ETAG)
            .cloned()
            .or_else(|| cached.origin_etag.clone());
        let origin_last_modified = validation_headers
            .get(::http::header::LAST_MODIFIED)
            .cloned()
            .or_else(|| cached.origin_last_modified.clone());
        let cost = hot_entry_cost(
            cached.body_allocated,
            &headers,
            &vary,
            &ticket.scope,
            &ticket.scheme,
            &ticket.authority,
            &ticket.domain,
            &ticket.path,
            &ticket.target,
            &tags,
            origin_etag.as_ref(),
            origin_last_modified.as_ref(),
        );
        let storage_key = variant_key(&ticket.key, &vary, &ticket.request_headers);
        let resident_age = cached.resident_age_at_lookup.saturating_add(response_delay);
        let entry = Entry {
            base_key: ticket.key,
            vary,
            scope: ticket.scope.clone(),
            scheme: ticket.scheme.clone(),
            authority: ticket.authority.clone(),
            domain: ticket.domain.clone(),
            path: ticket.path.clone(),
            target: ticket.target.clone(),
            tags,
            created_at: ticket.created_at,
            headers: Arc::new(headers.clone()),
            body: cached.body.clone(),
            body_allocated: cached.body_allocated,
            received: plan.received,
            last_access: Instant::now(),
            resident_since: Instant::now(),
            resident_age_base: resident_age,
            initial_age: plan.initial_age,
            ttl: plan.ttl,
            stale_while_revalidate: plan.stale_while_revalidate,
            stale_if_error: plan.stale_if_error,
            must_revalidate: plan.must_revalidate,
            proxy_revalidate: plan.proxy_revalidate,
            origin_etag,
            origin_last_modified,
            cost,
            sequence: 0,
            hits,
            transformed: cached.transformed,
            skip_compression: cached.skip_compression,
            requires_revalidation: plan.requires_revalidation,
        };
        if !self.insert(ticket, entry, settings) {
            return None;
        }
        if cached.storage_key != storage_key {
            self.discard_storage_key(&cached.storage_key);
        }
        self.counters.hits.fetch_add(1, Ordering::Relaxed);
        let mut response = cached.response(headers);
        response.headers_mut().insert(::http::header::AGE,
            HeaderValue::from(age),
        );
        Some(Hit {
            response,
            body: cached.body.clone(),
            hits,
            transformed: cached.transformed,
            remaining_ttl,
            cache_status_ttl: cache_status_ttl(plan.ttl, age),
            resident_age_seconds: resident_age.as_secs(),
            stale: false,
            background_revalidate: false,
        })
    }

    pub fn wrap(
        self: &Arc<Self>,
        response: Response,
        ticket: Ticket,
        plan: Plan,
        settings: &Settings,
    ) -> Response {
        self.wrap_with_fill(response, ticket, plan, settings, None)
    }

    pub fn wrap_with_fill(
        self: &Arc<Self>,
        response: Response,
        ticket: Ticket,
        plan: Plan,
        settings: &Settings,
        fill_leader: Option<FillLeaderGuard>,
    ) -> Response {
        let fill_size = plan.expected_length.unwrap_or(settings.max_object_bytes);
        let permits = if fill_size > 64 * 1024 * 1024 {
            FILL_SLOTS as u32
        } else if fill_size > 16 * 1024 * 1024 {
            4
        } else if fill_size > 4 * 1024 * 1024 {
            2
        } else {
            1
        };
        let permit = self.fills.clone().try_acquire_many_owned(permits);
        let Ok(permit) = permit else {
            self.counters.rejected_fills.fetch_add(1, Ordering::Relaxed);
            return response;
        };
        let (parts, body) = response.into_parts();
        let headers = parts.headers.clone();
        let origin_validators = parts
            .extensions
            .get::<crate::transform::OriginValidators>()
            .cloned();
        let origin_etag = match origin_validators.as_ref() {
            Some(validators) => validators.etag.clone(),
            None => headers.get(::http::header::ETAG).cloned(),
        };
        let origin_last_modified = match origin_validators.as_ref() {
            Some(validators) => validators.last_modified.clone(),
            None => headers.get(::http::header::LAST_MODIFIED).cloned(),
        };
        let transformed = parts
            .extensions
            .get::<crate::transform::Applied>()
            .is_some();
        let initial_capacity = if transformed {
            0
        } else {
            fill_size.min(64 * 1024)
        };
        let Some(mut fill_gauge) = FillGauge::new(self.clone(), initial_capacity, settings) else {
            self.counters.rejected_fills.fetch_add(1, Ordering::Relaxed);
            self.record_reason(&ticket.scope, &ticket.domain, "hot_memory_pressure");
            return Response::from_parts(parts, body);
        };
        let data = Vec::with_capacity(initial_capacity);
        let actual_capacity = data.capacity();
        if actual_capacity > initial_capacity
            && !fill_gauge.grow(actual_capacity - initial_capacity, settings)
        {
            self.counters.rejected_fills.fetch_add(1, Ordering::Relaxed);
            self.record_reason(&ticket.scope, &ticket.domain, "fill_allocator_overshoot");
            return Response::from_parts(parts, body);
        }
        let candidate = Candidate {
            cache: self.clone(),
            ticket,
            plan,
            headers,
            data,
            shared: None,
            shared_allocated: 0,
            fill_gauge,
            settings: settings.clone(),
            skip_compression: parts.extensions.get::<crate::compression::Skip>().is_some(),
            origin_etag,
            origin_last_modified,
            transformed,
            _permit: permit,
            _fill_leader: fill_leader,
        };
        if body.is_end_stream() {
            candidate.finish();
            return Response::from_parts(parts, body);
        }
        Response::from_parts(
            parts,
            Body::new(FillBody {
                inner: body,
                candidate: Some(candidate),
                done: false,
            }),
        )
    }
}