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::{AtomicBool, Ordering};
5use std::{
6    collections::{HashMap, HashSet},
7    sync::{Arc, LazyLock, Mutex as StdMutex},
8    time::Duration,
9};
10
11use anyhow::Result;
12use arc_swap::ArcSwap;
13use futures::StreamExt;
14
15use crate::component::{Endpoint, Instance};
16use crate::config::environment_names::runtime as env_runtime;
17use crate::discovery::{DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId};
18use crate::routing_policy::{RoutingOccupancyState, get_or_create_routing_occupancy_state};
19use crate::traits::DistributedRuntimeProvider;
20
21/// Default interval for periodic reconciliation of instance_avail with instance_source
22const DEFAULT_INHIBITED_DURATION_SECS: u64 = 5;
23
24/// Process-wide inhibited duration, resolved from the environment on first client construction.
25static INHIBITED_DURATION: LazyLock<Duration> =
26    LazyLock::new(|| inhibited_duration_from_env(|name| std::env::var(name).ok()));
27
28fn inhibited_duration_from_env(mut lookup: impl FnMut(&str) -> Option<String>) -> Duration {
29    let seconds = match lookup(env_runtime::DYN_RUNTIME_INHIBITED_DURATION_SECS) {
30        None => DEFAULT_INHIBITED_DURATION_SECS,
31        Some(raw) => match raw.parse::<u64>() {
32            Ok(seconds) => seconds,
33            Err(err) => {
34                tracing::warn!(
35                    value = raw,
36                    %err,
37                    "invalid {}; using the default of {} seconds",
38                    env_runtime::DYN_RUNTIME_INHIBITED_DURATION_SECS,
39                    DEFAULT_INHIBITED_DURATION_SECS,
40                );
41                DEFAULT_INHIBITED_DURATION_SECS
42            }
43        },
44    };
45    Duration::from_secs(seconds)
46}
47
48/// Shared endpoint discovery state for a single endpoint query.
49///
50/// This wraps both the coalesced instance snapshot used for routing decisions
51/// and a raw, lossless per-subscriber event feed used by the response-stream
52/// cancellation watcher. Both outputs are driven by a single underlying
53/// discovery `list_and_watch` task so clients do not multiply control-plane
54/// watches.
55#[derive(Debug)]
56pub(crate) struct EndpointDiscoverySource {
57    instance_source: tokio::sync::watch::Receiver<Vec<Instance>>,
58    event_subscribers: StdMutex<Vec<tokio::sync::mpsc::UnboundedSender<DiscoveryEvent>>>,
59}
60
61pub(crate) struct DiscoveryEventReceiver {
62    receiver: tokio::sync::mpsc::UnboundedReceiver<DiscoveryEvent>,
63    _source: Arc<EndpointDiscoverySource>,
64}
65
66impl std::ops::Deref for DiscoveryEventReceiver {
67    type Target = tokio::sync::mpsc::UnboundedReceiver<DiscoveryEvent>;
68
69    fn deref(&self) -> &Self::Target {
70        &self.receiver
71    }
72}
73
74impl std::ops::DerefMut for DiscoveryEventReceiver {
75    fn deref_mut(&mut self) -> &mut Self::Target {
76        &mut self.receiver
77    }
78}
79
80impl EndpointDiscoverySource {
81    fn new(instance_source: tokio::sync::watch::Receiver<Vec<Instance>>) -> Self {
82        Self {
83            instance_source,
84            event_subscribers: StdMutex::new(Vec::new()),
85        }
86    }
87
88    fn instance_receiver(&self) -> tokio::sync::watch::Receiver<Vec<Instance>> {
89        self.instance_source.clone()
90    }
91
92    fn subscribe_events(&self) -> tokio::sync::mpsc::UnboundedReceiver<DiscoveryEvent> {
93        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
94        self.event_subscribers.lock().unwrap().push(tx);
95        rx
96    }
97
98    fn broadcast_event(&self, event: &DiscoveryEvent) {
99        let subscribers = &mut *self.event_subscribers.lock().unwrap();
100        subscribers.retain(|tx| tx.send(event.clone()).is_ok());
101    }
102}
103
104#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
105pub struct RoutingInstanceCounts {
106    pub discovered: usize,
107    pub routable: usize,
108    pub overloaded: usize,
109    /// IDs not currently reported overloaded, derived from `discovered - overloaded`.
110    pub free: usize,
111}
112
113#[derive(Clone, Debug)]
114pub(crate) struct RoutingInstances {
115    discovered_ids: Vec<u64>,
116    routable_ids: Vec<u64>,
117    overloaded_ids: HashSet<u64>,
118    free_ids: Vec<u64>,
119    routable_id_set: Arc<HashSet<u64>>,
120    /// True after this client has observed at least one discovered instance.
121    /// Once set, a later empty snapshot is authoritative rather than startup
122    /// absence of information.
123    availability_initialized: bool,
124}
125
126impl RoutingInstances {
127    fn new(discovered_ids: Vec<u64>) -> Self {
128        let availability_initialized = !discovered_ids.is_empty();
129        Self::from_parts(
130            discovered_ids.clone(),
131            discovered_ids,
132            HashSet::new(),
133            availability_initialized,
134        )
135    }
136
137    fn from_parts(
138        mut discovered_ids: Vec<u64>,
139        mut routable_ids: Vec<u64>,
140        overloaded_ids: HashSet<u64>,
141        availability_initialized: bool,
142    ) -> Self {
143        discovered_ids.sort_unstable();
144        discovered_ids.dedup();
145        routable_ids.sort_unstable();
146        routable_ids.dedup();
147        let free_ids = Self::derive_free_ids(&routable_ids, &overloaded_ids);
148        let routable_id_set = Arc::new(routable_ids.iter().copied().collect());
149        Self {
150            discovered_ids,
151            routable_ids,
152            overloaded_ids,
153            free_ids,
154            routable_id_set,
155            availability_initialized,
156        }
157    }
158
159    pub(crate) fn discovered_ids(&self) -> &[u64] {
160        &self.discovered_ids
161    }
162
163    pub(crate) fn routable_ids(&self) -> &[u64] {
164        &self.routable_ids
165    }
166
167    fn available_ids(&self) -> Option<Arc<HashSet<u64>>> {
168        self.availability_initialized
169            .then(|| Arc::clone(&self.routable_id_set))
170    }
171
172    pub(crate) fn free_ids(&self) -> &[u64] {
173        &self.free_ids
174    }
175
176    pub(crate) fn counts(&self) -> RoutingInstanceCounts {
177        RoutingInstanceCounts {
178            discovered: self.discovered_ids.len(),
179            routable: self.routable_ids.len(),
180            overloaded: self.overloaded_ids.len(),
181            free: self.free_ids.len(),
182        }
183    }
184
185    pub(crate) fn is_overloaded(&self, instance_id: u64) -> bool {
186        self.overloaded_ids.contains(&instance_id)
187    }
188
189    fn overloaded_ids(&self) -> Option<HashSet<u64>> {
190        if self.overloaded_ids.is_empty() {
191            return None;
192        }
193
194        Some(self.overloaded_ids.clone())
195    }
196
197    fn reconcile_discovered(&self, discovered_ids: Vec<u64>) -> Self {
198        let old_discovered_ids = self.discovered_ids.iter().copied().collect::<HashSet<_>>();
199        let new_discovered_ids = discovered_ids.iter().copied().collect::<HashSet<_>>();
200        let mut overloaded_ids = self.overloaded_ids.clone();
201        overloaded_ids
202            .retain(|id| !old_discovered_ids.contains(id) || new_discovered_ids.contains(id));
203
204        let availability_initialized = self.availability_initialized || !discovered_ids.is_empty();
205        Self::from_parts(
206            discovered_ids.clone(),
207            discovered_ids,
208            overloaded_ids,
209            availability_initialized,
210        )
211    }
212
213    fn report_instance_down(&self, instance_id: u64) -> Self {
214        let routable_ids: Vec<u64> = self
215            .routable_ids
216            .iter()
217            .copied()
218            .filter(|id| *id != instance_id)
219            .collect();
220
221        Self::from_parts(
222            self.discovered_ids.clone(),
223            routable_ids,
224            self.overloaded_ids.clone(),
225            self.availability_initialized,
226        )
227    }
228
229    #[cfg(any(test, feature = "testing"))]
230    fn override_routable_ids(&self, routable_ids: Vec<u64>) -> Self {
231        // Route through from_parts so `free_ids` is recomputed from the new
232        // routable set instead of carrying the stale value forward.
233        Self::from_parts(
234            self.discovered_ids.clone(),
235            routable_ids,
236            self.overloaded_ids.clone(),
237            self.availability_initialized,
238        )
239    }
240
241    fn set_overloaded(&self, overloaded_ids: HashSet<u64>) -> Self {
242        Self::from_parts(
243            self.discovered_ids.clone(),
244            self.routable_ids.clone(),
245            overloaded_ids,
246            self.availability_initialized,
247        )
248    }
249
250    /// Add a single instance to the overloaded set (immediate
251    /// backpressure mark). Short-lived: the next metric-driven
252    /// `set_overloaded` recompute overwrites the whole set.
253    fn mark_overloaded(&self, instance_id: u64) -> Self {
254        let mut overloaded_ids = self.overloaded_ids.clone();
255        overloaded_ids.insert(instance_id);
256        Self::from_parts(
257            self.discovered_ids.clone(),
258            self.routable_ids.clone(),
259            overloaded_ids,
260            self.availability_initialized,
261        )
262    }
263
264    fn clear_overloaded_for_removed(&self, removed_ids: &HashSet<u64>) -> Self {
265        let mut overloaded_ids = self.overloaded_ids.clone();
266        overloaded_ids.retain(|id| !removed_ids.contains(id));
267        Self::from_parts(
268            self.discovered_ids.clone(),
269            self.routable_ids.clone(),
270            overloaded_ids,
271            self.availability_initialized,
272        )
273    }
274
275    fn derive_free_ids(routable_ids: &[u64], overloaded_ids: &HashSet<u64>) -> Vec<u64> {
276        if overloaded_ids.is_empty() {
277            return routable_ids.to_vec();
278        }
279
280        routable_ids
281            .iter()
282            .copied()
283            .filter(|id| !overloaded_ids.contains(id))
284            .collect()
285    }
286}
287
288#[derive(Debug)]
289struct RoutingInstancesState {
290    snapshot: ArcSwap<RoutingInstances>,
291    update_lock: StdMutex<()>,
292    overload_reconciliation_needed: AtomicBool,
293    instance_avail_tx: tokio::sync::watch::Sender<Vec<u64>>,
294}
295
296impl RoutingInstancesState {
297    fn new(discovered_ids: Vec<u64>) -> (Self, tokio::sync::watch::Receiver<Vec<u64>>) {
298        let snapshot = RoutingInstances::new(discovered_ids);
299        let (instance_avail_tx, instance_avail_rx) =
300            tokio::sync::watch::channel(snapshot.routable_ids().to_vec());
301        (
302            Self {
303                snapshot: ArcSwap::from_pointee(snapshot),
304                update_lock: StdMutex::new(()),
305                overload_reconciliation_needed: AtomicBool::new(false),
306                instance_avail_tx,
307            },
308            instance_avail_rx,
309        )
310    }
311
312    fn snapshot(&self) -> arc_swap::Guard<Arc<RoutingInstances>> {
313        self.snapshot.load()
314    }
315
316    fn update(
317        &self,
318        update: impl FnOnce(&RoutingInstances) -> RoutingInstances,
319        publish_routable_ids: bool,
320    ) -> Arc<RoutingInstances> {
321        let _guard = self.update_lock.lock().unwrap();
322        let current = self.snapshot.load();
323        let next = Arc::new(update(&current));
324        self.snapshot.store(next.clone());
325        if publish_routable_ids {
326            self.publish_routable_ids(&next);
327        }
328        next
329    }
330
331    fn publish_routable_ids(&self, routing_instances: &RoutingInstances) {
332        let _ = self
333            .instance_avail_tx
334            .send(routing_instances.routable_ids().to_vec());
335    }
336
337    fn routable_ids(&self) -> Vec<u64> {
338        self.snapshot().routable_ids().to_vec()
339    }
340
341    fn available_ids(&self) -> Option<Arc<HashSet<u64>>> {
342        self.snapshot().available_ids()
343    }
344
345    fn free_ids(&self) -> Vec<u64> {
346        self.snapshot().free_ids.clone()
347    }
348
349    fn counts(&self) -> RoutingInstanceCounts {
350        self.snapshot().counts()
351    }
352
353    fn overloaded_ids(&self) -> Option<HashSet<u64>> {
354        self.snapshot().overloaded_ids()
355    }
356
357    fn report_instance_down(&self, instance_id: u64) {
358        self.update(|current| current.report_instance_down(instance_id), true);
359    }
360
361    fn set_overloaded_instances(&self, overloaded_instance_ids: &[u64]) -> bool {
362        let overloaded_ids = overloaded_instance_ids
363            .iter()
364            .copied()
365            .collect::<HashSet<_>>();
366        let _guard = self.update_lock.lock().unwrap();
367        self.overload_reconciliation_needed
368            .store(false, Ordering::Release);
369        let current = self.snapshot.load();
370        if current.overloaded_ids == overloaded_ids {
371            return false;
372        }
373
374        let next = Arc::new(current.set_overloaded(overloaded_ids));
375        self.snapshot.store(next);
376        true
377    }
378
379    fn mark_overloaded_immediate(&self, instance_id: u64) {
380        let _guard = self.update_lock.lock().unwrap();
381        let current = self.snapshot.load();
382        let next = Arc::new(current.mark_overloaded(instance_id));
383        self.snapshot.store(next);
384        self.overload_reconciliation_needed
385            .store(true, Ordering::Release);
386    }
387
388    fn overload_reconciliation_needed(&self) -> bool {
389        self.overload_reconciliation_needed.load(Ordering::Acquire)
390    }
391
392    fn clear_overloaded_for_removed(&self, removed_instance_ids: &[u64]) {
393        if removed_instance_ids.is_empty() {
394            return;
395        }
396
397        let removed_ids = removed_instance_ids.iter().copied().collect::<HashSet<_>>();
398        self.update(
399            move |current| current.clear_overloaded_for_removed(&removed_ids),
400            false,
401        );
402    }
403
404    fn reconcile_discovered(&self, discovered_ids: Vec<u64>) -> Arc<RoutingInstances> {
405        self.update(
406            move |current| current.reconcile_discovered(discovered_ids),
407            true,
408        )
409    }
410
411    #[cfg(any(test, feature = "testing"))]
412    fn override_routable_ids(&self, ids: Vec<u64>) {
413        self.update(move |current| current.override_routable_ids(ids), true);
414    }
415}
416
417#[derive(Clone, Debug)]
418pub struct Client {
419    // This is me
420    pub endpoint: Endpoint,
421    // Shared endpoint discovery source backing both snapshots and raw events.
422    endpoint_discovery_source: Arc<EndpointDiscoverySource>,
423    // These are the remotes I know about from watching key-value store
424    pub instance_source: Arc<tokio::sync::watch::Receiver<Vec<Instance>>>,
425    // Immutable routing snapshot. Free IDs are derived from discovered IDs and overloaded IDs.
426    routing_instances: Arc<RoutingInstancesState>,
427    // Client clones and standalone watchers jointly own the reconciliation task.
428    instance_avail_owner: Arc<tokio::sync::watch::Receiver<Vec<u64>>>,
429    /// Interval for periodic reconciliation of instance_avail with instance_source.
430    /// This ensures instances removed via `report_instance_down` are eventually restored.
431    /// A zero value disables local worker inhibition.
432    reconcile_interval: Duration,
433}
434
435impl Client {
436    // Client with auto-discover instances using key-value store
437    pub(crate) async fn new(endpoint: Endpoint) -> Result<Self> {
438        Self::with_reconcile_interval(endpoint, *INHIBITED_DURATION).await
439    }
440
441    /// Like [`Self::new`], but the `monitor_instance_source` background task
442    /// is bound to `cancel_token` instead of the process-wide primary token.
443    /// See [`Self::with_reconcile_interval_and_cancellation`] for why a
444    /// caller whose own scope is narrower than the process needs this.
445    pub(crate) async fn with_cancellation(
446        endpoint: Endpoint,
447        cancel_token: tokio_util::sync::CancellationToken,
448    ) -> Result<Self> {
449        Self::with_reconcile_interval_and_cancellation(endpoint, *INHIBITED_DURATION, cancel_token)
450            .await
451    }
452
453    /// Create a client with a custom reconcile interval.
454    /// The reconcile interval controls how often `instance_avail` is reset to match
455    /// `instance_source`, restoring any instances removed via `report_instance_down`.
456    pub(crate) async fn with_reconcile_interval(
457        endpoint: Endpoint,
458        reconcile_interval: Duration,
459    ) -> Result<Self> {
460        let cancel_token = endpoint.drt().primary_token();
461        Self::with_reconcile_interval_and_cancellation(endpoint, reconcile_interval, cancel_token)
462            .await
463    }
464
465    /// Like [`Self::with_reconcile_interval`], but the `monitor_instance_source`
466    /// background task is bound to `cancel_token` rather than the process-wide
467    /// primary token.
468    ///
469    /// A caller that builds a `Client` scoped to something narrower than the
470    /// process — a monitor bound to one `WorkerSet`'s lifecycle, say — must use
471    /// this constructor. `Client` is `Clone`, and `monitor_instance_source`
472    /// captures its own clone before returning, so dropping every `Client`
473    /// handle the caller holds does not stop that task; only cancelling its
474    /// token does. Built through [`Self::new`] or [`Self::with_reconcile_interval`]
475    /// instead, that task runs until process shutdown regardless of how long
476    /// the caller actually keeps the `Client` around.
477    pub(crate) async fn with_reconcile_interval_and_cancellation(
478        endpoint: Endpoint,
479        reconcile_interval: Duration,
480        cancel_token: tokio_util::sync::CancellationToken,
481    ) -> Result<Self> {
482        tracing::trace!(
483            "Client::new_dynamic: Creating dynamic client for endpoint: {}",
484            endpoint.id()
485        );
486        let endpoint_discovery_source =
487            Self::get_or_create_dynamic_discovery_source(&endpoint).await?;
488        let instance_source = Arc::new(endpoint_discovery_source.instance_receiver());
489
490        // Seed instance_avail from the current instance_source snapshot so that
491        // callers who proceed immediately after wait_for_instances (which reads
492        // instance_source directly) will also find instances in instance_avail
493        // (which is read by the routing methods like random/round_robin).
494        let initial_ids: Vec<u64> = instance_source
495            .borrow()
496            .iter()
497            .map(|instance| instance.id())
498            .collect();
499        let (routing_instances, instance_avail_owner) = RoutingInstancesState::new(initial_ids);
500        let client = Client {
501            endpoint: endpoint.clone(),
502            endpoint_discovery_source,
503            instance_source: instance_source.clone(),
504            routing_instances: Arc::new(routing_instances),
505            instance_avail_owner: Arc::new(instance_avail_owner),
506            reconcile_interval,
507        };
508        client.monitor_instance_source_with_cancellation(cancel_token, true);
509        Ok(client)
510    }
511
512    /// Instances available from watching key-value store
513    pub fn instances(&self) -> Vec<Instance> {
514        self.instance_source.borrow().clone()
515    }
516
517    pub fn instance_ids(&self) -> Vec<u64> {
518        self.instances().into_iter().map(|ep| ep.id()).collect()
519    }
520
521    /// Whether the latest discovery snapshot contains this instance, including inhibited workers.
522    pub fn is_instance_discovered(&self, instance_id: u64) -> bool {
523        self.routing_instances
524            .snapshot()
525            .discovered_ids()
526            .binary_search(&instance_id)
527            .is_ok()
528    }
529
530    pub fn instance_ids_avail(&self) -> Vec<u64> {
531        self.routing_instances.routable_ids()
532    }
533
534    /// Routable instance ids excluding those currently flagged overloaded — the set used
535    /// for load-aware (random / round-robin) worker selection.
536    pub fn instance_ids_free(&self) -> Vec<u64> {
537        self.routing_instances.free_ids()
538    }
539
540    pub(crate) fn routing_instances(&self) -> arc_swap::Guard<Arc<RoutingInstances>> {
541        self.routing_instances.snapshot()
542    }
543
544    pub fn routing_instance_counts(&self) -> RoutingInstanceCounts {
545        self.routing_instances.counts()
546    }
547
548    /// Get a watcher for available instance IDs
549    pub fn instance_avail_watcher(&self) -> tokio::sync::watch::Receiver<Vec<u64>> {
550        self.instance_avail_owner.as_ref().clone()
551    }
552
553    /// Create a client view whose routable instances are restricted by a caller-owned
554    /// admission set.
555    ///
556    /// Endpoint discovery remains the source of connection metadata and hard availability. The
557    /// returned client publishes only the intersection of that endpoint membership and
558    /// `admitted_ids`, allowing a higher-level controller to keep discovered-but-unvalidated
559    /// instances out of a routing group. The view has independent overload and fault-inhibition
560    /// state, just like a freshly constructed client.
561    pub fn with_admitted_instances(
562        &self,
563        admitted_ids: tokio::sync::watch::Receiver<Vec<u64>>,
564    ) -> Self {
565        self.with_admitted_instances_and_cancellation(
566            admitted_ids,
567            self.endpoint.drt().primary_token(),
568        )
569    }
570
571    /// Like [`Self::with_admitted_instances`], with a lifecycle token for construction-time
572    /// cancellation by an owning controller.
573    pub fn with_admitted_instances_and_cancellation(
574        &self,
575        mut admitted_ids: tokio::sync::watch::Receiver<Vec<u64>>,
576        cancel_token: tokio_util::sync::CancellationToken,
577    ) -> Self {
578        let mut endpoint_instances = self.instance_source.as_ref().clone();
579        let initial = Self::filter_admitted_instances(
580            endpoint_instances.borrow().as_slice(),
581            admitted_ids.borrow().as_slice(),
582        );
583        let initial_ids = initial.iter().map(Instance::id).collect::<Vec<_>>();
584        let (instance_tx, instance_rx) = tokio::sync::watch::channel(initial);
585        let updater_cancel = cancel_token.clone();
586
587        tokio::spawn(async move {
588            loop {
589                tokio::select! {
590                    _ = updater_cancel.cancelled() => break,
591                    _ = instance_tx.closed() => break,
592                    result = endpoint_instances.changed() => {
593                        if result.is_err() {
594                            break;
595                        }
596                    }
597                    result = admitted_ids.changed() => {
598                        if result.is_err() {
599                            break;
600                        }
601                    }
602                }
603
604                let next = Self::filter_admitted_instances(
605                    endpoint_instances.borrow_and_update().as_slice(),
606                    admitted_ids.borrow_and_update().as_slice(),
607                );
608                let changed = *instance_tx.borrow() != next;
609                if changed && instance_tx.send(next).is_err() {
610                    break;
611                }
612            }
613        });
614
615        let (routing_instances, instance_avail_owner) = RoutingInstancesState::new(initial_ids);
616        let client = Self {
617            endpoint: self.endpoint.clone(),
618            endpoint_discovery_source: self.endpoint_discovery_source.clone(),
619            instance_source: Arc::new(instance_rx),
620            routing_instances: Arc::new(routing_instances),
621            instance_avail_owner: Arc::new(instance_avail_owner),
622            reconcile_interval: self.reconcile_interval,
623        };
624        client.monitor_instance_source_with_cancellation(cancel_token, false);
625        client
626    }
627
628    fn filter_admitted_instances(instances: &[Instance], admitted_ids: &[u64]) -> Vec<Instance> {
629        if admitted_ids.is_empty() {
630            return Vec::new();
631        }
632
633        let admitted = admitted_ids.iter().copied().collect::<HashSet<_>>();
634        instances
635            .iter()
636            .filter(|instance| admitted.contains(&instance.id()))
637            .cloned()
638            .collect()
639    }
640
641    /// Subscribe to raw discovery events for this endpoint.
642    ///
643    /// Unlike `instance_source`, this feed does not coalesce remove→add pairs,
644    /// so consumers can react to every removal event exactly once.
645    pub(crate) fn subscribe_discovery_events(&self) -> DiscoveryEventReceiver {
646        DiscoveryEventReceiver {
647            receiver: self.endpoint_discovery_source.subscribe_events(),
648            _source: self.endpoint_discovery_source.clone(),
649        }
650    }
651
652    /// Wait for at least one Instance to be available for this Endpoint
653    pub async fn wait_for_instances(&self) -> Result<Vec<Instance>> {
654        tracing::trace!(
655            "wait_for_instances: Starting wait for endpoint: {}",
656            self.endpoint.id()
657        );
658        let mut rx = self.instance_source.as_ref().clone();
659        // wait for there to be 1 or more endpoints
660        let mut instances: Vec<Instance>;
661        loop {
662            instances = rx.borrow_and_update().to_vec();
663            if instances.is_empty() {
664                rx.changed().await?;
665            } else {
666                tracing::info!(
667                    "wait_for_instances: Found {} instance(s) for endpoint: {}",
668                    instances.len(),
669                    self.endpoint.id()
670                );
671                break;
672            }
673        }
674        Ok(instances)
675    }
676
677    /// Mark an instance as down/unavailable
678    pub fn report_instance_down(&self, instance_id: u64) {
679        if self.reconcile_interval.is_zero() {
680            tracing::debug!(
681                instance_id,
682                "local worker inhibition is disabled; leaving instance routable"
683            );
684            return;
685        }
686
687        self.routing_instances.report_instance_down(instance_id);
688        tracing::debug!("inhibiting instance {instance_id}");
689    }
690
691    /// Replace the set of overloaded instances reported by the worker monitor.
692    /// Returns true when this changes the routing snapshot.
693    pub fn set_overloaded_instances(&self, overloaded_instance_ids: &[u64]) -> bool {
694        self.routing_instances
695            .set_overloaded_instances(overloaded_instance_ids)
696    }
697
698    /// Whether request-path backpressure changed overload state after the monitor's
699    /// most recent metric publication.
700    pub fn overload_reconciliation_needed(&self) -> bool {
701        self.routing_instances.overload_reconciliation_needed()
702    }
703
704    /// Mark an instance overloaded immediately after a worker-scoped
705    /// `WorkerOverloaded` response. This is backpressure, not a fault, so it
706    /// does not call `report_instance_down`. The next worker-monitor
707    /// reconciliation replaces this short-lived global routing hint.
708    pub fn mark_overloaded_immediate(&self, instance_id: u64) {
709        self.routing_instances
710            .mark_overloaded_immediate(instance_id);
711        tracing::debug!(
712            instance_id,
713            "marking instance overloaded (backpressure); next metric event will re-evaluate"
714        );
715    }
716
717    pub fn clear_overloaded_instances_for_removed(&self, removed_instance_ids: &[u64]) {
718        self.routing_instances
719            .clear_overloaded_for_removed(removed_instance_ids);
720    }
721
722    pub fn overloaded_instance_ids(&self) -> Option<HashSet<u64>> {
723        self.routing_instances.overloaded_ids()
724    }
725
726    /// Workers currently eligible for selection: discovered and not locally
727    /// inhibited by [`Self::report_instance_down`].
728    ///
729    /// This hard-availability snapshot is separate from transient overload.
730    /// `None` means this client has not discovered an instance yet. After the
731    /// first discovery, `Some` is authoritative, including `Some(empty)` when
732    /// the last previously discovered worker is removed.
733    pub fn available_instance_ids(&self) -> Option<Arc<HashSet<u64>>> {
734        self.routing_instances.available_ids()
735    }
736
737    /// Monitor the key-value instance source and update instance_avail.
738    ///
739    /// This function also performs periodic reconciliation: if `instance_source` hasn't
740    /// changed for `reconcile_interval`, we reset `instance_avail` to match
741    /// `instance_source`. This ensures instances removed via `report_instance_down`
742    /// are eventually restored even if the discovery source doesn't emit updates.
743    ///
744    /// The spawned task runs until `cancel_token` cancels. A caller that wants
745    /// this task to outlive nothing shorter than the process should pass
746    /// `self.endpoint.drt().primary_token()`, as [`Self::new`] does.
747    fn monitor_instance_source_with_cancellation(
748        &self,
749        cancel_token: tokio_util::sync::CancellationToken,
750        prune_shared_occupancy: bool,
751    ) {
752        let reconcile_interval = self.reconcile_interval;
753        let endpoint = self.endpoint.clone();
754        let endpoint_discovery_source = self.endpoint_discovery_source.clone();
755        let routing_instances = self.routing_instances.clone();
756        let instance_source = self.instance_source.clone();
757        let endpoint_id = self.endpoint.id();
758        tokio::task::spawn(async move {
759            let mut rx = instance_source.as_ref().clone();
760            while !cancel_token.is_cancelled() {
761                let instance_ids: Vec<u64> = rx
762                    .borrow_and_update()
763                    .iter()
764                    .map(|instance| instance.id())
765                    .collect();
766
767                let snapshot = routing_instances.reconcile_discovered(instance_ids);
768
769                // Clean up stale occupancy counters for instances that no longer exist.
770                if prune_shared_occupancy {
771                    let registry = endpoint.drt().routing_occupancy_states();
772                    if let Ok(registry) = registry.try_lock()
773                        && let Some(weak) = registry.get(&endpoint)
774                        && let Some(state) = weak.upgrade()
775                    {
776                        state.retain(snapshot.discovered_ids());
777                    }
778                }
779
780                tokio::select! {
781                    _ = cancel_token.cancelled() => break,
782                    _ = routing_instances.instance_avail_tx.closed() => break,
783                    result = rx.changed() => {
784                        if let Err(err) = result {
785                            tracing::error!(
786                                "monitor_instance_source: The Sender is dropped: {err}, endpoint={endpoint_id}",
787                            );
788                            cancel_token.cancel();
789                        }
790                    }
791                    _ = tokio::time::sleep(reconcile_interval), if !reconcile_interval.is_zero() => {
792                        tracing::trace!(
793                            "monitor_instance_source: periodic reconciliation for endpoint={endpoint_id}",
794                        );
795                    }
796                }
797            }
798            drop(endpoint_discovery_source);
799        });
800    }
801
802    /// Simulate a complete discovery snapshot for testing.
803    #[cfg(any(test, feature = "testing"))]
804    pub fn override_discovered_instances(&self, ids: Vec<u64>) {
805        self.reconcile_discovered_instances(ids);
806    }
807
808    /// Override routable IDs for testing while preserving discovery membership.
809    #[cfg(any(test, feature = "testing"))]
810    pub fn override_instance_avail(&self, ids: Vec<u64>) {
811        self.routing_instances.override_routable_ids(ids);
812    }
813
814    fn reconcile_discovered_instances(&self, discovered_ids: Vec<u64>) -> Arc<RoutingInstances> {
815        self.routing_instances.reconcile_discovered(discovered_ids)
816    }
817
818    async fn get_or_create_dynamic_discovery_source(
819        endpoint: &Endpoint,
820    ) -> Result<Arc<EndpointDiscoverySource>> {
821        let drt = endpoint.drt();
822        let sources = drt.endpoint_discovery_sources();
823        let mut sources = sources.lock().await;
824
825        if let Some(source) = sources.get(endpoint) {
826            if let Some(source) = source.upgrade() {
827                return Ok(source);
828            } else {
829                sources.remove(endpoint);
830            }
831        }
832
833        let discovery = drt.discovery();
834        let discovery_query = crate::discovery::DiscoveryQuery::Endpoint {
835            namespace: endpoint.component.namespace.name.clone(),
836            component: endpoint.component.name.clone(),
837            endpoint: endpoint.name.clone(),
838        };
839
840        let mut discovery_stream = discovery
841            .list_and_watch(discovery_query.clone(), None)
842            .await?;
843        let (watch_tx, watch_rx) = tokio::sync::watch::channel(vec![]);
844        let discovery_source = Arc::new(EndpointDiscoverySource::new(watch_rx));
845
846        let secondary = endpoint.component.drt.runtime().secondary().clone();
847        let discovery_source_task = Arc::downgrade(&discovery_source);
848
849        secondary.spawn(async move {
850            tracing::trace!("endpoint_watcher: Starting for discovery query: {:?}", discovery_query);
851            let mut map: HashMap<u64, Instance> = HashMap::new();
852
853            loop {
854                let discovery_event = tokio::select! {
855                    _ = watch_tx.closed() => {
856                        break;
857                    }
858                    discovery_event = discovery_stream.next() => {
859                        match discovery_event {
860                            Some(Ok(event)) => {
861                                event
862                            },
863                            Some(Err(e)) => {
864                                tracing::error!("endpoint_watcher: discovery stream error: {}; shutting down for discovery query: {:?}", e, discovery_query);
865                                break;
866                            }
867                            None => {
868                                break;
869                            }
870                        }
871                    }
872                };
873
874                if let Some(discovery_source) = discovery_source_task.upgrade() {
875                    discovery_source.broadcast_event(&discovery_event);
876                }
877
878                match discovery_event {
879                    DiscoveryEvent::Added(DiscoveryInstance::Endpoint(instance)) => {
880                        map.insert(instance.instance_id, instance);
881                    }
882                    DiscoveryEvent::Added(_) => {}
883                    DiscoveryEvent::ModelTaintsUpdated(_) => {}
884                    DiscoveryEvent::Removed(id) => {
885                        if let DiscoveryInstanceId::Endpoint(endpoint_id) = id {
886                            map.remove(&endpoint_id.instance_id);
887                        }
888                    }
889                }
890
891                let instances: Vec<Instance> = map.values().cloned().collect();
892                if watch_tx.send(instances).is_err() {
893                    break;
894                }
895            }
896            let _ = watch_tx.send(vec![]);
897        });
898
899        sources.insert(endpoint.clone(), Arc::downgrade(&discovery_source));
900        Ok(discovery_source)
901    }
902}
903
904#[cfg(test)]
905mod tests {
906    use super::*;
907    use crate::{DistributedRuntime, Runtime, distributed::DistributedConfig};
908
909    async fn wait_for_discovery_event(
910        receiver: &mut DiscoveryEventReceiver,
911        predicate: impl Fn(&DiscoveryEvent) -> bool,
912    ) {
913        tokio::time::timeout(Duration::from_secs(1), async {
914            loop {
915                let event = receiver.recv().await.expect("discovery event feed closed");
916                if predicate(&event) {
917                    return;
918                }
919            }
920        })
921        .await
922        .expect("expected discovery event was not received");
923    }
924
925    async fn wait_for_watch_state<T>(
926        receiver: &mut tokio::sync::watch::Receiver<Vec<T>>,
927        predicate: impl Fn(&[T]) -> bool,
928    ) {
929        tokio::time::timeout(Duration::from_secs(1), async {
930            loop {
931                if predicate(receiver.borrow_and_update().as_slice()) {
932                    return;
933                }
934                receiver
935                    .changed()
936                    .await
937                    .expect("instance availability feed closed");
938            }
939        })
940        .await
941        .expect("expected instance availability state was not observed");
942    }
943
944    #[test]
945    fn test_inhibited_duration_from_env() {
946        assert_eq!(
947            inhibited_duration_from_env(|_| None),
948            Duration::from_secs(DEFAULT_INHIBITED_DURATION_SECS)
949        );
950        assert_eq!(
951            inhibited_duration_from_env(|_| Some("17".to_string())),
952            Duration::from_secs(17)
953        );
954        assert_eq!(
955            inhibited_duration_from_env(|_| Some("0".to_string())),
956            Duration::ZERO
957        );
958        assert_eq!(
959            inhibited_duration_from_env(|_| Some("invalid".to_string())),
960            Duration::from_secs(DEFAULT_INHIBITED_DURATION_SECS)
961        );
962    }
963
964    #[tokio::test]
965    async fn dropping_last_client_releases_routing_state() {
966        let rt = Runtime::from_current().unwrap();
967        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
968            .await
969            .unwrap();
970        let endpoint = drt
971            .namespace("test_client_lifecycle".to_string())
972            .unwrap()
973            .component("test_component".to_string())
974            .unwrap()
975            .endpoint("decode".to_string());
976        let client = endpoint.client().await.unwrap();
977        let routing_instances = Arc::downgrade(&client.routing_instances);
978        let discovery_source = Arc::downgrade(&client.endpoint_discovery_source);
979        let mut raw_watcher = client.instance_source.as_ref().clone();
980        let mut watcher = client.instance_avail_watcher();
981        let mut event_watcher = client.subscribe_discovery_events();
982        raw_watcher.borrow_and_update();
983        watcher.borrow_and_update();
984
985        drop(client);
986        assert!(routing_instances.upgrade().is_some());
987        assert!(discovery_source.upgrade().is_some());
988
989        endpoint.register_endpoint_instance().await.unwrap();
990        wait_for_watch_state(&mut watcher, |instances| instances.len() == 1).await;
991        wait_for_discovery_event(&mut event_watcher, |event| {
992            matches!(event, DiscoveryEvent::Added(DiscoveryInstance::Endpoint(_)))
993        })
994        .await;
995
996        endpoint.unregister_endpoint_instance().await.unwrap();
997        wait_for_watch_state(&mut watcher, |instances| instances.is_empty()).await;
998        assert!(raw_watcher.borrow_and_update().is_empty());
999        wait_for_discovery_event(&mut event_watcher, |event| {
1000            matches!(event, DiscoveryEvent::Removed(_))
1001        })
1002        .await;
1003
1004        drop(watcher);
1005
1006        tokio::time::timeout(Duration::from_secs(1), async {
1007            while routing_instances.strong_count() != 0 {
1008                tokio::task::yield_now().await;
1009            }
1010        })
1011        .await
1012        .expect("client monitor retained state after its last observer was dropped");
1013        assert!(discovery_source.upgrade().is_some());
1014
1015        endpoint.register_endpoint_instance().await.unwrap();
1016        wait_for_watch_state(&mut raw_watcher, |instances| instances.len() == 1).await;
1017        wait_for_discovery_event(&mut event_watcher, |event| {
1018            matches!(event, DiscoveryEvent::Added(DiscoveryInstance::Endpoint(_)))
1019        })
1020        .await;
1021        endpoint.unregister_endpoint_instance().await.unwrap();
1022        wait_for_watch_state(&mut raw_watcher, |instances| instances.is_empty()).await;
1023        wait_for_discovery_event(&mut event_watcher, |event| {
1024            matches!(event, DiscoveryEvent::Removed(_))
1025        })
1026        .await;
1027
1028        drop(event_watcher);
1029        tokio::time::timeout(Duration::from_secs(1), async {
1030            while discovery_source.strong_count() != 0 {
1031                tokio::task::yield_now().await;
1032            }
1033        })
1034        .await
1035        .expect("discovery event receiver retained its source after being dropped");
1036
1037        rt.shutdown();
1038    }
1039
1040    /// Test that instances removed via report_instance_down are restored after
1041    /// the reconciliation interval elapses.
1042    #[tokio::test]
1043    async fn test_instance_reconciliation() {
1044        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(100);
1045
1046        let rt = Runtime::from_current().unwrap();
1047        // Use process_local config to avoid needing etcd/nats
1048        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1049            .await
1050            .unwrap();
1051        let ns = drt.namespace("test_reconciliation".to_string()).unwrap();
1052        let component = ns.component("test_component".to_string()).unwrap();
1053        let endpoint = component.endpoint("test_endpoint".to_string());
1054
1055        // Use a short reconcile interval for faster tests
1056        let client = Client::with_reconcile_interval(endpoint, TEST_RECONCILE_INTERVAL)
1057            .await
1058            .unwrap();
1059
1060        // Initially, instance_avail should be empty (no registered instances)
1061        assert!(client.instance_ids_avail().is_empty());
1062
1063        // For this test, we'll directly manipulate instance_avail and verify reconciliation
1064        // Store some test IDs
1065        client.override_instance_avail(vec![1, 2, 3]);
1066
1067        assert_eq!(client.instance_ids_avail(), vec![1u64, 2, 3]);
1068
1069        // Simulate report_instance_down removing instance 2
1070        client.report_instance_down(2);
1071        assert_eq!(client.instance_ids_avail(), vec![1u64, 3]);
1072
1073        // Wait for reconciliation interval + buffer
1074        // The monitor_instance_source will reset instance_avail to match instance_source
1075        // Since instance_source is empty, after reconciliation instance_avail should be empty
1076        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;
1077
1078        // After reconciliation, instance_avail should match instance_source (which is empty)
1079        assert!(
1080            client.instance_ids_avail().is_empty(),
1081            "After reconciliation, instance_avail should match instance_source"
1082        );
1083
1084        rt.shutdown();
1085    }
1086
1087    /// A zero inhibited duration disables local worker inhibition.
1088    #[tokio::test]
1089    async fn test_zero_inhibited_duration_leaves_instance_routable() {
1090        let rt = Runtime::from_current().unwrap();
1091        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1092            .await
1093            .unwrap();
1094        let ns = drt
1095            .namespace("test_disabled_inhibition".to_string())
1096            .unwrap();
1097        let component = ns.component("test_component".to_string()).unwrap();
1098        let endpoint = component.endpoint("test_endpoint".to_string());
1099
1100        let client = Client::with_reconcile_interval(endpoint, Duration::ZERO)
1101            .await
1102            .unwrap();
1103
1104        client.override_instance_avail(vec![1, 2, 3]);
1105        client.report_instance_down(2);
1106
1107        assert_eq!(
1108            client.instance_ids_avail(),
1109            vec![1, 2, 3],
1110            "a zero inhibited duration should leave the reported instance routable"
1111        );
1112
1113        rt.shutdown();
1114    }
1115
1116    /// Test that report_instance_down correctly removes an instance from instance_avail.
1117    #[tokio::test]
1118    async fn test_report_instance_down() {
1119        let rt = Runtime::from_current().unwrap();
1120        // Use process_local config to avoid needing etcd/nats
1121        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1122            .await
1123            .unwrap();
1124        let ns = drt.namespace("test_report_down".to_string()).unwrap();
1125        let component = ns.component("test_component".to_string()).unwrap();
1126        let endpoint = component.endpoint("test_endpoint".to_string());
1127
1128        let client = endpoint.client().await.unwrap();
1129
1130        // Manually set up instance_avail with test instances
1131        client.override_instance_avail(vec![1, 2, 3]);
1132        assert_eq!(client.instance_ids_avail(), vec![1u64, 2, 3]);
1133
1134        // Report instance 2 as down
1135        client.report_instance_down(2);
1136
1137        // Verify instance 2 is removed
1138        let avail = client.instance_ids_avail();
1139        assert!(avail.contains(&1), "Instance 1 should still be available");
1140        assert!(
1141            !avail.contains(&2),
1142            "Instance 2 should be removed after report_instance_down"
1143        );
1144        assert!(avail.contains(&3), "Instance 3 should still be available");
1145
1146        rt.shutdown();
1147    }
1148
1149    #[tokio::test]
1150    async fn test_overloaded_instance_ids_returns_none_when_empty() {
1151        let rt = Runtime::from_current().unwrap();
1152        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1153            .await
1154            .unwrap();
1155        let ns = drt.namespace("test_overloaded_ids".to_string()).unwrap();
1156        let component = ns.component("test_component".to_string()).unwrap();
1157        let endpoint = component.endpoint("test_endpoint".to_string());
1158        let client = endpoint.client().await.unwrap();
1159
1160        assert_eq!(client.overloaded_instance_ids(), None);
1161
1162        assert!(client.set_overloaded_instances(&[7]));
1163        assert_eq!(client.overloaded_instance_ids(), Some(HashSet::from([7])));
1164        assert!(!client.set_overloaded_instances(&[7]));
1165
1166        assert!(client.set_overloaded_instances(&[]));
1167        assert_eq!(client.overloaded_instance_ids(), None);
1168        assert!(!client.set_overloaded_instances(&[]));
1169
1170        rt.shutdown();
1171    }
1172
1173    #[tokio::test]
1174    async fn test_instance_reconciliation_preserves_overloaded_existing_instances() {
1175        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1176
1177        let rt = Runtime::from_current().unwrap();
1178        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1179            .await
1180            .unwrap();
1181        let ns = drt
1182            .namespace("test_overloaded_reconciliation".to_string())
1183            .unwrap();
1184        let component = ns.component("test_component".to_string()).unwrap();
1185        let endpoint = component.endpoint("test_endpoint".to_string());
1186
1187        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1188            .await
1189            .unwrap();
1190        endpoint.register_endpoint_instance().await.unwrap();
1191        let instances = client.wait_for_instances().await.unwrap();
1192        let worker_id = instances[0].id();
1193
1194        for _ in 0..10 {
1195            if client.instance_ids_free().contains(&worker_id) {
1196                break;
1197            }
1198            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1199        }
1200        assert!(
1201            client.instance_ids_free().contains(&worker_id),
1202            "worker should be free after initial discovery reconciliation"
1203        );
1204
1205        client.set_overloaded_instances(&[worker_id]);
1206        assert!(
1207            client.instance_ids_free().is_empty(),
1208            "worker should be overloaded before periodic reconciliation"
1209        );
1210
1211        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;
1212
1213        assert!(
1214            client.instance_ids_free().is_empty(),
1215            "periodic reconciliation should not mark an existing overloaded worker free"
1216        );
1217
1218        rt.shutdown();
1219    }
1220
1221    #[tokio::test]
1222    async fn test_report_instance_down_preserves_overloaded_state() {
1223        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1224
1225        let rt = Runtime::from_current().unwrap();
1226        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1227            .await
1228            .unwrap();
1229        let ns = drt
1230            .namespace("test_report_down_preserves_overloaded".to_string())
1231            .unwrap();
1232        let component = ns.component("test_component".to_string()).unwrap();
1233        let endpoint = component.endpoint("test_endpoint".to_string());
1234
1235        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1236            .await
1237            .unwrap();
1238        endpoint.register_endpoint_instance().await.unwrap();
1239        let instances = client.wait_for_instances().await.unwrap();
1240        let worker_id = instances[0].id();
1241
1242        for _ in 0..10 {
1243            if client.instance_ids_avail().contains(&worker_id) {
1244                break;
1245            }
1246            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1247        }
1248
1249        client.set_overloaded_instances(&[worker_id]);
1250        client.report_instance_down(worker_id);
1251
1252        assert!(
1253            !client.instance_ids_avail().contains(&worker_id),
1254            "reported-down worker should leave routable availability"
1255        );
1256        assert_eq!(
1257            client.routing_instance_counts().overloaded,
1258            1,
1259            "reported-down worker should remain overloaded while still discovered"
1260        );
1261        assert!(
1262            client.instance_ids_free().is_empty(),
1263            "reported-down overloaded worker should not become free"
1264        );
1265
1266        endpoint.unregister_endpoint_instance().await.unwrap();
1267        for _ in 0..10 {
1268            if client.routing_instance_counts().overloaded == 0 {
1269                break;
1270            }
1271            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1272        }
1273
1274        assert_eq!(
1275            client.routing_instance_counts().overloaded,
1276            0,
1277            "stable discovery removal should clear overloaded state"
1278        );
1279
1280        rt.shutdown();
1281    }
1282
1283    #[tokio::test]
1284    async fn test_instance_reconciliation_prunes_removed_overloaded_instances() {
1285        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1286
1287        let rt = Runtime::from_current().unwrap();
1288        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1289            .await
1290            .unwrap();
1291        let ns = drt
1292            .namespace("test_removed_overloaded_cleanup".to_string())
1293            .unwrap();
1294        let component = ns.component("test_component".to_string()).unwrap();
1295        let endpoint = component.endpoint("test_endpoint".to_string());
1296
1297        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1298            .await
1299            .unwrap();
1300        endpoint.register_endpoint_instance().await.unwrap();
1301        let instances = client.wait_for_instances().await.unwrap();
1302        let worker_id = instances[0].id();
1303
1304        client.set_overloaded_instances(&[worker_id]);
1305        assert_eq!(client.routing_instance_counts().overloaded, 1);
1306        assert!(client.instance_ids_free().is_empty());
1307
1308        endpoint.unregister_endpoint_instance().await.unwrap();
1309        for _ in 0..10 {
1310            if client.routing_instance_counts().overloaded == 0 {
1311                break;
1312            }
1313            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1314        }
1315
1316        assert_eq!(
1317            client.routing_instance_counts().overloaded,
1318            0,
1319            "removed discovered workers should not remain in overloaded state"
1320        );
1321
1322        rt.shutdown();
1323    }
1324
1325    #[tokio::test]
1326    async fn test_instance_ids_free_excludes_overloaded_new_instances() {
1327        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1328
1329        let rt = Runtime::from_current().unwrap();
1330        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1331            .await
1332            .unwrap();
1333        let worker_id = drt.connection_id();
1334        let ns = drt
1335            .namespace("test_new_overloaded_reconciliation".to_string())
1336            .unwrap();
1337        let component = ns.component("test_component".to_string()).unwrap();
1338        let endpoint = component.endpoint("test_endpoint".to_string());
1339
1340        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1341            .await
1342            .unwrap();
1343        client.set_overloaded_instances(&[worker_id]);
1344
1345        endpoint.register_endpoint_instance().await.unwrap();
1346        let instances = client.wait_for_instances().await.unwrap();
1347        assert_eq!(instances[0].id(), worker_id);
1348        assert!(
1349            client.instance_ids_free().is_empty(),
1350            "newly discovered overloaded worker should not be free"
1351        );
1352
1353        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;
1354
1355        assert!(
1356            client.instance_ids_free().is_empty(),
1357            "discovery reconciliation should not affect recomputed free workers"
1358        );
1359
1360        rt.shutdown();
1361    }
1362
1363    #[tokio::test]
1364    async fn test_discovery_add_updates_free_without_overloaded_publish() {
1365        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1366
1367        let rt = Runtime::from_current().unwrap();
1368        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1369            .await
1370            .unwrap();
1371        let ns = drt
1372            .namespace("test_free_updates_on_discovery_add".to_string())
1373            .unwrap();
1374        let component = ns.component("test_component".to_string()).unwrap();
1375        let endpoint = component.endpoint("test_endpoint".to_string());
1376
1377        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1378            .await
1379            .unwrap();
1380        endpoint.register_endpoint_instance().await.unwrap();
1381        let instances = client.wait_for_instances().await.unwrap();
1382        let worker_id = instances[0].id();
1383
1384        for _ in 0..10 {
1385            if client.instance_ids_free().contains(&worker_id) {
1386                break;
1387            }
1388            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1389        }
1390
1391        assert_eq!(
1392            client.instance_ids_free(),
1393            vec![worker_id],
1394            "newly discovered non-overloaded workers should appear free without an overload update"
1395        );
1396
1397        rt.shutdown();
1398    }
1399
1400    /// Test that instance_avail_watcher receives updates when instances change.
1401    #[tokio::test]
1402    async fn test_instance_avail_watcher() {
1403        let rt = Runtime::from_current().unwrap();
1404        // Use process_local config to avoid needing etcd/nats
1405        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1406            .await
1407            .unwrap();
1408        let ns = drt.namespace("test_watcher".to_string()).unwrap();
1409        let component = ns.component("test_component".to_string()).unwrap();
1410        let endpoint = component.endpoint("test_endpoint".to_string());
1411
1412        let client = endpoint.client().await.unwrap();
1413        let watcher = client.instance_avail_watcher();
1414
1415        // Set initial instances
1416        client.override_instance_avail(vec![1, 2, 3]);
1417
1418        // Report instance down - this should notify the watcher
1419        client.report_instance_down(2);
1420
1421        // The watcher should receive the update
1422        // Note: We need to check if changed() was signaled
1423        let current = watcher.borrow().clone();
1424        assert_eq!(current, vec![1, 3]);
1425
1426        rt.shutdown();
1427    }
1428
1429    /// Regression test: `monitor_instance_source_with_cancellation`'s task must
1430    /// exit on its own `cancel_token`, not only at process shutdown.
1431    ///
1432    /// `Client::new` bound this task to the process-wide primary token
1433    /// unconditionally. A caller building a `Client` scoped to something
1434    /// narrower — a monitor bound to one `WorkerSet`'s lifecycle, say — had no
1435    /// way to stop the task before then: dropping every `Client` handle does
1436    /// not stop it, since it holds its own clone. Every WorkerSet rebuild
1437    /// leaked one.
1438    ///
1439    /// The observable is the strong count of `routing_instances`: the spawned
1440    /// task captures a clone of it, so the count returning to 1 proves the
1441    /// task actually exited and dropped that capture.
1442    #[tokio::test]
1443    async fn monitor_instance_source_exits_on_its_own_cancellation_token() {
1444        let rt = Runtime::from_current().unwrap();
1445        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1446            .await
1447            .unwrap();
1448        let ns = drt
1449            .namespace("test_monitor_instance_source_cancellation".to_string())
1450            .unwrap();
1451        let component = ns.component("test_component".to_string()).unwrap();
1452        let endpoint = component.endpoint("test_endpoint".to_string());
1453
1454        let cancel_token = tokio_util::sync::CancellationToken::new();
1455        let client = Client::with_cancellation(endpoint.clone(), cancel_token.clone())
1456            .await
1457            .unwrap();
1458
1459        // Negative control, first: the task must still be alive, and still
1460        // holding its capture, before cancellation.
1461        tokio::time::sleep(Duration::from_millis(200)).await;
1462        assert!(
1463            Arc::strong_count(&client.routing_instances) > 1,
1464            "monitor task must be running (and holding its capture) before cancellation"
1465        );
1466
1467        cancel_token.cancel();
1468        tokio::time::timeout(Duration::from_secs(5), async {
1469            while Arc::strong_count(&client.routing_instances) > 1 {
1470                tokio::time::sleep(Duration::from_millis(10)).await;
1471            }
1472        })
1473        .await
1474        .expect("monitor_instance_source task must exit when its cancel_token cancels");
1475
1476        rt.shutdown();
1477    }
1478
1479    #[tokio::test]
1480    async fn admitted_client_never_routes_unadmitted_endpoint_instances() {
1481        let rt = Runtime::from_current().unwrap();
1482        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1483            .await
1484            .unwrap();
1485        let ns = drt.namespace("test_admitted_client".to_string()).unwrap();
1486        let component = ns.component("test_component".to_string()).unwrap();
1487        let endpoint = component.endpoint("test_endpoint".to_string());
1488        let endpoint_client = endpoint.client().await.unwrap();
1489        endpoint.register_endpoint_instance().await.unwrap();
1490        let worker_id = endpoint_client.wait_for_instances().await.unwrap()[0].id();
1491
1492        let (admission_tx, admission_rx) = tokio::sync::watch::channel(Vec::new());
1493        let admitted_client = endpoint_client.with_admitted_instances(admission_rx);
1494        let mut admitted = admitted_client.instance_avail_watcher();
1495        assert!(admitted.borrow().is_empty());
1496
1497        admission_tx.send_replace(vec![worker_id]);
1498        tokio::time::timeout(Duration::from_secs(1), admitted.changed())
1499            .await
1500            .unwrap()
1501            .unwrap();
1502        assert_eq!(admitted.borrow_and_update().clone(), vec![worker_id]);
1503
1504        admission_tx.send_replace(Vec::new());
1505        tokio::time::timeout(Duration::from_secs(1), admitted.changed())
1506            .await
1507            .unwrap()
1508            .unwrap();
1509        assert!(admitted.borrow_and_update().is_empty());
1510
1511        rt.shutdown();
1512    }
1513
1514    /// Test that concurrent select_and_increment distributes load correctly.
1515    #[tokio::test]
1516    async fn test_concurrent_select_and_increment() {
1517        let state = Arc::new(RoutingOccupancyState::default());
1518        let instance_ids: Vec<u64> = vec![100, 200, 300];
1519        let num_requests = 90;
1520
1521        let mut handles = Vec::new();
1522        for _ in 0..num_requests {
1523            let state = state.clone();
1524            let ids = instance_ids.clone();
1525            handles.push(tokio::spawn(async move {
1526                state.select_exact_min_and_increment(&ids).await
1527            }));
1528        }
1529
1530        for handle in handles {
1531            handle.await.unwrap();
1532        }
1533
1534        assert_eq!(state.load(100), 30);
1535        assert_eq!(state.load(200), 30);
1536        assert_eq!(state.load(300), 30);
1537    }
1538
1539    #[tokio::test]
1540    async fn test_select_exact_min_and_increment_randomizes_ties() {
1541        let mut selected = [false; 3];
1542
1543        for _ in 0..120 {
1544            let state = RoutingOccupancyState::default();
1545            let picked = state
1546                .select_exact_min_and_increment(&[10, 20, 30])
1547                .await
1548                .unwrap();
1549            match picked {
1550                10 => selected[0] = true,
1551                20 => selected[1] = true,
1552                30 => selected[2] = true,
1553                _ => panic!("unexpected worker id: {picked}"),
1554            }
1555        }
1556
1557        let selected_count = selected.into_iter().filter(|seen| *seen).count();
1558        assert!(
1559            selected_count > 1,
1560            "tie-breaking should not always select the first minimum-load worker"
1561        );
1562    }
1563
1564    #[tokio::test]
1565    async fn test_connection_counts() {
1566        let rt = Runtime::from_current().unwrap();
1567        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1568            .await
1569            .unwrap();
1570        let ns = drt.namespace("test_ll_counts".to_string()).unwrap();
1571        let component = ns.component("test_component".to_string()).unwrap();
1572        let endpoint = component.endpoint("test_endpoint".to_string());
1573
1574        let state1 = get_or_create_routing_occupancy_state(&endpoint).await;
1575        let state2 = get_or_create_routing_occupancy_state(&endpoint).await;
1576
1577        let picked1 = state1
1578            .select_exact_min_and_increment(&[10, 20, 30])
1579            .await
1580            .unwrap();
1581        assert_eq!(state1.load(picked1), 1);
1582
1583        let picked2 = state1
1584            .select_exact_min_and_increment(&[10, 20, 30])
1585            .await
1586            .unwrap();
1587        assert_ne!(picked1, picked2);
1588
1589        // state2 should see the same counts (same underlying Arc)
1590        assert_eq!(state2.load(10), state1.load(10));
1591        assert_eq!(state2.load(20), state1.load(20));
1592        assert_eq!(state2.load(30), state1.load(30));
1593
1594        state2.decrement(picked1);
1595        assert_eq!(state1.load(picked1), if picked1 == picked2 { 1 } else { 0 });
1596
1597        rt.shutdown();
1598    }
1599
1600    #[tokio::test]
1601    async fn test_least_loaded_state_retain_preserves_live_counts() {
1602        let state = RoutingOccupancyState::default();
1603
1604        // Add some connections
1605        state.select_exact_min_and_increment(&[1, 2, 3]).await;
1606        state.select_exact_min_and_increment(&[1, 2, 3]).await;
1607        state.select_exact_min_and_increment(&[1, 2, 3]).await;
1608        // Each instance should have 1 connection
1609        assert_eq!(state.load(1), 1);
1610        assert_eq!(state.load(2), 1);
1611        assert_eq!(state.load(3), 1);
1612
1613        // Discovery removal must not delete guard-owned accounting.
1614        state.retain(&[1, 3]);
1615
1616        assert_eq!(state.load(1), 1);
1617        assert_eq!(state.load(2), 1);
1618        assert_eq!(state.load(3), 1);
1619    }
1620
1621    #[tokio::test]
1622    async fn test_monitor_instance_source_defers_removed_worker_cleanup() {
1623        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(50);
1624
1625        let rt = Runtime::from_current().unwrap();
1626        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
1627            .await
1628            .unwrap();
1629        let ns = drt.namespace("test_occupancy_cleanup".to_string()).unwrap();
1630        let component = ns.component("test_component".to_string()).unwrap();
1631        let endpoint = component.endpoint("test_endpoint".to_string());
1632
1633        let client = Client::with_reconcile_interval(endpoint.clone(), TEST_RECONCILE_INTERVAL)
1634            .await
1635            .unwrap();
1636        endpoint.register_endpoint_instance().await.unwrap();
1637        client.wait_for_instances().await.unwrap();
1638
1639        let worker_id = client.instance_ids_avail()[0];
1640        let state = get_or_create_routing_occupancy_state(&endpoint).await;
1641        state.increment(worker_id);
1642        assert_eq!(state.load(worker_id), 1);
1643
1644        endpoint.unregister_endpoint_instance().await.unwrap();
1645
1646        for _ in 0..10 {
1647            if !client.instance_ids().contains(&worker_id) {
1648                break;
1649            }
1650            tokio::time::sleep(TEST_RECONCILE_INTERVAL).await;
1651        }
1652
1653        assert_eq!(
1654            state.load(worker_id),
1655            1,
1656            "discovery absence must retain live accounting"
1657        );
1658        state.decrement(worker_id);
1659        assert_eq!(state.load(worker_id), 0);
1660
1661        rt.shutdown();
1662    }
1663}
dynamo_runtime/component/client.rs

dynamo_runtime/component/
client.rs
