RustProxy/src/cache/tests/cold_tier_tests/part_01.rs

use super::*;
use http_body_util::BodyExt;

fn cold_settings() -> Settings {
    Settings {
        enabled: true,
        max_entries: 1,
        max_bytes: 256 * 1024,
        max_object_bytes: 64 * 1024,
        cold_enabled: true,
        cold_path: "cache/cold-test".into(),
        cold_max_bytes: 2 * 1024 * 1024,
        cold_max_entries: 16,
        cold_after_seconds: 300,
        ..Settings::default()
    }
}

async fn fill(cache: &Arc<Cache>, settings: &Settings, path: &str, body: &'static str) -> Ticket {
    let ticket = cache
        .ticket_scoped(
            11,
            "cold-test",
            "http",
            "cold.example",
            "origin@127.0.0.1:8080",
            path,
            &HeaderMap::new(),
        )
        .unwrap();
    let mut response = Response::new(Body::from(body));
    response.headers_mut().insert(
        "cache-control",
        HeaderValue::from_static("public, max-age=120"),
    );
    response
        .headers_mut()
        .insert("etag", HeaderValue::from_static("\"cold-v1\""));
    let plan = plan(StatusCode::OK, response.headers(), settings, Duration::ZERO).unwrap();
    cache
        .wrap(response, ticket.clone(), plan, settings)
        .into_body()
        .collect()
        .await
        .unwrap();
    ticket
}

async fn wait_for_cold(cache: &Cache, expected: usize) {
    for _ in 0..100 {
        if cache.view().cold_entries >= expected {
            return;
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
    panic!("cold tier did not persist in time");
}

#[tokio::test]
async fn selective_purge_removes_matching_cold_object() {
    let dir = tempfile::tempdir().unwrap();
    let settings = cold_settings();
    let shutdown = CancellationToken::new();
    let cache = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        shutdown.clone(),
    );
    let first = fill(&cache, &settings, "/one?v=1", "first-payload").await;
    let _second = fill(&cache, &settings, "/two", "second-payload").await;
    wait_for_cold(&cache, 1).await;
    assert_eq!(cache.view().cold_entries, 1);
    let result = cache
        .invalidate(InvalidationRequest {
            operation: InvalidationOperation::Purge,
            selector: InvalidationSelector::Exact,
            value: "/one?v=1".into(),
            host: Some("cold.example".into()),
            vhost: Some("cold-test".into()),
            ttl_seconds: None,
            confirm_wide: false,
        })
        .unwrap();
    assert_eq!(result.affected_entries, 1);
    assert_eq!(cache.view().cold_entries, 0);
    assert!(cache.get_tiered(&first, &settings).await.is_none());
    shutdown.cancel();
}

#[tokio::test]
async fn lru_demotion_persists_and_survives_restart() {
    let dir = tempfile::tempdir().unwrap();
    let settings = cold_settings();
    settings.validate().unwrap();
    let shutdown = CancellationToken::new();
    let cache = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        shutdown.clone(),
    );
    let first = fill(&cache, &settings, "/one", "first-payload").await;
    let _second = fill(&cache, &settings, "/two", "second-payload").await;
    wait_for_cold(&cache, 1).await;
    assert_eq!(cache.view().entries, 1);
    assert_eq!(cache.view().cold_entries, 1);
    let cold = cache.dashboard(Some("cold-test"), 0, 50);
    assert!(cold
        .objects
        .iter()
        .any(|row| row.path == "/one" && row.tier == "cold"));

    shutdown.cancel();
    drop(cache);

    let restart_shutdown = CancellationToken::new();
    let restarted = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        restart_shutdown.clone(),
    );
    assert_eq!(restarted.view().cold_entries, 1);
    let hit = restarted
        .get_tiered(&first, &settings)
        .await
        .expect("cold object should be readable after restart");
    let payload = hit.response.into_body().collect().await.unwrap().to_bytes();
    assert_eq!(payload, Bytes::from_static(b"first-payload"));
    assert_eq!(restarted.view().cold_hits, 1);
    assert_eq!(restarted.view().cold_promotions, 1);
    restart_shutdown.cancel();
}

#[tokio::test]
async fn cold_restart_preserves_transformed_validators_stale_metadata_and_no_transform() {
    let dir = tempfile::tempdir().unwrap();
    let mut settings = cold_settings();
    settings.max_entries = 8;
    let shutdown = CancellationToken::new();
    let cache = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        shutdown.clone(),
    );

    let transformed_ticket = cache
        .ticket_scoped(
            11,
            "cold-test",
            "http",
            "cold.example",
            "origin@127.0.0.1:8080",
            "/transformed.json",
            &HeaderMap::new(),
        )
        .unwrap();
    let mut transformed = Response::new(Body::from(r#"{"a":1}"#));
    transformed.headers_mut().insert(
        "cache-control",
        HeaderValue::from_static(
            "public, max-age=120, stale-while-revalidate=30, stale-if-error=90",
        ),
    );
    transformed.headers_mut().insert(
        "etag",
        HeaderValue::from_static("\"rustproxy-representation\""),
    );
    transformed
        .headers_mut()
        .insert("content-length", HeaderValue::from_static("7"));
    transformed
        .extensions_mut()
        .insert(crate::transform::OriginValidators {
            etag: Some(HeaderValue::from_static("\"origin-v1\"")),
            last_modified: Some(HeaderValue::from_static("Mon, 26 Jan 2026 23:11:40 GMT")),
        });
    transformed
        .extensions_mut()
        .insert(crate::transform::Applied {
            kind: crate::transform::Kind::Json,
            original_bytes: 9,
            transformed_bytes: 7,
        });
    let transformed_plan = plan(
        StatusCode::OK,
        transformed.headers(),
        &settings,
        Duration::ZERO,
    )
    .unwrap();
    cache
        .wrap(
            transformed,
            transformed_ticket.clone(),
            transformed_plan,
            &settings,
        )
        .into_body()
        .collect()
        .await
        .unwrap();

    let protected_ticket = cache
        .ticket_scoped(
            11,
            "cold-test",
            "http",
            "cold.example",
            "origin@127.0.0.1:8080",
            "/protected.css",
            &HeaderMap::new(),
        )
        .unwrap();
    let original_body = "body { color: red; }";
    let mut protected = Response::new(Body::from(original_body));
    protected.headers_mut().insert(
        "cache-control",
        HeaderValue::from_static("public, max-age=120, immutable, no-transform"),
    );
    protected
        .headers_mut()
        .insert("etag", HeaderValue::from_static("\"origin-css\""));
    protected.extensions_mut().insert(crate::compression::Skip);
    let protected_plan = plan(
        StatusCode::OK,
        protected.headers(),
        &settings,
        Duration::ZERO,
    )
    .unwrap();
    cache
        .wrap(
            protected,
            protected_ticket.clone(),
            protected_plan,
            &settings,
        )
        .into_body()
        .collect()
        .await
        .unwrap();

    assert_eq!(cache.flush_hot_to_cold().await, 2);
    wait_for_cold(&cache, 2).await;
    shutdown.cancel();
    drop(cache);

    let restart_shutdown = CancellationToken::new();
    let restarted = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        restart_shutdown.clone(),
    );
    assert_eq!(restarted.view().cold_entries, 2);
    // COLD entries intentionally keep validators and headers on disk to bound RAM.
    // Inspect them through the same verified metadata loader used on COLD hits.
    let (key, transformed) = {
        let cold = restarted.cold.lock().unwrap();
        cold.entries
            .iter()
            .find(|(_, entry)| entry.path == "/transformed.json")
            .map(|(key, entry)| (*key, entry.clone()))
            .expect("transformed cold entry")
    };
    assert_eq!(transformed.stale_while_revalidate, 30);
    assert_eq!(transformed.stale_if_error, 90);
    assert!(transformed.has_validator);
    assert!(transformed.transformed);
    let metadata = read_cold_metadata(&key, &transformed)
        .await
        .expect("persisted transformed metadata");
    assert_eq!(metadata.origin_etag.as_ref().unwrap(), "\"origin-v1\"");
    assert_eq!(
        metadata.origin_last_modified.as_ref().unwrap(),
        "Mon, 26 Jan 2026 23:11:40 GMT"
    );
    assert_eq!(metadata.headers["etag"], "\"rustproxy-representation\"");

    let transformed_hit = restarted
        .get_tiered(&transformed_ticket, &settings)
        .await
        .expect("transformed object after restart");
    assert_eq!(
        transformed_hit.response.headers()["etag"],
        "\"rustproxy-representation\""
    );
    assert!(transformed_hit.transformed);
    assert!(transformed_hit
        .response
        .extensions()
        .get::<crate::transform::CachedTransformed>()
        .is_some());
    assert_eq!(
        transformed_hit
            .response
            .into_body()
            .collect()
            .await
            .unwrap()
            .to_bytes(),
        Bytes::from_static(br#"{"a":1}"#)
    );

    let protected_hit = restarted
        .get_tiered(&protected_ticket, &settings)
        .await
        .expect("no-transform object after restart");
    assert_eq!(protected_hit.response.headers()["etag"], "\"origin-css\"");
    assert!(crate::transform::no_transform(
        protected_hit.response.headers()
    ));
    assert!(protected_hit
        .response
        .extensions()
        .get::<crate::compression::Skip>()
        .is_some());
    assert_eq!(
        protected_hit
            .response
            .into_body()
            .collect()
            .await
            .unwrap()
            .to_bytes(),
        Bytes::from_static(b"body { color: red; }")
    );
    restart_shutdown.cancel();
}

#[tokio::test]
async fn cold_policy_can_disable_lru_demotion() {
    let dir = tempfile::tempdir().unwrap();
    let mut settings = cold_settings();
    settings.cold_on_lru = false;
    let shutdown = CancellationToken::new();
    let cache = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        shutdown.clone(),
    );
    fill(&cache, &settings, "/one", "first-payload").await;
    fill(&cache, &settings, "/two", "second-payload").await;
    tokio::time::sleep(Duration::from_millis(30)).await;
    let view = cache.view();
    assert_eq!(view.entries, 1);
    assert_eq!(view.cold_entries, 0);
    assert_eq!(view.evictions, 1);
    shutdown.cancel();
}

#[tokio::test]
async fn cold_policy_can_disable_idle_demotion() {
    let dir = tempfile::tempdir().unwrap();
    let mut settings = cold_settings();
    settings.max_entries = 8;
    settings.cold_on_idle = false;
    let shutdown = CancellationToken::new();
    let cache = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        shutdown.clone(),
    );
    fill(&cache, &settings, "/idle", "idle-payload").await;
    {
        let mut store = cache.store.lock().unwrap();
        for entry in store.entries.values_mut() {
            entry.last_access = Instant::now() - Duration::from_secs(600);
        }
    }
    cache.demote_idle();
    assert_eq!(cache.view().entries, 1);
    assert_eq!(cache.view().cold_entries, 0);

    settings.cold_on_idle = true;
    cache.configure(settings.clone());
    cache.demote_idle();
    wait_for_cold(&cache, 1).await;
    assert_eq!(cache.view().entries, 0);
    shutdown.cancel();
}

#[tokio::test]
async fn cold_policy_can_demote_by_total_hot_residency_age() {
    let dir = tempfile::tempdir().unwrap();
    let mut settings = cold_settings();
    settings.max_entries = 8;
    settings.cold_on_idle = false;
    settings.cold_on_age = true;
    settings.cold_max_hot_seconds = 5;
    let shutdown = CancellationToken::new();
    let cache = Cache::start(
        11,
        Some(dir.path().to_path_buf()),
        settings.clone(),
        shutdown.clone(),
    );
    fill(&cache, &settings, "/aged", "aged-payload").await;
    {
        let mut store = cache.store.lock().unwrap();
        for entry in store.entries.values_mut() {
            entry.resident_since = Instant::now() - Duration::from_secs(60);
            entry.last_access = Instant::now();
        }
    }
    cache.demote_idle();
    wait_for_cold(&cache, 1).await;
    assert_eq!(cache.view().entries, 0);
    assert_eq!(cache.view().cold_entries, 1);
    shutdown.cancel();
}