dynamo_runtime/component/
client.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::{
6    collections::{HashMap, HashSet},
7    sync::{Arc, Mutex as StdMutex},
8    time::Duration,
9};
10
11use anyhow::Result;
12use arc_swap::ArcSwap;
13use dashmap::DashMap;
14use futures::StreamExt;
15use rand::Rng;
16
17use crate::component::{Endpoint, Instance};
18use crate::discovery::{DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId};
19use crate::traits::DistributedRuntimeProvider;
20
21/// Shared occupancy state for routing modes that track per-worker in-flight requests.
22#[derive(Debug, Default)]
23pub(crate) struct RoutingOccupancyState {
24    counts: DashMap<u64, AtomicU64>,
25    exact_selection_lock: tokio::sync::Mutex<()>,
26}
27
28impl RoutingOccupancyState {
29    pub(crate) fn increment(&self, instance_id: u64) {
30        self.counts
31            .entry(instance_id)
32            .or_insert_with(|| AtomicU64::new(0))
33            .fetch_add(1, Ordering::Relaxed);
34    }
35
36    pub(crate) async fn select_exact_min(&self, instance_ids: &[u64]) -> Option<u64> {
37        instance_ids
38            .iter()
39            .min_by_key(|&&id| self.load(id))
40            .copied()
41    }
42
43    pub(crate) async fn select_exact_min_and_increment(&self, instance_ids: &[u64]) -> Option<u64> {
44        let _guard = self.exact_selection_lock.lock().await;
45
46        let mut min_load = u64::MAX;
47        let mut selected = None;
48        let mut tie_count = 0usize;
49        let mut rng = rand::rng();
50        for &id in instance_ids {
51            let load = self.load(id);
52            if load < min_load {
53                min_load = load;
54                selected = Some(id);
55                tie_count = 1;
56                continue;
57            }
58
59            if load == min_load {
60                tie_count += 1;
61                // Reservoir sampling keeps tied minima uniform without allocating in this locked hot path.
62                if rng.random_range(0..tie_count) == 0 {
63                    selected = Some(id);
64                }
65            }
66        }
67
68        let id = selected?;
69        self.increment(id);
70        Some(id)
71    }
72
73    /// Least-loaded selection without the increment. Same tie-break policy as
74    /// [`Self::select_exact_min_and_increment`] so peek and select share a
75    /// distribution.
76    pub(crate) fn peek_min(&self, instance_ids: &[u64]) -> Option<u64> {
77        let mut min_load = u64::MAX;
78        let mut selected = None;
79        let mut tie_count = 0usize;
80        let mut rng = rand::rng();
81        for &id in instance_ids {
82            let load = self.load(id);
83            if load < min_load {
84                min_load = load;
85                selected = Some(id);
86                tie_count = 1;
87                continue;
88            }
89
90            if load == min_load {
91                tie_count += 1;
92                // Reservoir sampling keeps tied minima uniform; matches select_exact_min_and_increment.
93                if rng.random_range(0..tie_count) == 0 {
94                    selected = Some(id);
95                }
96            }
97        }
98
99        selected
100    }
101
102    pub(crate) fn decrement(&self, instance_id: u64) {
103        if let Some(count) = self.counts.get(&instance_id) {
104            let _ = count.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
105                Some(current.saturating_sub(1))
106            });
107        }
108    }
109
110    pub(crate) fn load(&self, instance_id: u64) -> u64 {
111        self.counts
112            .get(&instance_id)
113            .map(|c| c.load(Ordering::Relaxed))
114            .unwrap_or(0)
115    }
116
117    pub(crate) fn retain(&self, instance_ids: &[u64]) {
118        let live: HashSet<u64> = instance_ids.iter().copied().collect();
119        self.counts.retain(|id, _| live.contains(id));
120    }
121}
122
123/// Get or create the shared routing occupancy state for an endpoint.
124pub(crate) async fn get_or_create_routing_occupancy_state(
125    endpoint: &Endpoint,
126) -> Arc<RoutingOccupancyState> {
127    let drt = endpoint.drt();
128    let registry = drt.routing_occupancy_states();
129    let mut registry = registry.lock().await;
130
131    if let Some(weak) = registry.get(endpoint) {
132        if let Some(state) = weak.upgrade() {
133            return state;
134        } else {
135            registry.remove(endpoint);
136        }
137    }
138
139    let state = Arc::new(RoutingOccupancyState::default());
140    registry.insert(endpoint.clone(), Arc::downgrade(&state));
141    state
142}
143
144/// Default interval for periodic reconciliation of instance_avail with instance_source
145const DEFAULT_RECONCILE_INTERVAL: Duration = Duration::from_secs(5);
146
147/// Shared endpoint discovery state for a single endpoint query.
148///
149/// This wraps both the coalesced instance snapshot used for routing decisions
150/// and a raw, lossless per-subscriber event feed used by the response-stream
151/// cancellation watcher. Both outputs are driven by a single underlying
152/// discovery `list_and_watch` task so clients do not multiply control-plane
153/// watches.
154#[derive(Debug)]
155pub(crate) struct EndpointDiscoverySource {
156    instance_source: tokio::sync::watch::Receiver<Vec<Instance>>,
157    event_subscribers: StdMutex<Vec<tokio::sync::mpsc::UnboundedSender<DiscoveryEvent>>>,
158}
159
160impl EndpointDiscoverySource {
161    fn new(instance_source: tokio::sync::watch::Receiver<Vec<Instance>>) -> Self {
162        Self {
163            instance_source,
164            event_subscribers: StdMutex::new(Vec::new()),
165        }
166    }
167
168    fn instance_receiver(&self) -> tokio::sync::watch::Receiver<Vec<Instance>> {
169        self.instance_source.clone()
170    }
171
172    fn subscribe_events(&self) -> tokio::sync::mpsc::UnboundedReceiver<DiscoveryEvent> {
173        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
174        self.event_subscribers.lock().unwrap().push(tx);
175        rx
176    }
177
178    fn broadcast_event(&self, event: &DiscoveryEvent) {
179        let subscribers = &mut *self.event_subscribers.lock().unwrap();
180        subscribers.retain(|tx| tx.send(event.clone()).is_ok());
181    }
182}
183
184#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
185pub struct RoutingInstanceCounts {
186    pub discovered: usize,
187    pub routable: usize,
188    pub overloaded: usize,
189    /// IDs not currently reported overloaded, derived from `discovered - overloaded`.
190    pub free: usize,
191}
192
193#[derive(Clone, Debug)]
194pub(crate) struct RoutingInstances {
195    discovered_ids: Vec<u64>,
196    routable_ids: Vec<u64>,
197    overloaded_ids: HashSet<u64>,
198    free_ids: Vec<u64>,
199}
200
201impl RoutingInstances {
202    fn new(discovered_ids: Vec<u64>) -> Self {
203        Self::from_parts(discovered_ids.clone(), discovered_ids, HashSet::new())
204    }
205
206    fn from_parts(
207        discovered_ids: Vec<u64>,
208        routable_ids: Vec<u64>,
209        overloaded_ids: HashSet<u64>,
210    ) -> Self {
211        let free_ids = Self::derive_free_ids(&routable_ids, &overloaded_ids);
212        Self {
213            discovered_ids,
214            routable_ids,
215            overloaded_ids,
216            free_ids,
217        }
218    }
219
220    pub(crate) fn discovered_ids(&self) -> &[u64] {
221        &self.discovered_ids
222    }
223
224    pub(crate) fn routable_ids(&self) -> &[u64] {
225        &self.routable_ids
226    }
227
228    pub(crate) fn free_ids(&self) -> &[u64] {
229        &self.free_ids
230    }
231
232    pub(crate) fn counts(&self) -> RoutingInstanceCounts {
233        RoutingInstanceCounts {
234            discovered: self.discovered_ids.len(),
235            routable: self.routable_ids.len(),
236            overloaded: self.overloaded_ids.len(),
237            free: self.free_ids.len(),
238        }
239    }
240
241    pub(crate) fn is_overloaded(&self, instance_id: u64) -> bool {
242        self.overloaded_ids.contains(&instance_id)
243    }
244
245    fn overloaded_ids(&self) -> Option<HashSet<u64>> {
246        if self.overloaded_ids.is_empty() {
247            return None;
248        }
249
250        Some(self.overloaded_ids.clone())
251    }
252
253    fn reconcile_discovered(&self, discovered_ids: Vec<u64>) -> Self {
254        let old_discovered_ids = self.discovered_ids.iter().copied().collect::<HashSet<_>>();
255        let new_discovered_ids = discovered_ids.iter().copied().collect::<HashSet<_>>();
256        let mut overloaded_ids = self.overloaded_ids.clone();
257        overloaded_ids
258            .retain(|id| !old_discovered_ids.contains(id) || new_discovered_ids.contains(id));
259
260        Self::from_parts(discovered_ids.clone(), discovered_ids, overloaded_ids)
261    }
262
263    fn report_instance_down(&self, instance_id: u64) -> Self {
264        let routable_ids: Vec<u64> = self
265            .routable_ids
266            .iter()
267            .copied()
268            .filter(|id| *id != instance_id)
269            .collect();
270
271        Self::from_parts(
272            self.discovered_ids.clone(),
273            routable_ids,
274            self.overloaded_ids.clone(),
275        )
276    }
277
278    #[cfg(test)]
279    fn override_routable_ids(&self, routable_ids: Vec<u64>) -> Self {
280        // Route through from_parts so `free_ids` is recomputed from the new
281        // routable set instead of carrying the stale value forward.
282        Self::from_parts(
283            self.discovered_ids.clone(),
284            routable_ids,
285            self.overloaded_ids.clone(),
286        )
287    }
288
289    fn set_overloaded(&self, overloaded_ids: HashSet<u64>) -> Self {
290        Self::from_parts(
291            self.discovered_ids.clone(),
292            self.routable_ids.clone(),
293            overloaded_ids,
294        )
295    }
296
297    /// Add a single instance to the overloaded set (immediate
298    /// backpressure mark). Short-lived: the next metric-driven
299    /// `set_overloaded` recompute overwrites the whole set.
300    fn mark_overloaded(&self, instance_id: u64) -> Self {
301        let mut overloaded_ids = self.overloaded_ids.clone();
302        overloaded_ids.insert(instance_id);
303        Self::from_parts(
304            self.discovered_ids.clone(),
305            self.routable_ids.clone(),
306            overloaded_ids,
307        )
308    }
309
310    fn clear_overloaded_for_removed(&self, removed_ids: &HashSet<u64>) -> Self {
311        let mut overloaded_ids = self.overloaded_ids.clone();
312        overloaded_ids.retain(|id| !removed_ids.contains(id));
313        Self::from_parts(
314            self.discovered_ids.clone(),
315            self.routable_ids.clone(),
316            overloaded_ids,
317        )
318    }
319
320    fn derive_free_ids(routable_ids: &[u64], overloaded_ids: &HashSet<u64>) -> Vec<u64> {
321        if overloaded_ids.is_empty() {
322            return routable_ids.to_vec();
323        }
324
325        routable_ids
326            .iter()
327            .copied()
328            .filter(|id| !overloaded_ids.contains(id))
329            .collect()
330    }
331}
332
333#[derive(Debug)]
334struct RoutingInstancesState {
335    snapshot: ArcSwap<RoutingInstances>,
336    update_lock: StdMutex<()>,
337    instance_avail_tx: tokio::sync::watch::Sender<Vec<u64>>,
338    instance_avail_rx: tokio::sync::watch::Receiver<Vec<u64>>,
339}
340
341impl RoutingInstancesState {
342    fn new(discovered_ids: Vec<u64>) -> Self {
343        let snapshot = RoutingInstances::new(discovered_ids);
344        let (instance_avail_tx, instance_avail_rx) =
345            tokio::sync::watch::channel(snapshot.routable_ids().to_vec());
346        Self {
347            snapshot: ArcSwap::from_pointee(snapshot),
348            update_lock: StdMutex::new(()),
349            instance_avail_tx,
350            instance_avail_rx,
351        }
352    }
353
354    fn snapshot(&self) -> arc_swap::Guard<Arc<RoutingInstances>> {
355        self.snapshot.load()
356    }
357
358    fn update(
359        &self,
360        update: impl FnOnce(&RoutingInstances) -> RoutingInstances,
361        publish_routable_ids: bool,
362    ) -> Arc<RoutingInstances> {
363        let _guard = self.update_lock.lock().unwrap();
364        let current = self.snapshot.load();
365        let next = Arc::new(update(&current));
366        self.snapshot.store(next.clone());
367        if publish_routable_ids {
368            self.publish_routable_ids(&next);
369        }
370        next
371    }
372
373    fn publish_routable_ids(&self, routing_instances: &RoutingInstances) {
374        let _ = self
375            .instance_avail_tx
376            .send(routing_instances.routable_ids().to_vec());
377    }
378
379    fn routable_ids(&self) -> Vec<u64> {
380        self.snapshot().routable_ids().to_vec()
381    }
382
383    fn free_ids(&self) -> Vec<u64> {
384        self.snapshot().free_ids.clone()
385    }
386
387    fn counts(&self) -> RoutingInstanceCounts {
388        self.snapshot().counts()
389    }
390
391    fn overloaded_ids(&self) -> Option<HashSet<u64>> {
392        self.snapshot().overloaded_ids()
393    }
394
395    fn instance_avail_watcher(&self) -> tokio::sync::watch::Receiver<Vec<u64>> {
396        self.instance_avail_rx.clone()
397    }
398
399    fn report_instance_down(&self, instance_id: u64) {
400        self.update(|current| current.report_instance_down(instance_id), true);
401    }
402
403    fn set_overloaded_instances(&self, overloaded_instance_ids: &[u64]) -> bool {
404        let overloaded_ids = overloaded_instance_ids
405            .iter()
406            .copied()
407            .collect::<HashSet<_>>();
408        let _guard = self.update_lock.lock().unwrap();
409        let current = self.snapshot.load();
410        if current.overloaded_ids == overloaded_ids {
411            return false;
412        }
413
414        let next = Arc::new(current.set_overloaded(overloaded_ids));
415        self.snapshot.store(next);
416        true
417    }
418
419    fn mark_overloaded_immediate(&self, instance_id: u64) {
420        self.update(
421            move |current| current.mark_overloaded(instance_id),
422            // Routable set is unchanged — only the derived free set shrinks —
423            // so there's no need to republish routable_ids.
424            false,
425        );
426    }
427
428    fn clear_overloaded_for_removed(&self, removed_instance_ids: &[u64]) {
429        if removed_instance_ids.is_empty() {
430            return;
431        }
432
433        let removed_ids = removed_instance_ids.iter().copied().collect::<HashSet<_>>();
434        self.update(
435            move |current| current.clear_overloaded_for_removed(&removed_ids),
436            false,
437        );
438    }
439
440    fn reconcile_discovered(&self, discovered_ids: Vec<u64>) -> Arc<RoutingInstances> {
441        self.update(
442            move |current| current.reconcile_discovered(discovered_ids),
443            true,
444        )
445    }
446
447    #[cfg(test)]
448    fn override_routable_ids(&self, ids: Vec<u64>) {
449        self.update(move |current| current.override_routable_ids(ids), true);
450    }
451}
452
453#[derive(Clone, Debug)]
454pub struct Client {
455    // This is me
456    pub endpoint: Endpoint,
457    // Shared endpoint discovery source backing both snapshots and raw events.
458    endpoint_discovery_source: Arc<EndpointDiscoverySource>,
459    // These are the remotes I know about from watching key-value store
460    pub instance_source: Arc<tokio::sync::watch::Receiver<Vec<Instance>>>,
461    // Immutable routing snapshot. Free IDs are derived from discovered IDs and overloaded IDs.
462    routing_instances: Arc<RoutingInstancesState>,
463    /// Interval for periodic reconciliation of instance_avail with instance_source.
464    /// This ensures instances removed via `report_instance_down` are eventually restored.
465    reconcile_interval: Duration,
466}
467
468impl Client {
469    // Client with auto-discover instances using key-value store
470    pub(crate) async fn new(endpoint: Endpoint) -> Result<Self> {
471        Self::with_reconcile_interval(endpoint, DEFAULT_RECONCILE_INTERVAL).await
472    }
473
474    /// Create a client with a custom reconcile interval.
475    /// The reconcile interval controls how often `instance_avail` is reset to match
476    /// `instance_source`, restoring any instances removed via `report_instance_down`.
477    pub(crate) async fn with_reconcile_interval(
478        endpoint: Endpoint,
479        reconcile_interval: Duration,
480    ) -> Result<Self> {
481        tracing::trace!(
482            "Client::new_dynamic: Creating dynamic client for endpoint: {}",
483            endpoint.id()
484        );
485        let endpoint_discovery_source =
486            Self::get_or_create_dynamic_discovery_source(&endpoint).await?;
487        let instance_source = Arc::new(endpoint_discovery_source.instance_receiver());
488
489        // Seed instance_avail from the current instance_source snapshot so that
490        // callers who proceed immediately after wait_for_instances (which reads
491        // instance_source directly) will also find instances in instance_avail
492        // (which is read by the routing methods like random/round_robin).
493        let initial_ids: Vec<u64> = instance_source
494            .borrow()
495            .iter()
496            .map(|instance| instance.id())
497            .collect();
498        let client = Client {
499            endpoint: endpoint.clone(),
500            endpoint_discovery_source,
501            instance_source: instance_source.clone(),
502            routing_instances: Arc::new(RoutingInstancesState::new(initial_ids)),
503            reconcile_interval,
504        };
505        client.monitor_instance_source();
506        Ok(client)
507    }
508
509    /// Instances available from watching key-value store
510    pub fn instances(&self) -> Vec<Instance> {
511        self.instance_source.borrow().clone()
512    }
513
514    pub fn instance_ids(&self) -> Vec<u64> {
515        self.instances().into_iter().map(|ep| ep.id()).collect()
516    }
517
518    pub fn instance_ids_avail(&self) -> Vec<u64> {
519        self.routing_instances.routable_ids()
520    }
521
522    /// Routable instance ids excluding those currently flagged overloaded — the set used
523    /// for load-aware (random / round-robin) worker selection.
524    pub fn instance_ids_free(&self) -> Vec<u64> {
525        self.routing_instances.free_ids()
526    }
527
528    pub(crate) fn routing_instances(&self) -> arc_swap::Guard<Arc<RoutingInstances>> {
529        self.routing_instances.snapshot()
530    }
531
532    pub fn routing_instance_counts(&self) -> RoutingInstanceCounts {
533        self.routing_instances.counts()
534    }
535
536    /// Get a watcher for available instance IDs
537    pub fn instance_avail_watcher(&self) -> tokio::sync::watch::Receiver<Vec<u64>> {
538        self.routing_instances.instance_avail_watcher()
539    }
540
541    /// Subscribe to raw discovery events for this endpoint.
542    ///
543    /// Unlike `instance_source`, this feed does not coalesce remove→add pairs,
544    /// so consumers can react to every removal event exactly once.
545    pub(crate) fn subscribe_discovery_events(
546        &self,
547    ) -> tokio::sync::mpsc::UnboundedReceiver<DiscoveryEvent> {
548        self.endpoint_discovery_source.subscribe_events()
549    }
550
551    /// Wait for at least one Instance to be available for this Endpoint
552    pub async fn wait_for_instances(&self) -> Result<Vec<Instance>> {
553        tracing::trace!(
554            "wait_for_instances: Starting wait for endpoint: {}",
555            self.endpoint.id()
556        );
557        let mut rx = self.instance_source.as_ref().clone();
558        // wait for there to be 1 or more endpoints
559        let mut instances: Vec<Instance>;
560        loop {
561            instances = rx.borrow_and_update().to_vec();
562            if instances.is_empty() {
563                rx.changed().await?;
564            } else {
565                tracing::info!(
566                    "wait_for_instances: Found {} instance(s) for endpoint: {}",
567                    instances.len(),
568                    self.endpoint.id()
569                );
570                break;
571            }
572        }
573        Ok(instances)
574    }
575
576    /// Mark an instance as down/unavailable
577    pub fn report_instance_down(&self, instance_id: u64) {
578        self.routing_instances.report_instance_down(instance_id);
579        tracing::debug!("inhibiting instance {instance_id}");
580    }
581
582    /// Replace the set of overloaded instances reported by the worker monitor.
583    /// Returns true when this changes the routing snapshot.
584    pub fn set_overloaded_instances(&self, overloaded_instance_ids: &[u64]) -> bool {
585        self.routing_instances
586            .set_overloaded_instances(overloaded_instance_ids)
587    }
588
589    /// Mark an instance overloaded immediately. A worker returning
590    /// `ResourceExhausted` is busy ("queue full, retry later"), not faulted, so
591    /// this is the overload path, NOT `report_instance_down`. Short-lived: the
592    /// next `set_overloaded_instances` recompute overwrites the overloaded set.
593    pub fn mark_overloaded_immediate(&self, instance_id: u64) {
594        self.routing_instances
595            .mark_overloaded_immediate(instance_id);
596        tracing::debug!(
597            instance_id,
598            "marking instance overloaded (backpressure); next metric event will re-evaluate"
599        );
600    }
601
602    pub fn clear_overloaded_instances_for_removed(&self, removed_instance_ids: &[u64]) {
603        self.routing_instances
604            .clear_overloaded_for_removed(removed_instance_ids);
605    }
606
607    pub fn overloaded_instance_ids(&self) -> Option<HashSet<u64>> {
608        self.routing_instances.overloaded_ids()
609    }
610
611    /// Monitor the key-value instance source and update instance_avail.
612    ///
613    /// This function also performs periodic reconciliation: if `instance_source` hasn't
614    /// changed for `reconcile_interval`, we reset `instance_avail` to match
615    /// `instance_source`. This ensures instances removed via `report_instance_down`
616    /// are eventually restored even if the discovery source doesn't emit updates.
617    fn monitor_instance_source(&self) {
618        let reconcile_interval = self.reconcile_interval;
619        let cancel_token = self.endpoint.drt().primary_token();
620        let client = self.clone();
621        let endpoint_id = self.endpoint.id();
622        tokio::task::spawn(async move {
623            let mut rx = client.instance_source.as_ref().clone();
624            while !cancel_token.is_cancelled() {
625                let instance_ids: Vec<u64> = rx
626                    .borrow_and_update()
627                    .iter()
628                    .map(|instance| instance.id())
629                    .collect();
630
631                let routing_instances = client.reconcile_discovered_instances(instance_ids);
632
633                // Clean up stale occupancy counters for instances that no longer exist.
634                let registry = client.endpoint.drt().routing_occupancy_states();
635                if let Ok(registry) = registry.try_lock()
636                    && let Some(weak) = registry.get(&client.endpoint)
637                    && let Some(state) = weak.upgrade()
638                {
639                    state.retain(routing_instances.discovered_ids());
640                }
641
642                tokio::select! {
643                    result = rx.changed() => {
644                        if let Err(err) = result {
645                            tracing::error!(
646                                "monitor_instance_source: The Sender is dropped: {err}, endpoint={endpoint_id}",
647                            );
648                            cancel_token.cancel();
649                        }
650                    }
651                    _ = tokio::time::sleep(reconcile_interval) => {
652                        tracing::trace!(
653                            "monitor_instance_source: periodic reconciliation for endpoint={endpoint_id}",
654                        );
655                    }
656                }
657            }
658        });
659    }
660
661    /// Override routable IDs for testing. This allows creating an inconsistency
662    /// between `instance_ids_avail()` and `instances()` to simulate downed workers.
663    #[cfg(test)]
664    pub(crate) fn override_instance_avail(&self, ids: Vec<u64>) {
665        self.routing_instances.override_routable_ids(ids);
666    }
667
668    fn reconcile_discovered_instances(&self, discovered_ids: Vec<u64>) -> Arc<RoutingInstances> {
669        self.routing_instances.reconcile_discovered(discovered_ids)
670    }
671
672    async fn get_or_create_dynamic_discovery_source(
673        endpoint: &Endpoint,
674    ) -> Result<Arc<EndpointDiscoverySource>> {
675        let drt = endpoint.drt();
676        let sources = drt.endpoint_discovery_sources();
677        let mut sources = sources.lock().await;
678
679        if let Some(source) = sources.get(endpoint) {
680            if let Some(source) = source.upgrade() {
681                return Ok(source);
682            } else {
683                sources.remove(endpoint);
684            }
685        }
686
687        let discovery = drt.discovery();
688        let discovery_query = crate::discovery::DiscoveryQuery::Endpoint {
689            namespace: endpoint.component.namespace.name.clone(),
690            component: endpoint.component.name.clone(),
691            endpoint: endpoint.name.clone(),
692        };
693
694        let mut discovery_stream = discovery
695            .list_and_watch(discovery_query.clone(), None)
696            .await?;
697        let (watch_tx, watch_rx) = tokio::sync::watch::channel(vec![]);
698        let discovery_source = Arc::new(EndpointDiscoverySource::new(watch_rx));
699
700        let secondary = endpoint.component.drt.runtime().secondary().clone();
701        let discovery_source_task = discovery_source.clone();
702
703        secondary.spawn(async move {
704            tracing::trace!("endpoint_watcher: Starting for discovery query: {:?}", discovery_query);
705            let mut map: HashMap<u64, Instance> = HashMap::new();
706
707            loop {
708                let discovery_event = tokio::select! {
709                    _ = watch_tx.closed() => {
710                        break;
711                    }
712                    discovery_event = discovery_stream.next() => {
713                        match discovery_event {
714                            Some(Ok(event)) => {
715                                event
716                            },
717                            Some(Err(e)) => {
718                                tracing::error!("endpoint_watcher: discovery stream error: {}; shutting down for discovery query: {:?}", e, discovery_query);
719                                break;
720                            }
721                            None => {
722                                break;
723                            }
724                        }
725                    }
726                };
727
728                discovery_source_task.broadcast_event(&discovery_event);
729
730                match discovery_event {
731                    DiscoveryEvent::Added(DiscoveryInstance::Endpoint(instance)) => {
732                        map.insert(instance.instance_id, instance);
733                    }
734                    DiscoveryEvent::Added(_) => {}
735                    DiscoveryEvent::Removed(id) => {
736                        if let DiscoveryInstanceId::Endpoint(endpoint_id) = id {
737                            map.remove(&endpoint_id.instance_id);
738                        }
739                    }
740                }
741
742                let instances: Vec<Instance> = map.values().cloned().collect();
743                if watch_tx.send(instances).is_err() {
744                    break;
745                }
746            }
747            let _ = watch_tx.send(vec![]);
748        });
749
750        sources.insert(endpoint.clone(), Arc::downgrade(&discovery_source));
751        Ok(discovery_source)
752    }
753}
754
755#[cfg(test)]
756mod tests {
757    use super::*;
758    use crate::{DistributedRuntime, Runtime, distributed::DistributedConfig};
759
760    /// Test that instances removed via report_instance_down are restored after
761    /// the reconciliation interval elapses.
762    #[tokio::test]
763    async fn test_instance_reconciliation() {
764        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(100);
765
766        let rt = Runtime::from_current().unwrap();
767        // Use process_local config to avoid needing etcd/nats
768        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
769            .await
770            .unwrap();
771        let ns = drt.namespace("test_reconciliation".to_string()).unwrap();
772        let component = ns.component("test_component".to_string()).unwrap();
773        let endpoint = component.endpoint("test_endpoint".to_string());
774
775        // Use a short reconcile interval for faster tests
776        let client = Client::with_reconcile_interval(endpoint, TEST_RECONCILE_INTERVAL)
777            .await
778            .unwrap();
779
780        // Initially, instance_avail should be empty (no registered instances)
781        assert!(client.instance_ids_avail().is_empty());
782
783        // For this test, we'll directly manipulate instance_avail and verify reconciliation
784        // Store some test IDs
785        client.override_instance_avail(vec![1, 2, 3]);
786
787        assert_eq!(client.instance_ids_avail(), vec![1u64, 2, 3]);
788
789        // Simulate report_instance_down removing instance 2
790        client.report_instance_down(2);
791        assert_eq!(client.instance_ids_avail(), vec![1u64, 3]);
792
793        // Wait for reconciliation interval + buffer
794        // The monitor_instance_source will reset instance_avail to match instance_source
795        // Since instance_source is empty, after reconciliation instance_avail should be empty
796        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;
797
798        // After reconciliation, instance_avail should match instance_source (which is empty)
799        assert!(
800            client.instance_ids_avail().is_empty(),
801            "After reconciliation, instance_avail should match instance_source"
802        );
803
804        rt.shutdown();
805    }
806
807    /// Test that report_instance_down correctly removes an instance from instance_avail.
808    #[tokio::test]
809    async fn test_report_instance_down() {
810        let rt = Runtime::from_current().unwrap();
811        // Use process_local config to avoid needing etcd/nats
812        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
813            .await
814            .unwrap();
815        let ns = drt.namespace("test_report_down".to_string()).unwrap();
816        let component = ns.component("test_component".to_string()).unwrap();
817        let endpoint = component.endpoint("test_endpoint".to_string());
818
819        let client = endpoint.client().await.unwrap();
820
821        // Manually set up instance_avail with test instances
822        client.override_instance_avail(vec![1, 2, 3]);
823        assert_eq!(client.instance_ids_avail(), vec![1u64, 2, 3]);
824
825        // Report instance 2 as down
826        client.report_instance_down(2);
827
828        // Verify instance 2 is removed
829        let avail = client.instance_ids_avail();
830        assert!(avail.contains(&1), "Instance 1 should still be available");
831        assert!(
832            !avail.contains(&2),
833            "Instance 2 should be removed after report_instance_down"
834        );
835        assert!(avail.contains(&3), "Instance 3 should still be available");
836
837        rt.shutdown();
838    }
839
840    #[tokio::test]
841    async fn test_overloaded_instance_ids_returns_none_when_empty() {
842        let rt = Runtime::from_current().unwrap();
843        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
844            .await
845            .unwrap();
846        let ns = drt.namespace("test_overloaded_ids".to_string()).unwrap();
847        let component = ns.component("test_component".to_string()).unwrap();
848        let endpoint = component.endpoint("test_endpoint".to_string());
849        let client = endpoint.client().await.unwrap();
850
851        assert_eq!(client.overloaded_instance_ids(), None);
852
853        assert!(client.set_overloaded_instances(&[7]));
854        assert_eq!(client.overloaded_instance_ids(), Some(HashSet::from([7])));
855        assert!(!client.set_overloaded_instances(&[7]));
856
857        assert!(client.set_overloaded_instances(&[]));
858        assert_eq!(client.overloaded_instance_ids(), None);
859        assert!(!client.set_overloaded_instances(&[]));
860
861        rt.shutdown();
862    }
863
864    #[tokio::test]
865    async fn test_instance_reconciliation_preserves_overloaded_existing_instances() {
866        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
867
868        let rt = Runtime::from_current().unwrap();
869        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
870            .await
871            .unwrap();
872        let ns = drt
873            .namespace("test_overloaded_reconciliation".to_string())
874            .unwrap();
875        let component = ns.component("test_component".to_string()).unwrap();
876        let endpoint = component.endpoint("test_endpoint".to_string());
877
878        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
879            .await
880            .unwrap();
881        endpoint.register_endpoint_instance().await.unwrap();
882        let instances = client.wait_for_instances().await.unwrap();
883        let worker_id = instances[0].id();
884
885        for _ in 0..10 {
886            if client.instance_ids_free().contains(&worker_id) {
887                break;
888            }
889            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
890        }
891        assert!(
892            client.instance_ids_free().contains(&worker_id),
893            "worker should be free after initial discovery reconciliation"
894        );
895
896        client.set_overloaded_instances(&[worker_id]);
897        assert!(
898            client.instance_ids_free().is_empty(),
899            "worker should be overloaded before periodic reconciliation"
900        );
901
902        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;
903
904        assert!(
905            client.instance_ids_free().is_empty(),
906            "periodic reconciliation should not mark an existing overloaded worker free"
907        );
908
909        rt.shutdown();
910    }
911
912    #[tokio::test]
913    async fn test_report_instance_down_preserves_overloaded_state() {
914        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
915
916        let rt = Runtime::from_current().unwrap();
917        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
918            .await
919            .unwrap();
920        let ns = drt
921            .namespace("test_report_down_preserves_overloaded".to_string())
922            .unwrap();
923        let component = ns.component("test_component".to_string()).unwrap();
924        let endpoint = component.endpoint("test_endpoint".to_string());
925
926        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
927            .await
928            .unwrap();
929        endpoint.register_endpoint_instance().await.unwrap();
930        let instances = client.wait_for_instances().await.unwrap();
931        let worker_id = instances[0].id();
932
933        for _ in 0..10 {
934            if client.instance_ids_avail().contains(&worker_id) {
935                break;
936            }
937            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
938        }
939
940        client.set_overloaded_instances(&[worker_id]);
941        client.report_instance_down(worker_id);
942
943        assert!(
944            !client.instance_ids_avail().contains(&worker_id),
945            "reported-down worker should leave routable availability"
946        );
947        assert_eq!(
948            client.routing_instance_counts().overloaded,
949            1,
950            "reported-down worker should remain overloaded while still discovered"
951        );
952        assert!(
953            client.instance_ids_free().is_empty(),
954            "reported-down overloaded worker should not become free"
955        );
956
957        endpoint.unregister_endpoint_instance().await.unwrap();
958        for _ in 0..10 {
959            if client.routing_instance_counts().overloaded == 0 {
960                break;
961            }
962            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
963        }
964
965        assert_eq!(
966            client.routing_instance_counts().overloaded,
967            0,
968            "stable discovery removal should clear overloaded state"
969        );
970
971        rt.shutdown();
972    }
973
974    #[tokio::test]
975    async fn test_instance_reconciliation_prunes_removed_overloaded_instances() {
976        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
977
978        let rt = Runtime::from_current().unwrap();
979        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
980            .await
981            .unwrap();
982        let ns = drt
983            .namespace("test_removed_overloaded_cleanup".to_string())
984            .unwrap();
985        let component = ns.component("test_component".to_string()).unwrap();
986        let endpoint = component.endpoint("test_endpoint".to_string());
987
988        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
989            .await
990            .unwrap();
991        endpoint.register_endpoint_instance().await.unwrap();
992        let instances = client.wait_for_instances().await.unwrap();
993        let worker_id = instances[0].id();
994
995        client.set_overloaded_instances(&[worker_id]);
996        assert_eq!(client.routing_instance_counts().overloaded, 1);
997        assert!(client.instance_ids_free().is_empty());
998
999        endpoint.unregister_endpoint_instance().await.unwrap();
1000        for _ in 0..10 {
1001            if client.routing_instance_counts().overloaded == 0 {
1002                break;
1003            }
1004            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1005        }
1006
1007        assert_eq!(
1008            client.routing_instance_counts().overloaded,
1009            0,
1010            "removed discovered workers should not remain in overloaded state"
1011        );
1012
1013        rt.shutdown();
1014    }
1015
1016    #[tokio::test]
1017    async fn test_instance_ids_free_excludes_overloaded_new_instances() {
1018        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1019
1020        let rt = Runtime::from_current().unwrap();
1021        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1022            .await
1023            .unwrap();
1024        let worker_id = drt.connection_id();
1025        let ns = drt
1026            .namespace("test_new_overloaded_reconciliation".to_string())
1027            .unwrap();
1028        let component = ns.component("test_component".to_string()).unwrap();
1029        let endpoint = component.endpoint("test_endpoint".to_string());
1030
1031        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1032            .await
1033            .unwrap();
1034        client.set_overloaded_instances(&[worker_id]);
1035
1036        endpoint.register_endpoint_instance().await.unwrap();
1037        let instances = client.wait_for_instances().await.unwrap();
1038        assert_eq!(instances[0].id(), worker_id);
1039        assert!(
1040            client.instance_ids_free().is_empty(),
1041            "newly discovered overloaded worker should not be free"
1042        );
1043
1044        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;
1045
1046        assert!(
1047            client.instance_ids_free().is_empty(),
1048            "discovery reconciliation should not affect recomputed free workers"
1049        );
1050
1051        rt.shutdown();
1052    }
1053
1054    #[tokio::test]
1055    async fn test_discovery_add_updates_free_without_overloaded_publish() {
1056        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1057
1058        let rt = Runtime::from_current().unwrap();
1059        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1060            .await
1061            .unwrap();
1062        let ns = drt
1063            .namespace("test_free_updates_on_discovery_add".to_string())
1064            .unwrap();
1065        let component = ns.component("test_component".to_string()).unwrap();
1066        let endpoint = component.endpoint("test_endpoint".to_string());
1067
1068        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1069            .await
1070            .unwrap();
1071        endpoint.register_endpoint_instance().await.unwrap();
1072        let instances = client.wait_for_instances().await.unwrap();
1073        let worker_id = instances[0].id();
1074
1075        for _ in 0..10 {
1076            if client.instance_ids_free().contains(&worker_id) {
1077                break;
1078            }
1079            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1080        }
1081
1082        assert_eq!(
1083            client.instance_ids_free(),
1084            vec![worker_id],
1085            "newly discovered non-overloaded workers should appear free without an overload update"
1086        );
1087
1088        rt.shutdown();
1089    }
1090
1091    /// Test that instance_avail_watcher receives updates when instances change.
1092    #[tokio::test]
1093    async fn test_instance_avail_watcher() {
1094        let rt = Runtime::from_current().unwrap();
1095        // Use process_local config to avoid needing etcd/nats
1096        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1097            .await
1098            .unwrap();
1099        let ns = drt.namespace("test_watcher".to_string()).unwrap();
1100        let component = ns.component("test_component".to_string()).unwrap();
1101        let endpoint = component.endpoint("test_endpoint".to_string());
1102
1103        let client = endpoint.client().await.unwrap();
1104        let watcher = client.instance_avail_watcher();
1105
1106        // Set initial instances
1107        client.override_instance_avail(vec![1, 2, 3]);
1108
1109        // Report instance down - this should notify the watcher
1110        client.report_instance_down(2);
1111
1112        // The watcher should receive the update
1113        // Note: We need to check if changed() was signaled
1114        let current = watcher.borrow().clone();
1115        assert_eq!(current, vec![1, 3]);
1116
1117        rt.shutdown();
1118    }
1119
1120    /// Test that concurrent select_and_increment distributes load correctly.
1121    #[tokio::test]
1122    async fn test_concurrent_select_and_increment() {
1123        let state = Arc::new(RoutingOccupancyState::default());
1124        let instance_ids: Vec<u64> = vec![100, 200, 300];
1125        let num_requests = 90;
1126
1127        let mut handles = Vec::new();
1128        for _ in 0..num_requests {
1129            let state = state.clone();
1130            let ids = instance_ids.clone();
1131            handles.push(tokio::spawn(async move {
1132                state.select_exact_min_and_increment(&ids).await
1133            }));
1134        }
1135
1136        for handle in handles {
1137            handle.await.unwrap();
1138        }
1139
1140        assert_eq!(state.load(100), 30);
1141        assert_eq!(state.load(200), 30);
1142        assert_eq!(state.load(300), 30);
1143    }
1144
1145    #[tokio::test]
1146    async fn test_select_exact_min_and_increment_randomizes_ties() {
1147        let mut selected = [false; 3];
1148
1149        for _ in 0..120 {
1150            let state = RoutingOccupancyState::default();
1151            let picked = state
1152                .select_exact_min_and_increment(&[10, 20, 30])
1153                .await
1154                .unwrap();
1155            match picked {
1156                10 => selected[0] = true,
1157                20 => selected[1] = true,
1158                30 => selected[2] = true,
1159                _ => panic!("unexpected worker id: {picked}"),
1160            }
1161        }
1162
1163        let selected_count = selected.into_iter().filter(|seen| *seen).count();
1164        assert!(
1165            selected_count > 1,
1166            "tie-breaking should not always select the first minimum-load worker"
1167        );
1168    }
1169
1170    #[tokio::test]
1171    async fn test_connection_counts() {
1172        let rt = Runtime::from_current().unwrap();
1173        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1174            .await
1175            .unwrap();
1176        let ns = drt.namespace("test_ll_counts".to_string()).unwrap();
1177        let component = ns.component("test_component".to_string()).unwrap();
1178        let endpoint = component.endpoint("test_endpoint".to_string());
1179
1180        let state1 = get_or_create_routing_occupancy_state(&endpoint).await;
1181        let state2 = get_or_create_routing_occupancy_state(&endpoint).await;
1182
1183        let picked1 = state1
1184            .select_exact_min_and_increment(&[10, 20, 30])
1185            .await
1186            .unwrap();
1187        assert_eq!(state1.load(picked1), 1);
1188
1189        let picked2 = state1
1190            .select_exact_min_and_increment(&[10, 20, 30])
1191            .await
1192            .unwrap();
1193        assert_ne!(picked1, picked2);
1194
1195        // state2 should see the same counts (same underlying Arc)
1196        assert_eq!(state2.load(10), state1.load(10));
1197        assert_eq!(state2.load(20), state1.load(20));
1198        assert_eq!(state2.load(30), state1.load(30));
1199
1200        state2.decrement(picked1);
1201        assert_eq!(state1.load(picked1), if picked1 == picked2 { 1 } else { 0 });
1202
1203        rt.shutdown();
1204    }
1205
1206    #[tokio::test]
1207    async fn test_least_loaded_state_retain() {
1208        let state = RoutingOccupancyState::default();
1209
1210        // Add some connections
1211        state.select_exact_min_and_increment(&[1, 2, 3]).await;
1212        state.select_exact_min_and_increment(&[1, 2, 3]).await;
1213        state.select_exact_min_and_increment(&[1, 2, 3]).await;
1214        // Each instance should have 1 connection
1215        assert_eq!(state.load(1), 1);
1216        assert_eq!(state.load(2), 1);
1217        assert_eq!(state.load(3), 1);
1218
1219        // Retain only instances 1 and 3 (instance 2 was removed)
1220        state.retain(&[1, 3]);
1221
1222        assert_eq!(state.load(1), 1);
1223        assert_eq!(state.load(2), 0);
1224        assert_eq!(state.load(3), 1);
1225    }
1226
1227    #[tokio::test]
1228    async fn test_monitor_instance_source_cleans_up_removed_worker_counts() {
1229        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1230
1231        let rt = Runtime::from_current().unwrap();
1232        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1233            .await
1234            .unwrap();
1235        let ns = drt.namespace("test_occupancy_cleanup".to_string()).unwrap();
1236        let component = ns.component("test_component".to_string()).unwrap();
1237        let endpoint = component.endpoint("test_endpoint".to_string());
1238
1239        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1240            .await
1241            .unwrap();
1242        endpoint.register_endpoint_instance().await.unwrap();
1243        client.wait_for_instances().await.unwrap();
1244
1245        let worker_id = client.instance_ids_avail()[0];
1246        let state = get_or_create_routing_occupancy_state(&endpoint).await;
1247        state.increment(worker_id);
1248        assert_eq!(state.load(worker_id), 1);
1249
1250        endpoint.unregister_endpoint_instance().await.unwrap();
1251
1252        for _ in 0..10 {
1253            if state.load(worker_id) == 0 {
1254                break;
1255            }
1256            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1257        }
1258
1259        assert_eq!(state.load(worker_id), 0);
1260
1261        rt.shutdown();
1262    }
1263}
dynamo_runtime/component/client.rs

dynamo_runtime/component/
client.rs
