perf(gateway): scale request hot paths for 20k streams

Shard and singleflight hot-path caches, batch and prioritize candidate and usage lifecycle persistence, and extend database and pressure-test instrumentation for 20k concurrent streams.

apps/aether-gateway/src/data/state/candidate_cache.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 use std::future::Future;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
@@ -12,12 +12,13 @@ use aether_data_contracts::repository::candidate_selection::{
     StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
 };
 use async_trait::async_trait;
-use tokio::sync::Notify;
+use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore};
 use tokio::time::timeout;
 use tracing::warn;
 
 const CANDIDATE_SELECTION_CACHE_TTL: Duration = Duration::from_secs(5);
 const CANDIDATE_SELECTION_CACHE_MAX_ENTRIES: usize = 4096;
+const CANDIDATE_SELECTION_CACHE_MAX_INFLIGHT: usize = 4096;
 #[cfg(not(test))]
 const CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT: Duration = Duration::from_secs(10);
 #[cfg(test)]
@@ -30,10 +31,10 @@ const CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from
 pub(super) struct CachedMinimalCandidateSelectionReadRepository {
     inner: Arc<dyn MinimalCandidateSelectionReadRepository>,
     entries: ExpiringMap<CandidateSelectionCacheKey, Vec<StoredMinimalCandidateSelectionRow>>,
-    inflight: Mutex<HashMap<CandidateSelectionCacheKey, u64>>,
-    inflight_notify: Notify,
-    next_inflight_token: AtomicU64,
+    inflight: Mutex<HashMap<CandidateSelectionCacheKey, Arc<InflightState>>>,
     epoch: AtomicU64,
+    mutation: Mutex<()>,
+    admission: Arc<Semaphore>,
 }
 
 impl CachedMinimalCandidateSelectionReadRepository {
@@ -42,9 +43,9 @@ impl CachedMinimalCandidateSelectionReadRepository {
             inner,
             entries: ExpiringMap::new(),
             inflight: Mutex::new(HashMap::new()),
-            inflight_notify: Notify::new(),
-            next_inflight_token: AtomicU64::new(1),
             epoch: AtomicU64::new(0),
+            mutation: Mutex::new(()),
+            admission: Arc::new(Semaphore::new(CANDIDATE_SELECTION_CACHE_MAX_INFLIGHT)),
         }
     }
 
@@ -61,18 +62,35 @@ impl CachedMinimalCandidateSelectionReadRepository {
             return Ok(rows);
         }
 
+        self.load_after_cache_miss(key, load).await
+    }
+
+    async fn load_after_cache_miss<F, Fut>(
+        &self,
+        key: CandidateSelectionCacheKey,
+        load: F,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>
+    where
+        F: Fn() -> Fut,
+        Fut: Future<Output = Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>>,
+    {
         loop {
-            let notified = self.inflight_notify.notified();
             match self.register_inflight(&key) {
-                InflightRegistration::Bypass => {
-                    return load_candidate_selection_rows_with_timeout(&key, load()).await;
+                InflightRegistration::Saturated => {
+                    return Err(DataLayerError::TimedOut(format!(
+                        "candidate selection cache admission saturated for {key:?}"
+                    )));
                 }
-                InflightRegistration::Follower => {
-                    if timeout(CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT, notified)
-                        .await
-                        .is_err()
+                InflightRegistration::Follower(state) => {
+                    match timeout(
+                        CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT,
+                        state.wait(),
+                    )
+                    .await
                     {
-                        self.expire_inflight(&key);
+                        Ok(Ok(())) => {}
+                        Ok(Err(error)) => return Err(error),
+                        Err(_) => self.expire_inflight(&key, &state),
                     }
                     if let Some(rows) = self.entries.get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
                     {
@@ -80,60 +98,107 @@ impl CachedMinimalCandidateSelectionReadRepository {
                     }
                     continue;
                 }
-                InflightRegistration::Leader(token) => {
-                    let mut guard = InflightGuard::new(self, key.clone(), token);
-                    let load_epoch = self.epoch.load(Ordering::Acquire);
+                InflightRegistration::Leader(mut guard) => {
+                    // A writer may have populated the cache after the first
+                    // miss but before this flight was registered.
+                    if let Some(rows) = self.entries.get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
+                    {
+                        return Ok(rows);
+                    }
+
                     let result = load_candidate_selection_rows_with_timeout(&key, load()).await;
-                    if let Ok(rows) = &result {
-                        if load_epoch == self.epoch.load(Ordering::Acquire) {
-                            self.entries.insert(
-                                key.clone(),
-                                rows.clone(),
-                                CANDIDATE_SELECTION_CACHE_TTL,
-                                CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
-                            );
-                        }
+                    match &result {
+                        Ok(rows) => guard.finish_loaded(rows.clone()),
+                        Err(error) => guard.finish(Some(error.clone())),
                     }
-                    guard.finish();
                     return result;
                 }
             }
         }
     }
 
-    fn register_inflight(&self, key: &CandidateSelectionCacheKey) -> InflightRegistration {
-        match self.inflight.lock() {
-            Ok(mut inflight) => {
-                if inflight.contains_key(key) {
-                    return InflightRegistration::Follower;
-                }
-                let token = self.next_inflight_token.fetch_add(1, Ordering::AcqRel);
-                inflight.insert(key.clone(), token);
-                InflightRegistration::Leader(token)
-            }
-            Err(_) => InflightRegistration::Bypass,
+    fn register_inflight(&self, key: &CandidateSelectionCacheKey) -> InflightRegistration<'_> {
+        let mut inflight = self
+            .inflight
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if let Some(state) = inflight.get(key) {
+            return InflightRegistration::Follower(Arc::clone(state));
         }
+        if inflight.len() >= CANDIDATE_SELECTION_CACHE_MAX_INFLIGHT {
+            return InflightRegistration::Saturated;
+        }
+        let Ok(admission) = Arc::clone(&self.admission).try_acquire_owned() else {
+            return InflightRegistration::Saturated;
+        };
+        let state = Arc::new(InflightState {
+            epoch: self.epoch.load(Ordering::Acquire),
+            notify: Notify::new(),
+            completed: AtomicBool::new(false),
+            error: Mutex::new(None),
+        });
+        inflight.insert(key.clone(), Arc::clone(&state));
+        InflightRegistration::Leader(InflightGuard::new(self, key.clone(), state, admission))
     }
 
-    fn finish_inflight(&self, key: &CandidateSelectionCacheKey, token: u64) {
-        let mut removed = false;
-        if let Ok(mut inflight) = self.inflight.lock() {
-            if inflight.get(key).copied() == Some(token) {
-                inflight.remove(key);
-                removed = true;
-            }
+    fn finish_inflight(
+        &self,
+        key: &CandidateSelectionCacheKey,
+        state: &Arc<InflightState>,
+        rows: Option<Vec<StoredMinimalCandidateSelectionRow>>,
+        admission: Option<OwnedSemaphorePermit>,
+    ) -> Option<Arc<InflightState>> {
+        let _mutation = self
+            .mutation
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let mut inflight = self
+            .inflight
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        drop(admission);
+        if !inflight
+            .get(key)
+            .is_some_and(|current| Arc::ptr_eq(current, state))
+            || self.epoch.load(Ordering::Acquire) != state.epoch
+        {
+            return None;
         }
-        if removed {
-            self.inflight_notify.notify_waiters();
+        if let Some(rows) = rows {
+            self.entries.insert(
+                key.clone(),
+                rows,
+                CANDIDATE_SELECTION_CACHE_TTL,
+                CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
+            );
         }
+        let removed = inflight.remove(key);
+        drop(inflight);
+        drop(_mutation);
+        removed
     }
 
-    fn expire_inflight(&self, key: &CandidateSelectionCacheKey) {
-        let mut removed = false;
-        if let Ok(mut inflight) = self.inflight.lock() {
-            removed = inflight.remove(key).is_some();
-        }
-        if removed {
+    fn expire_inflight(&self, key: &CandidateSelectionCacheKey, state: &Arc<InflightState>) {
+        let _mutation = self
+            .mutation
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let removed = {
+            let mut inflight = self
+                .inflight
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            if inflight
+                .get(key)
+                .is_some_and(|current| Arc::ptr_eq(current, state))
+            {
+                inflight.remove(key)
+            } else {
+                None
+            }
+        };
+        drop(_mutation);
+        if let Some(state) = removed {
             warn!(
                 event_name = "candidate_selection_cache_inflight_expired",
                 log_type = "ops",
@@ -141,64 +206,143 @@ impl CachedMinimalCandidateSelectionReadRepository {
                 wait_timeout_ms = CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT.as_millis() as u64,
                 "gateway candidate selection cache expired stale inflight load"
             );
-            self.inflight_notify.notify_waiters();
+            state.complete(None);
         }
     }
 
     fn clear(&self) {
+        let _mutation = self
+            .mutation
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
         self.epoch.fetch_add(1, Ordering::AcqRel);
         self.entries.clear();
-        let mut cleared_inflight = false;
-        if let Ok(mut inflight) = self.inflight.lock() {
-            cleared_inflight = !inflight.is_empty();
-            inflight.clear();
-        }
-        if cleared_inflight {
+        let states = self
+            .inflight
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .drain()
+            .map(|(_, state)| state)
+            .collect::<Vec<_>>();
+        if !states.is_empty() {
             warn!(
                 event_name = "candidate_selection_cache_inflight_cleared",
                 log_type = "ops",
                 "gateway candidate selection cache cleared in-flight loads"
             );
-            self.inflight_notify.notify_waiters();
+            for state in states {
+                state.complete(None);
+            }
         }
     }
 }
 
-enum InflightRegistration {
-    Leader(u64),
-    Follower,
-    Bypass,
+struct InflightState {
+    epoch: u64,
+    notify: Notify,
+    completed: AtomicBool,
+    error: Mutex<Option<DataLayerError>>,
+}
+
+impl InflightState {
+    fn complete(&self, error: Option<DataLayerError>) {
+        if let Some(error) = error {
+            if let Ok(mut current) = self.error.lock() {
+                *current = Some(error);
+            }
+        }
+        if !self.completed.swap(true, Ordering::AcqRel) {
+            self.notify.notify_waiters();
+        }
+    }
+
+    async fn wait(&self) -> Result<(), DataLayerError> {
+        loop {
+            if self.completed.load(Ordering::Acquire) {
+                return self
+                    .error
+                    .lock()
+                    .ok()
+                    .and_then(|error| error.clone())
+                    .map_or(Ok(()), Err);
+            }
+
+            // Register before checking completion a second time. Creating a
+            // Notified future alone is insufficient because notify_waiters()
+            // can otherwise run before the future's first poll.
+            let mut notified = Box::pin(self.notify.notified());
+            notified.as_mut().enable();
+            if self.completed.load(Ordering::Acquire) {
+                return self
+                    .error
+                    .lock()
+                    .ok()
+                    .and_then(|error| error.clone())
+                    .map_or(Ok(()), Err);
+            }
+            notified.await;
+        }
+    }
+}
+
+enum InflightRegistration<'a> {
+    Leader(InflightGuard<'a>),
+    Follower(Arc<InflightState>),
+    Saturated,
 }
 
 struct InflightGuard<'a> {
     cache: &'a CachedMinimalCandidateSelectionReadRepository,
     key: Option<CandidateSelectionCacheKey>,
-    token: u64,
+    state: Arc<InflightState>,
+    admission: Option<OwnedSemaphorePermit>,
 }
 
 impl<'a> InflightGuard<'a> {
     fn new(
         cache: &'a CachedMinimalCandidateSelectionReadRepository,
         key: CandidateSelectionCacheKey,
-        token: u64,
+        state: Arc<InflightState>,
+        admission: OwnedSemaphorePermit,
     ) -> Self {
         Self {
             cache,
             key: Some(key),
-            token,
+            state,
+            admission: Some(admission),
         }
     }
 
-    fn finish(&mut self) {
-        if let Some(key) = self.key.take() {
-            self.cache.finish_inflight(&key, self.token);
+    fn epoch(&self) -> u64 {
+        self.state.epoch
+    }
+
+    fn finish_loaded(&mut self, rows: Vec<StoredMinimalCandidateSelectionRow>) {
+        let removed = self.key.take().and_then(|key| {
+            self.cache
+                .finish_inflight(&key, &self.state, Some(rows), self.admission.take())
+        });
+        self.admission.take();
+        if let Some(removed) = removed {
+            removed.complete(None);
+        }
+    }
+
+    fn finish(&mut self, error: Option<DataLayerError>) {
+        let removed = self.key.take().and_then(|key| {
+            self.cache
+                .finish_inflight(&key, &self.state, None, self.admission.take())
+        });
+        self.admission.take();
+        if let Some(removed) = removed {
+            removed.complete(error);
         }
     }
 }
 
 impl Drop for InflightGuard<'_> {
     fn drop(&mut self) {
-        self.finish();
+        self.finish(None);
     }
 }
 
@@ -573,6 +717,202 @@ mod tests {
         assert_eq!(inner.calls(), 1);
     }
 
+    #[tokio::test]
+    async fn candidate_selection_cache_only_notifies_matching_inflight_key() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key_a = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let key_b = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "anthropic:messages".to_string(),
+        };
+
+        let mut leader_a = match cache.register_inflight(&key_a) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first key registration should lead"),
+        };
+        let mut leader_b = match cache.register_inflight(&key_b) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("second key registration should lead independently"),
+        };
+        let state_a = match cache.register_inflight(&key_a) {
+            InflightRegistration::Follower(state) => state,
+            _ => panic!("duplicate key registration should follow"),
+        };
+
+        leader_b.finish(None);
+        assert!(
+            tokio::time::timeout(Duration::from_millis(10), state_a.wait())
+                .await
+                .is_err(),
+            "completing another key must not wake this follower"
+        );
+        leader_a.finish(None);
+        tokio::time::timeout(Duration::from_millis(100), state_a.wait())
+            .await
+            .expect("completing the matching key must wake its follower")
+            .expect("successful completion should not publish an error");
+        assert!(cache.inflight.lock().unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn candidate_selection_cache_follower_observes_completion_before_first_poll() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+
+        let mut leader = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+        let state = match cache.register_inflight(&key) {
+            InflightRegistration::Follower(state) => state,
+            _ => panic!("second registration should follow"),
+        };
+
+        // Complete before wait() is constructed or polled. notify_waiters()
+        // alone would lose this notification and wait for the full timeout.
+        leader.finish(None);
+        tokio::time::timeout(Duration::from_millis(100), state.wait())
+            .await
+            .expect("completed follower must not miss the broadcast")
+            .expect("successful completion should not publish an error");
+    }
+
+    #[tokio::test]
+    async fn candidate_selection_cache_shares_leader_failure() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let mut leader = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+        let state = match cache.register_inflight(&key) {
+            InflightRegistration::Follower(state) => state,
+            _ => panic!("second registration should follow"),
+        };
+
+        leader.finish(Some(DataLayerError::Sql(
+            "forced candidate cache load failure".to_string(),
+        )));
+        let error = tokio::time::timeout(Duration::from_millis(100), state.wait())
+            .await
+            .expect("failed load should release its follower")
+            .expect_err("follower should observe the leader failure");
+        assert_eq!(
+            error.to_string(),
+            "sql error: forced candidate cache load failure"
+        );
+        assert!(cache.inflight.lock().unwrap().is_empty());
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+    async fn candidate_selection_cache_broadcasts_to_all_same_key_followers() {
+        const FOLLOWERS: usize = 64;
+
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let mut leader = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+
+        let mut tasks = Vec::with_capacity(FOLLOWERS);
+        for _ in 0..FOLLOWERS {
+            let state = match cache.register_inflight(&key) {
+                InflightRegistration::Follower(state) => state,
+                _ => panic!("same-key registration should follow"),
+            };
+            tasks.push(tokio::spawn(async move { state.wait().await }));
+        }
+
+        tokio::task::yield_now().await;
+        leader.finish(None);
+        for task in tasks {
+            tokio::time::timeout(Duration::from_millis(250), task)
+                .await
+                .expect("all same-key followers should receive completion")
+                .expect("follower task should finish")
+                .expect("successful completion should not publish an error");
+        }
+    }
+
+    #[tokio::test]
+    async fn candidate_selection_cache_cancelled_guard_wakes_registered_follower() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let guard = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+        let state = match cache.register_inflight(&key) {
+            InflightRegistration::Follower(state) => state,
+            _ => panic!("second registration should follow"),
+        };
+
+        drop(guard);
+        tokio::time::timeout(Duration::from_millis(100), state.wait())
+            .await
+            .expect("leader cancellation must wake an existing follower")
+            .expect("leader cancellation should allow a retry");
+        assert!(cache.inflight.lock().unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn candidate_selection_cache_clear_wakes_follower_and_old_guard_keeps_new_flight() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let old_guard = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+        let old_state = match cache.register_inflight(&key) {
+            InflightRegistration::Follower(state) => state,
+            _ => panic!("second registration should follow"),
+        };
+        let old_epoch = old_guard.epoch();
+
+        cache.clear();
+        assert!(cache.epoch.load(Ordering::Acquire) > old_epoch);
+        let mut new_guard = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("registration after clear should lead"),
+        };
+        assert_ne!(old_guard.epoch(), new_guard.epoch());
+
+        // Dropping the invalidated leader's RAII guard must not remove the
+        // new flight created for the same key after clear().
+        drop(old_guard);
+        assert!(cache
+            .inflight
+            .lock()
+            .unwrap()
+            .get(&key)
+            .is_some_and(|current| Arc::ptr_eq(current, &new_guard.state)));
+        tokio::time::timeout(Duration::from_millis(100), old_state.wait())
+            .await
+            .expect("clear must wake followers of the invalidated flight")
+            .expect("clear should allow a retry");
+
+        new_guard.finish(None);
+        assert!(cache.inflight.lock().unwrap().is_empty());
+    }
+
     #[tokio::test]
     async fn candidate_selection_cache_clear_invalidates_entries() {
         let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
@@ -587,6 +927,123 @@ mod tests {
         assert_eq!(inner.calls(), 2);
     }
 
+    #[test]
+    fn candidate_selection_cache_rejects_publication_from_pre_clear_flight() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let mut stale_leader = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+
+        cache.clear();
+        stale_leader.finish_loaded(Vec::new());
+
+        assert!(cache
+            .entries
+            .get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
+            .is_none());
+    }
+
+    #[test]
+    fn candidate_selection_cache_clear_keeps_active_load_admission_bounded() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let mut cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        cache.admission = Arc::new(Semaphore::new(2));
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let mut detached_leaders = Vec::new();
+
+        for _ in 0..2 {
+            let leader = match cache.register_inflight(&key) {
+                InflightRegistration::Leader(guard) => guard,
+                _ => panic!("load below the hard active limit should lead"),
+            };
+            cache.clear();
+            detached_leaders.push(leader);
+        }
+
+        assert_eq!(cache.admission.available_permits(), 0);
+        assert!(matches!(
+            cache.register_inflight(&key),
+            InflightRegistration::Saturated
+        ));
+        drop(detached_leaders);
+        assert_eq!(cache.admission.available_permits(), 2);
+        assert!(matches!(
+            cache.register_inflight(&key),
+            InflightRegistration::Leader(_)
+        ));
+    }
+
+    #[test]
+    fn candidate_selection_cache_expired_leader_cannot_publish_over_replacement() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let mut old_leader = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("first registration should lead"),
+        };
+        let old_state = Arc::clone(&old_leader.state);
+        cache.expire_inflight(&key, &old_state);
+        let mut replacement = match cache.register_inflight(&key) {
+            InflightRegistration::Leader(guard) => guard,
+            _ => panic!("expiration should allow a replacement leader"),
+        };
+        assert_eq!(old_leader.epoch(), replacement.epoch());
+
+        old_leader.finish_loaded(vec![sample_row("stale-key", 1)]);
+        assert!(cache
+            .entries
+            .get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
+            .is_none());
+
+        replacement.finish_loaded(vec![sample_row("fresh-key", 2)]);
+        let cached = cache
+            .entries
+            .get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
+            .expect("replacement should publish");
+        assert_eq!(cached[0].key_id, "fresh-key");
+    }
+
+    #[tokio::test]
+    async fn candidate_selection_cache_rechecks_fresh_entry_after_leader_registration() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner);
+        let key = CandidateSelectionCacheKey::ApiFormat {
+            api_format: "openai:chat".to_string(),
+        };
+        let expected = vec![sample_row("cached-key", 1)];
+        cache.entries.insert(
+            key.clone(),
+            expected.clone(),
+            CANDIDATE_SELECTION_CACHE_TTL,
+            CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
+        );
+        let loads = AtomicUsize::new(0);
+
+        // Exercise the post-initial-miss path directly to model a concurrent
+        // writer filling the cache immediately before flight registration.
+        let rows = cache
+            .load_after_cache_miss(key, || async {
+                loads.fetch_add(1, Ordering::SeqCst);
+                Ok(Vec::new())
+            })
+            .await
+            .expect("fresh entry should satisfy the lookup");
+
+        assert_eq!(rows, expected);
+        assert_eq!(loads.load(Ordering::SeqCst), 0);
+        assert!(cache.inflight.lock().unwrap().is_empty());
+    }
+
     #[tokio::test]
     async fn candidate_selection_cache_releases_inflight_when_leader_is_cancelled() {
         let inner = Arc::new(FirstLoadPendingThenFastRepository::new());