perf(router): probe worker availability with an Fx set

lib/kv-router/src/scheduling/filter.rs

@@ -4,7 +4,7 @@
 use std::collections::{HashMap, HashSet};
 
 use super::policy_config::FullRankCounts;
-use super::types::KvSchedulerError;
+use super::types::{KvSchedulerError, WorkerIdSet};
 use crate::protocols::{
     DpRank, RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank,
 };
@@ -46,7 +46,7 @@ struct RejectCheck<'a> {
 pub struct RoutingEligibility<'a> {
     allowed_worker_ids: Option<&'a HashSet<WorkerId>>,
     overloaded_worker_ids: Option<&'a HashSet<WorkerId>>,
-    available_worker_ids: Option<&'a HashSet<WorkerId>>,
+    available_worker_ids: Option<&'a WorkerIdSet>,
     pinned_worker: Option<WorkerWithDpRank>,
     affinity_target: Option<WorkerAffinityTarget>,
     routing_constraints: &'a RoutingConstraints,
@@ -103,10 +103,7 @@ impl<'a> RoutingEligibility<'a> {
     /// Attach hard availability. Unlike transient overload, unavailability is
     /// enforced on every path, including affinity-derived pins.
     #[inline]
-    pub fn with_available_workers(
-        mut self,
-        available_worker_ids: Option<&'a HashSet<WorkerId>>,
-    ) -> Self {
+    pub fn with_available_workers(mut self, available_worker_ids: Option<&'a WorkerIdSet>) -> Self {
         self.available_worker_ids = available_worker_ids;
         self
     }
@@ -290,31 +287,29 @@ impl<'a> RoutingEligibility<'a> {
             return any_worker_rank(target.worker_id, config, &mut predicate);
         }
 
-        let inclusion = match (self.allowed_worker_ids, self.available_worker_ids) {
-            (Some(allowed), Some(available)) => {
-                Some(if allowed.capacity() <= available.capacity() {
-                    allowed
-                } else {
-                    available
-                })
-            }
-            (allowed, available) => allowed.or(available),
-        };
-        // Hash iteration scans capacity, not length. The half-capacity cutoff
-        // is a conservative heuristic for the extra worker lookups.
-        if let Some(inclusion) = inclusion.filter(|ids| ids.capacity() < workers.capacity() / 2) {
-            for &worker_id in inclusion {
-                let Some(config) = workers.get(&worker_id) else {
-                    continue;
-                };
-                if !self.allows_worker(worker_id, config) {
-                    continue;
-                }
-                if any_worker_rank(worker_id, config, &mut predicate) {
-                    return true;
-                }
-            }
-            return false;
+        // Walk the smaller inclusion set when it is much smaller than the workers. Hash
+        // iteration scans capacity, not length. The half-capacity cutoff is a
+        // conservative heuristic for the extra worker lookups.
+        let allowed = self
+            .allowed_worker_ids
+            .map_or(usize::MAX, |ids| ids.capacity());
+        let available = self
+            .available_worker_ids
+            .map_or(usize::MAX, |ids| ids.capacity());
+        if allowed.min(available) < workers.capacity() / 2 {
+            return if allowed <= available {
+                self.any_listed_worker_rank(
+                    workers,
+                    self.allowed_worker_ids.into_iter().flatten(),
+                    &mut predicate,
+                )
+            } else {
+                self.any_listed_worker_rank(
+                    workers,
+                    self.available_worker_ids.into_iter().flatten(),
+                    &mut predicate,
+                )
+            };
         }
 
         for (&worker_id, config) in workers {
@@ -330,6 +325,30 @@ impl<'a> RoutingEligibility<'a> {
         false
     }
 
+    /// `any_eligible_worker_rank` over the listed workers only.
+    fn any_listed_worker_rank<'i, C, F>(
+        &self,
+        workers: &HashMap<WorkerId, C>,
+        worker_ids: impl IntoIterator<Item = &'i WorkerId>,
+        predicate: &mut F,
+    ) -> bool
+    where
+        C: WorkerConfigLike,
+        F: FnMut(WorkerWithDpRank, &C) -> bool,
+    {
+        for &worker_id in worker_ids {
+            let Some(config) = workers.get(&worker_id) else {
+                continue;
+            };
+            if self.allows_worker(worker_id, config)
+                && any_worker_rank(worker_id, config, predicate)
+            {
+                return true;
+            }
+        }
+        false
+    }
+
     pub(crate) fn affinity_target_is_eligible<C: WorkerConfigLike>(
         &self,
         workers: &HashMap<WorkerId, C>,
@@ -569,7 +588,7 @@ mod tests {
         let workers = workers();
         let config = workers.get(&7).unwrap();
         let overloaded = HashSet::from([7]);
-        let available = HashSet::from([8]);
+        let available = WorkerIdSet::from_iter([8]);
         let constraints = RoutingConstraints::default();
         let worker = WorkerWithDpRank::new(7, 3);
         let eligibility =
@@ -582,7 +601,7 @@ mod tests {
             Some(WorkerEligibilityError::WorkerNotRoutable { worker_id: 7 })
         );
 
-        let empty = HashSet::new();
+        let empty = WorkerIdSet::default();
         let authoritative_empty = RoutingEligibility::new(None, None, None, &constraints)
             .with_available_workers(Some(&empty));
         assert!(!authoritative_empty.allows_worker(7, config));
@@ -767,7 +786,7 @@ mod tests {
         let mut workers: HashMap<_, _> = (0..64).map(|id| (id, config.clone())).collect();
         workers.get_mut(&1).unwrap().taints.clear();
         let allowed = HashSet::from([0, 1, 2, 3, 1000]);
-        let available = HashSet::from([0, 1, 2, 4, 1001]);
+        let available = WorkerIdSet::from_iter([0, 1, 2, 4, 1001]);
         let overloaded = HashSet::from([2]);
         let constraints = RoutingConstraints {
             required_taints: HashSet::from(["zone-a".to_string()]),
@@ -798,7 +817,7 @@ mod tests {
             ..Default::default()
         };
         let workers = (0..64).map(|id| (id, config.clone())).collect();
-        let available = HashSet::from([7, 1000]);
+        let available = WorkerIdSet::from_iter([7, 1000]);
         let constraints = RoutingConstraints::default();
         let eligibility = RoutingEligibility::new(None, None, None, &constraints)
             .with_available_workers(Some(&available));

lib/kv-router/src/scheduling/queue.rs

@@ -33,9 +33,10 @@ use super::selector::{
     CandidatePreparer, PreparedCandidates, WorkerSelectionInput, WorkerSelector, worker_input_rows,
 };
 use super::types::{
-    AdvisorySchedulingResponse, AdvisoryWorkerLoad, AttemptId, KvSchedulerError,
-    NonMaxOverlapSelection, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider,
-    SchedulingContext, SchedulingRequest, SchedulingResponse, WorkerAvailabilityProvider,
+    AdvisorySchedulingResponse, AdvisoryWorkerLoad, AttemptId, AvailableWorkerIds,
+    KvSchedulerError, NonMaxOverlapSelection, NonMaxOverlapSelectionObserver,
+    OverloadedWorkerProvider, SchedulingContext, SchedulingRequest, SchedulingResponse,
+    WorkerAvailabilityProvider, WorkerIdSet,
 };
 use crate::plugins::request_classifier::ClassifyRequest;
 use crate::plugins::worker_selection::{
@@ -103,7 +104,7 @@ struct PreparedAdmission {
     candidates: PreparedCandidates,
     /// Configuration generation, loaded before the configuration was read.
     workers_generation: u64,
-    available_worker_ids: Option<Arc<HashSet<WorkerId>>>,
+    available_worker_ids: Option<Arc<WorkerIdSet>>,
     overloaded_worker_ids: Option<Arc<HashSet<WorkerId>>>,
 }
 
@@ -112,7 +113,7 @@ impl PreparedAdmission {
     fn into_current(
         self,
         workers_generation: u64,
-        available_worker_ids: Option<&Arc<HashSet<WorkerId>>>,
+        available_worker_ids: Option<&Arc<WorkerIdSet>>,
         overloaded_worker_ids: Option<&Arc<HashSet<WorkerId>>>,
     ) -> Option<PreparedCandidates> {
         // Availability and overload are each one shared snapshot until they change,
@@ -130,7 +131,7 @@ fn selection_eligibility<'a, C: WorkerConfigLike>(
     request: &'a SchedulingRequest,
     workers: &HashMap<WorkerId, C>,
     overloaded_worker_ids: Option<&'a HashSet<WorkerId>>,
-    available_worker_ids: Option<&'a HashSet<WorkerId>>,
+    available_worker_ids: Option<&'a WorkerIdSet>,
     full_ranks: Option<&'a FullRankCounts>,
     exclusive_affinity: bool,
 ) -> RoutingEligibility<'a> {
@@ -599,7 +600,7 @@ struct SchedulerQueueActor<
     overlap_scores_refresh: Option<Arc<RF>>,
     overlap_refresh_after: Option<Duration>,
     overloaded_worker_provider: Option<OverloadedWorkerProvider>,
-    available_worker_provider: Option<WorkerAvailabilityProvider>,
+    available_worker_ids: Option<Arc<AvailableWorkerIds>>,
     non_max_overlap_selection_observer: Arc<OnceLock<NonMaxOverlapSelectionObserver>>,
     /// Per-candidate cost capture for admitted selections, installed once by the host.
     selection_trace: Arc<OnceLock<SelectionTraceSpec>>,
@@ -635,7 +636,7 @@ pub struct SchedulerQueue<
     block_size: u32,
     profile: Arc<PolicyProfile>,
     start_time: Instant,
-    available_worker_provider: Option<WorkerAvailabilityProvider>,
+    available_worker_ids: Option<Arc<AvailableWorkerIds>>,
     overloaded_worker_provider: Option<OverloadedWorkerProvider>,
     supports_overlap_refresh: bool,
     non_max_overlap_selection_observer: Arc<OnceLock<NonMaxOverlapSelectionObserver>>,
@@ -770,6 +771,8 @@ impl<
             })
         };
         let workers_generation = Arc::new(AtomicU64::new(0));
+        let available_worker_ids =
+            available_worker_provider.map(|provider| Arc::new(AvailableWorkerIds::new(provider)));
         let candidate_preparer = Arc::new(ArcSwapOption::from(
             selector.candidate_preparer().map(Arc::new),
         ));
@@ -797,7 +800,7 @@ impl<
             overlap_scores_refresh,
             overlap_refresh_after,
             overloaded_worker_provider: overloaded_worker_provider.clone(),
-            available_worker_provider: available_worker_provider.clone(),
+            available_worker_ids: available_worker_ids.clone(),
             non_max_overlap_selection_observer: Arc::clone(&non_max_overlap_selection_observer),
             selection_trace: Arc::clone(&selection_trace),
             reject_thresholds: reject_thresholds.clone(),
@@ -819,7 +822,7 @@ impl<
             block_size,
             profile,
             start_time,
-            available_worker_provider,
+            available_worker_ids,
             overloaded_worker_provider,
             supports_overlap_refresh: overlap_refresh_after.is_some(),
             non_max_overlap_selection_observer,
@@ -970,9 +973,9 @@ impl<
         worker_inputs: WorkerInputs,
     ) -> (ClassifyRequest, WorkerInputRows) {
         let available = self
-            .available_worker_provider
+            .available_worker_ids
             .as_ref()
-            .and_then(|provider| provider(request));
+            .and_then(|ids| ids.get(request));
         let workers = self.workers_with_configs.borrow();
         let cached_tokens = SchedulingContext::new(request, &workers)
             .with_available_workers(available.as_deref())
@@ -1039,9 +1042,9 @@ impl<
         ingress_at: Instant,
     ) -> QueueMetadata {
         let available = self
-            .available_worker_provider
+            .available_worker_ids
             .as_ref()
-            .and_then(|provider| provider(request));
+            .and_then(|ids| ids.get(request));
         let workers = self.workers_with_configs.borrow();
         let snapshot = QueueSnapshot::new(
             request.isl_tokens,
@@ -1251,9 +1254,9 @@ impl<
         // advances the generation before the actor selects with these rows.
         let workers_generation = self.workers_generation.load(AtomicOrdering::Acquire);
         let available_worker_ids = self
-            .available_worker_provider
+            .available_worker_ids
             .as_ref()
-            .and_then(|provider| provider(request));
+            .and_then(|ids| ids.get(request));
         let overloaded_worker_ids = self
             .overloaded_worker_provider
             .as_ref()
@@ -1469,9 +1472,9 @@ impl<
         // limit-zero rejection the queued path always has.
         if class.never_queues() {
             let available = self
-                .available_worker_provider
+                .available_worker_ids
                 .as_ref()
-                .and_then(|provider| provider(&request));
+                .and_then(|ids| ids.get(&request));
             let should_queue = class.queueing_enabled()
                 && (self.pending.has_backlog(class_index) || {
                     let configs = self.workers_with_configs.borrow();
@@ -1651,13 +1654,13 @@ impl<
     fn has_dispatchable_ready_head(&self) -> bool {
         let decay_now = Instant::now();
         let slots = &self.slots;
-        let provider = self.available_worker_provider.as_ref();
+        let available_ids = self.available_worker_ids.as_ref();
         let workers = &self.workers_with_configs;
         self.pending.any_ready_head(|_, class, queued| {
             if !class.queueing_enabled() {
                 return true;
             }
-            let available = provider.and_then(|provider| provider(&queued.request));
+            let available = available_ids.and_then(|ids| ids.get(&queued.request));
             let configs = workers.borrow();
             !Self::all_eligible_workers_full(
                 &slots.worker_load_reader(decay_now),
@@ -1747,13 +1750,13 @@ impl<
             // two reads, as they can right after the gate.
             let popped = {
                 let slots = &self.slots;
-                let provider = self.available_worker_provider.as_ref();
+                let available_ids = self.available_worker_ids.as_ref();
                 let workers = &self.workers_with_configs;
                 self.pending.pop_next(|_, class, queued| {
                     if !class.queueing_enabled() {
                         return true;
                     }
-                    let available = provider.and_then(|provider| provider(&queued.request));
+                    let available = available_ids.and_then(|ids| ids.get(&queued.request));
                     let configs = workers.borrow();
                     // TODO: This preserves head-of-line blocking within each policy
                     // class. A blocked constrained head can stall later entries in
@@ -1902,9 +1905,9 @@ impl<
         prepared: Option<Box<PreparedAdmission>>,
     ) -> Result<SelectedWorkerForRequest, KvSchedulerError> {
         let available_worker_ids = self
-            .available_worker_provider
+            .available_worker_ids
             .as_ref()
-            .and_then(|provider| provider(request));
+            .and_then(|ids| ids.get(request));
 
         {
             let workers = self.workers_with_configs.borrow();
@@ -3878,7 +3881,7 @@ policy_classes:
             overlap_scores_refresh: None::<Arc<NoopOverlapScoresRefresh>>,
             overlap_refresh_after: None,
             overloaded_worker_provider: None,
-            available_worker_provider: None,
+            available_worker_ids: None,
             non_max_overlap_selection_observer: Arc::clone(
                 &queue.non_max_overlap_selection_observer,
             ),