perf: resolve global lock contention in template server

examples/server/Cargo.toml

@@ -14,6 +14,7 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+dashmap = "5.5"
 
 [dev-dependencies]
 tower = { version = "0.5", features = ["util"] }

examples/server/src/main.rs

@@ -22,12 +22,14 @@ async fn ttl_sweeper(state: AppState) {
     loop {
         iv.tick().await;
         let now = Instant::now();
-        let mut sessions = state.sessions.write().await;
-        let before = sessions.len();
-        sessions.retain(|_, e| now.duration_since(e.last_touch) < SESSION_TTL);
-        let evicted = before - sessions.len();
+        let before = state.sessions.len();
+        state.sessions.retain(|_, e| {
+            let last_touch = *e.last_touch.lock().unwrap();
+            now.duration_since(last_touch) < SESSION_TTL
+        });
+        let evicted = before - state.sessions.len();
         if evicted > 0 {
-            tracing::info!(evicted, remaining = sessions.len(), "ttl sweep");
+            tracing::info!(evicted, remaining = state.sessions.len(), "ttl sweep");
         }
     }
 }

examples/server/src/routes.rs

@@ -223,13 +223,18 @@ async fn ingest(
     };
     let ingested = new_lines.len();
 
-    let mut sessions = state.sessions.write().await;
-    let entry = sessions.entry(id.clone()).or_insert_with(SessionEntry::new);
+    let entry = if let Some(e) = state.sessions.get(&id) {
+        e.clone()
+    } else {
+        state.sessions.entry(id.clone()).or_insert_with(|| std::sync::Arc::new(SessionEntry::new())).value().clone()
+    };
+
+    let mut index = entry.index.write().await;
     for l in new_lines {
-        entry.index.push(l);
+        index.push(l);
     }
-    entry.last_touch = Instant::now();
-    let total = entry.index.len();
+    *entry.last_touch.lock().unwrap() = Instant::now();
+    let total = index.len();
 
     tracing::info!(session = %id, ingested, total, "ingest");
     Json(json!({ "ingested": ingested, "total": total })).into_response()
@@ -244,11 +249,12 @@ async fn templates(
         Ok(cfg) => cfg,
         Err(resp) => return resp,
     };
-    let sessions = state.sessions.read().await;
-    let Some(entry) = sessions.get(&id) else {
-        return (StatusCode::NOT_FOUND, format!("no such session: {id}")).into_response();
+    let entry = match state.sessions.get(&id) {
+        Some(e) => e.clone(),
+        None => return (StatusCode::NOT_FOUND, format!("no such session: {id}")).into_response(),
     };
-    let result = entry.index.templates_with(&cfg);
+    let index = entry.index.read().await;
+    let result = index.templates_with(&cfg);
     if wants_json(q.format.as_deref()) {
         Json(result).into_response()
     } else {

examples/server/src/session.rs

@@ -5,12 +5,12 @@
 //! lives entirely in memory; sessions idle longer than [`SESSION_TTL`] are
 //! evicted by a background task (see `main.rs`).
 
-use std::collections::HashMap;
 use std::env;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use codag_drain::{TemplateIndex, TemplaterConfig};
+use dashmap::DashMap;
 use tokio::sync::{RwLock, Semaphore};
 
 pub use codag_drain::{parse_body, parse_json_line, parse_line, BodyFormat};
@@ -67,8 +67,8 @@ fn env_u64(name: &str, default: u64) -> u64 {
 /// One live session: a streaming index plus its last-touch timestamp.
 #[derive(Debug)]
 pub struct SessionEntry {
-    pub index: TemplateIndex,
-    pub last_touch: Instant,
+    pub index: RwLock<TemplateIndex>,
+    pub last_touch: std::sync::Mutex<Instant>,
 }
 
 impl SessionEntry {
@@ -80,16 +80,16 @@ impl SessionEntry {
 impl Default for SessionEntry {
     fn default() -> Self {
         SessionEntry {
-            index: TemplateIndex::new(TemplaterConfig::default()),
-            last_touch: Instant::now(),
+            index: RwLock::new(TemplateIndex::new(TemplaterConfig::default())),
+            last_touch: std::sync::Mutex::new(Instant::now()),
         }
     }
 }
 
 /// Shared, cloneable application state.
 #[derive(Clone)]
 pub struct AppState {
-    pub sessions: Arc<RwLock<HashMap<String, SessionEntry>>>,
+    pub sessions: Arc<DashMap<String, Arc<SessionEntry>>>,
     pub limits: Arc<Limits>,
     pub template_slots: Arc<Semaphore>,
     pub auth_token: Option<Arc<str>>,
@@ -100,7 +100,7 @@ impl AppState {
         let limits = Limits::from_env();
         let max_inflight = limits.max_inflight;
         AppState {
-            sessions: Arc::new(RwLock::new(HashMap::new())),
+            sessions: Arc::new(DashMap::new()),
             limits: Arc::new(limits),
             template_slots: Arc::new(Semaphore::new(max_inflight)),
             auth_token: auth_token_from_env(),

examples/server/tests/integration.rs

@@ -4,11 +4,11 @@
 use axum::body::Body;
 use axum::http::{HeaderValue, Request, StatusCode};
 use http_body_util::BodyExt;
-use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::sync::{RwLock, Semaphore};
+use tokio::sync::Semaphore;
 use tower::ServiceExt;
+use dashmap::DashMap;
 
 use codag_drain::{template_logs, LogLine, TemplaterConfig};
 use codag_drain_server::routes;
@@ -48,7 +48,7 @@ fn limited_app(limits: Limits) -> axum::Router {
 fn state_with_limits(limits: Limits) -> AppState {
     let max_inflight = limits.max_inflight;
     AppState {
-        sessions: Arc::new(RwLock::new(HashMap::new())),
+        sessions: Arc::new(DashMap::new()),
         limits: Arc::new(limits),
         template_slots: Arc::new(Semaphore::new(max_inflight)),
         auth_token: Some(Arc::<str>::from(TEST_TOKEN)),
@@ -105,7 +105,7 @@ async fn v1_routes_fail_closed_without_configured_auth_token() {
     let limits = Limits::from_env();
     let max_inflight = limits.max_inflight;
     let state = AppState {
-        sessions: Arc::new(RwLock::new(HashMap::new())),
+        sessions: Arc::new(DashMap::new()),
         limits: Arc::new(limits),
         template_slots: Arc::new(Semaphore::new(max_inflight)),
         auth_token: None,
@@ -345,7 +345,7 @@ async fn concurrency_overflow_returns_429() {
         template_timeout: Duration::from_secs(1),
     };
     let state = AppState {
-        sessions: Arc::new(RwLock::new(HashMap::new())),
+        sessions: Arc::new(DashMap::new()),
         limits: Arc::new(limits),
         template_slots: Arc::new(Semaphore::new(1)),
         auth_token: Some(Arc::<str>::from(TEST_TOKEN)),