Add in-process endpoint caching with moka

Cargo.toml

@@ -23,6 +23,7 @@ sqlx = { version = "0.8.6", features = ["runtime-tokio", "sqlite", "migrate"] }
 argon2 = "0.5.3"
 base64 = "0.22.1"
 clap = { version = "4.5.58", features = ["derive"] }
+moka = { version = "0.12", features = ["future"] }
 toml = "0.8"
 reqwest = { version = "0.13.2", features = ["json"] }
 rain_orderbook_js_api = { path = "lib/rain.orderbook/crates/js_api", default-features = false }

src/cache.rs

@@ -0,0 +1,158 @@
+use moka::future::Cache;
+use std::future::Future;
+use std::sync::Arc;
+use std::time::Duration;
+
+pub(crate) struct AppCache<K, V>(Cache<K, V>)
+where
+    K: std::hash::Hash + Eq + Send + Sync + 'static,
+    V: Clone + Send + Sync + 'static;
+
+impl<K, V> AppCache<K, V>
+where
+    K: std::hash::Hash + Eq + Send + Sync + 'static,
+    V: Clone + Send + Sync + 'static,
+{
+    pub(crate) fn new(max_capacity: u64, ttl: Duration) -> Self {
+        Self(
+            Cache::builder()
+                .max_capacity(max_capacity)
+                .time_to_live(ttl)
+                .build(),
+        )
+    }
+
+    pub(crate) async fn get(&self, key: &K) -> Option<V> {
+        self.0.get(key).await
+    }
+
+    pub(crate) async fn insert(&self, key: K, value: V) {
+        self.0.insert(key, value).await
+    }
+
+    pub(crate) async fn get_or_try_insert<F, Fut, E>(&self, key: K, fetch: F) -> Result<V, Arc<E>>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = Result<V, E>>,
+        E: Send + Sync + 'static,
+    {
+        self.0.try_get_with(key, fetch()).await
+    }
+
+    pub(crate) fn invalidate_all(&self) {
+        self.0.invalidate_all()
+    }
+}
+
+trait Invalidatable: Send + Sync {
+    fn invalidate_all(&self);
+}
+
+impl<K, V> Invalidatable for Cache<K, V>
+where
+    K: std::hash::Hash + Eq + Send + Sync + 'static,
+    V: Clone + Send + Sync + 'static,
+{
+    fn invalidate_all(&self) {
+        Cache::invalidate_all(self)
+    }
+}
+
+pub(crate) struct CacheGroup {
+    caches: Vec<Arc<dyn Invalidatable>>,
+}
+
+impl CacheGroup {
+    pub(crate) fn new() -> Self {
+        Self { caches: Vec::new() }
+    }
+
+    pub(crate) fn register<K, V>(&mut self, cache: &AppCache<K, V>)
+    where
+        K: std::hash::Hash + Eq + Send + Sync + 'static,
+        V: Clone + Send + Sync + 'static,
+    {
+        self.caches.push(Arc::new(cache.0.clone()));
+    }
+
+    pub(crate) fn invalidate_all(&self) {
+        for cache in &self.caches {
+            cache.invalidate_all();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[rocket::async_test]
+    async fn test_app_cache_insert_and_get() {
+        let cache: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        cache.insert("key", 42).await;
+        assert_eq!(cache.get(&"key").await, Some(42));
+    }
+
+    #[rocket::async_test]
+    async fn test_app_cache_get_returns_none_for_missing_key() {
+        let cache: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        assert!(cache.get(&"missing").await.is_none());
+    }
+
+    #[rocket::async_test]
+    async fn test_app_cache_invalidate_all_clears_entries() {
+        let cache: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        cache.insert("a", 1).await;
+        cache.insert("b", 2).await;
+        cache.invalidate_all();
+        tokio::task::yield_now().await;
+        assert!(cache.get(&"a").await.is_none());
+        assert!(cache.get(&"b").await.is_none());
+    }
+
+    #[rocket::async_test]
+    async fn test_get_or_try_insert_calls_fetch_on_miss() {
+        let cache: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        let result: Result<u32, Arc<String>> =
+            cache.get_or_try_insert("key", || async { Ok(42) }).await;
+        assert_eq!(result.unwrap(), 42);
+        assert_eq!(cache.get(&"key").await, Some(42));
+    }
+
+    #[rocket::async_test]
+    async fn test_get_or_try_insert_returns_cached_on_hit() {
+        let cache: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        cache.insert("key", 42).await;
+        let result: Result<u32, Arc<String>> = cache
+            .get_or_try_insert("key", || async { panic!("fetch should not be called") })
+            .await;
+        assert_eq!(result.unwrap(), 42);
+    }
+
+    #[rocket::async_test]
+    async fn test_get_or_try_insert_does_not_cache_errors() {
+        let cache: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        let result: Result<u32, Arc<String>> = cache
+            .get_or_try_insert("key", || async { Err("fail".to_string()) })
+            .await;
+        assert!(result.is_err());
+        assert!(cache.get(&"key").await.is_none());
+    }
+
+    #[rocket::async_test]
+    async fn test_cache_group_invalidate_all_clears_registered_caches() {
+        let cache_a: AppCache<&str, u32> = AppCache::new(10, Duration::from_secs(60));
+        let cache_b: AppCache<u32, String> = AppCache::new(10, Duration::from_secs(60));
+        cache_a.insert("x", 10).await;
+        cache_b.insert(1, "hello".into()).await;
+
+        let mut group = CacheGroup::new();
+        group.register(&cache_a);
+        group.register(&cache_b);
+        group.invalidate_all();
+
+        tokio::task::yield_now().await;
+        assert!(cache_a.get(&"x").await.is_none());
+        assert!(cache_b.get(&1).await.is_none());
+    }
+}

src/error.rs

@@ -36,6 +36,12 @@ pub enum ApiError {
     RateLimited(String),
 }
 
+impl From<std::sync::Arc<ApiError>> for ApiError {
+    fn from(arc: std::sync::Arc<ApiError>) -> Self {
+        (*arc).clone()
+    }
+}
+
 impl<'r> Responder<'r, 'static> for ApiError {
     fn respond_to(self, req: &'r Request<'_>) -> rocket::response::Result<'static> {
         let (status, code, message) = match &self {

src/main.rs

@@ -2,6 +2,7 @@
 extern crate rocket;
 
 mod auth;
+mod cache;
 mod catchers;
 mod cli;
 mod config;
@@ -121,10 +122,20 @@ pub(crate) fn rocket(
 
     let options = Options::Index | Options::NormalizeDirs;
 
+    let order_cache = routes::order::order_detail_cache();
+    let swap_cache = routes::swap::swap_quote_cache();
+
+    let mut registry_caches = cache::CacheGroup::new();
+    registry_caches.register(&order_cache);
+    registry_caches.register(&swap_cache);
+
     Ok(rocket::custom(figment)
         .manage(pool)
         .manage(rate_limiter)
         .manage(raindex_config)
+        .manage(order_cache)
+        .manage(swap_cache)
+        .manage(registry_caches)
         .mount("/", routes::health::routes())
         .mount("/v1/tokens", routes::tokens::routes())
         .mount("/v1/swap", routes::swap::routes())

src/routes/admin.rs

@@ -1,4 +1,5 @@
 use crate::auth::AdminKey;
+use crate::cache::CacheGroup;
 use crate::db::{settings, DbPool};
 use crate::error::{ApiError, ApiErrorResponse};
 use crate::fairings::{GlobalRateLimit, TracingSpan};
@@ -37,6 +38,7 @@ pub async fn put_registry(
     pool: &State<DbPool>,
     span: TracingSpan,
     request: Json<UpdateRegistryRequest>,
+    registry_caches: &State<CacheGroup>,
 ) -> Result<Json<RegistryResponse>, ApiError> {
     let req = request.into_inner();
     async move {
@@ -67,7 +69,9 @@ pub async fn put_registry(
         *guard = new_provider;
         drop(guard);
 
-        tracing::info!(registry_url = %req.registry_url, "registry updated");
+        registry_caches.invalidate_all();
+
+        tracing::info!(registry_url = %req.registry_url, "registry updated, caches invalidated");
 
         Ok(Json(RegistryResponse {
             registry_url: req.registry_url,
@@ -225,4 +229,92 @@ mod tests {
 
         assert_eq!(response.status(), Status::BadRequest);
     }
+
+    #[rocket::async_test]
+    async fn test_put_registry_invalidates_caches() {
+        use crate::routes::order::OrderDetailCache;
+        use crate::routes::swap::SwapQuoteCache;
+        use crate::types::common::TokenRef;
+        use crate::types::order::{OrderDetail, OrderDetailsInfo, OrderType};
+        use crate::types::swap::SwapQuoteResponse;
+        use alloy::primitives::{address, Address, U256};
+
+        let client = TestClientBuilder::new().build().await;
+        let (key_id, secret) = seed_admin_key(&client).await;
+        let header = basic_auth_header(&key_id, &secret);
+
+        let order_hash = "0x000000000000000000000000000000000000000000000000000000000000abcd"
+            .parse()
+            .unwrap();
+        let dummy_order = OrderDetail {
+            order_hash,
+            owner: Address::ZERO,
+            order_details: OrderDetailsInfo {
+                type_: OrderType::Solver,
+                io_ratio: "1.0".into(),
+            },
+            input_token: TokenRef {
+                address: Address::ZERO,
+                symbol: "USDC".into(),
+                decimals: 6,
+            },
+            output_token: TokenRef {
+                address: Address::ZERO,
+                symbol: "WETH".into(),
+                decimals: 18,
+            },
+            input_vault_id: U256::ZERO,
+            output_vault_id: U256::ZERO,
+            input_vault_balance: "0".into(),
+            output_vault_balance: "0".into(),
+            io_ratio: "1.0".into(),
+            created_at: 0,
+            orderbook_id: Address::ZERO,
+            trades: vec![],
+        };
+        let order_cache = client
+            .rocket()
+            .state::<OrderDetailCache>()
+            .expect("OrderDetailCache in state");
+        order_cache.insert(order_hash, dummy_order).await;
+        assert!(order_cache.get(&order_hash).await.is_some());
+
+        let usdc = address!("833589fCD6eDb6E08f4c7C32D4f71b54bdA02913");
+        let weth = address!("4200000000000000000000000000000000000006");
+        let cache_key = (usdc, weth, "100".to_string());
+        let dummy_quote = SwapQuoteResponse {
+            input_token: usdc,
+            output_token: weth,
+            output_amount: "100".into(),
+            estimated_output: "100".into(),
+            estimated_input: "150".into(),
+            estimated_io_ratio: "1.5".into(),
+        };
+        let swap_cache = client
+            .rocket()
+            .state::<SwapQuoteCache>()
+            .expect("SwapQuoteCache in state");
+        swap_cache.insert(cache_key.clone(), dummy_quote).await;
+        assert!(swap_cache.get(&cache_key).await.is_some());
+
+        let new_url = mock_raindex_registry_url().await;
+        let response = client
+            .put("/admin/registry")
+            .header(Header::new("Authorization", header))
+            .header(ContentType::JSON)
+            .body(format!(r#"{{"registry_url":"{new_url}"}}"#))
+            .dispatch()
+            .await;
+        assert_eq!(response.status(), Status::Ok);
+
+        tokio::task::yield_now().await;
+        assert!(
+            order_cache.get(&order_hash).await.is_none(),
+            "order cache must be empty after registry update"
+        );
+        assert!(
+            swap_cache.get(&cache_key).await.is_none(),
+            "swap cache must be empty after registry update"
+        );
+    }
 }