crank/crates/crank-runtime/tests/integration/valkey.rs

use std::{sync::Arc, time::Duration};

use crank_core::{
    CacheBackend, CacheScope, CoordinationStateReservation, CoordinationStateStore,
    CoordinationStateValue, RateLimitDecision, RateLimitStateStore,
};
use crank_runtime::RedisCacheStore;
use futures_util::future::join_all;
use serde_json::json;
use testcontainers::{
    GenericImage,
    core::{IntoContainerPort, WaitFor},
    runners::AsyncRunner,
};

#[tokio::test]
async fn valkey_coordination_and_rate_limit_operations_are_atomic() {
    let container = GenericImage::new("valkey/valkey", "8-alpine")
        .with_exposed_port(6379.tcp())
        .with_wait_for(WaitFor::message_on_stdout("Ready to accept connections"))
        .start()
        .await
        .expect("Valkey test container must start");
    let port = container
        .get_host_port_ipv4(6379.tcp())
        .await
        .expect("Valkey port must be mapped");
    let store = Arc::new(
        RedisCacheStore::connect(CacheBackend::Valkey, &format!("redis://127.0.0.1:{port}/0"))
            .await
            .expect("runtime store must connect to Valkey"),
    );

    verify_atomic_coordination(store.as_ref()).await;
    verify_atomic_rate_limit(store).await;
}

async fn verify_atomic_coordination(store: &RedisCacheStore) {
    let pending = CoordinationStateValue {
        payload: json!({ "state": "pending" }),
    };
    let attempts = (0..32).map(|_| {
        store.reserve_value(
            CacheScope::Coordination,
            "valkey-reservation",
            pending.clone(),
            Duration::from_secs(30),
        )
    });
    let results = join_all(attempts).await;
    assert_eq!(
        results
            .iter()
            .filter(|result| matches!(result, Ok(CoordinationStateReservation::Reserved)))
            .count(),
        1
    );

    let completed = CoordinationStateValue {
        payload: json!({ "state": "completed" }),
    };
    assert!(
        store
            .compare_and_set_value(
                CacheScope::Coordination,
                "valkey-reservation",
                &pending,
                completed.clone(),
                Duration::from_secs(30),
            )
            .await
            .unwrap()
    );
    assert_eq!(
        store
            .take_value(CacheScope::Coordination, "valkey-reservation")
            .await
            .unwrap(),
        Some(completed)
    );
    assert_eq!(
        store
            .take_value(CacheScope::Coordination, "valkey-reservation")
            .await
            .unwrap(),
        None
    );
}

async fn verify_atomic_rate_limit(store: Arc<RedisCacheStore>) {
    let attempts = (0..64).map(|_| {
        let store = Arc::clone(&store);
        async move {
            store
                .consume_token(
                    "valkey-burst",
                    8_000_000,
                    1_000_000,
                    0,
                    Duration::from_secs(30),
                )
                .await
        }
    });
    let results = join_all(attempts).await;
    assert_eq!(
        results
            .iter()
            .filter(|result| matches!(result, Ok(RateLimitDecision::Allowed)))
            .count(),
        8
    );
    assert!(
        results
            .iter()
            .filter_map(|result| result.as_ref().ok())
            .all(|decision| matches!(
                decision,
                RateLimitDecision::Allowed
                    | RateLimitDecision::Rejected {
                        retry_after_ms: 1_000
                    }
            ))
    );

    assert_eq!(
        store
            .consume_token(
                "valkey-retry-after",
                2_000_000,
                2_000_000,
                0,
                Duration::from_secs(30),
            )
            .await
            .unwrap(),
        RateLimitDecision::Allowed
    );
    assert_eq!(
        store
            .consume_token(
                "valkey-retry-after",
                2_000_000,
                2_000_000,
                0,
                Duration::from_secs(30),
            )
            .await
            .unwrap(),
        RateLimitDecision::Allowed
    );
    assert_eq!(
        store
            .consume_token(
                "valkey-retry-after",
                2_000_000,
                2_000_000,
                0,
                Duration::from_secs(30),
            )
            .await
            .unwrap(),
        RateLimitDecision::Rejected {
            retry_after_ms: 500
        }
    );
}