Skip to main content

toolkit/runtime/
gear_manager.rs

1//! Gear Manager - tracks and manages all live gear instances in the runtime
2
3use dashmap::DashMap;
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7use uuid::Uuid;
8
9/// Represents an endpoint where a gear instance can be reached
10#[derive(Clone, Debug, PartialEq, Eq, Hash)]
11pub struct Endpoint {
12    pub uri: String,
13}
14
15/// Typed view of an endpoint for parsing and matching
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub enum EndpointKind {
18    /// TCP endpoint with resolved socket address
19    Tcp(std::net::SocketAddr),
20    /// Unix domain socket with file path
21    Uds(std::path::PathBuf),
22    /// Other/unparsed endpoint URI
23    Other(String),
24}
25
26impl Endpoint {
27    pub fn from_uri<S: Into<String>>(s: S) -> Self {
28        Self { uri: s.into() }
29    }
30
31    pub fn uds(path: impl AsRef<std::path::Path>) -> Self {
32        Self {
33            uri: format!("unix://{}", path.as_ref().display()),
34        }
35    }
36
37    #[must_use]
38    pub fn http(host: &str, port: u16) -> Self {
39        Self {
40            uri: format!("http://{host}:{port}"),
41        }
42    }
43
44    #[must_use]
45    pub fn https(host: &str, port: u16) -> Self {
46        Self {
47            uri: format!("https://{host}:{port}"),
48        }
49    }
50
51    /// Parse the endpoint URI into a typed view
52    #[must_use]
53    pub fn kind(&self) -> EndpointKind {
54        if let Some(rest) = self.uri.strip_prefix("unix://") {
55            return EndpointKind::Uds(std::path::PathBuf::from(rest));
56        }
57        if let Some(rest) = self.uri.strip_prefix("http://")
58            && let Ok(addr) = rest.parse::<std::net::SocketAddr>()
59        {
60            return EndpointKind::Tcp(addr);
61        }
62        if let Some(rest) = self.uri.strip_prefix("https://")
63            && let Ok(addr) = rest.parse::<std::net::SocketAddr>()
64        {
65            return EndpointKind::Tcp(addr);
66        }
67        EndpointKind::Other(self.uri.clone())
68    }
69}
70
71#[derive(Clone, Copy, Debug, PartialEq, Eq)]
72pub enum InstanceState {
73    Registered,
74    Ready,
75    Healthy,
76    Quarantined,
77    Draining,
78}
79
80/// Runtime state of an instance (guarded by `RwLock` for safe mutation)
81#[derive(Clone, Debug)]
82pub struct InstanceRuntimeState {
83    pub last_heartbeat: Instant,
84    pub state: InstanceState,
85}
86
87/// Represents a single instance of a gear
88#[derive(Debug)]
89#[must_use]
90pub struct GearInstance {
91    pub gear: String,
92    pub instance_id: Uuid,
93    pub control: Option<Endpoint>,
94    pub grpc_services: HashMap<String, Endpoint>,
95    pub version: Option<String>,
96    inner: Arc<parking_lot::RwLock<InstanceRuntimeState>>,
97}
98
99impl Clone for GearInstance {
100    fn clone(&self) -> Self {
101        Self {
102            gear: self.gear.clone(),
103            instance_id: self.instance_id,
104            control: self.control.clone(),
105            grpc_services: self.grpc_services.clone(),
106            version: self.version.clone(),
107            inner: Arc::clone(&self.inner),
108        }
109    }
110}
111
112impl GearInstance {
113    pub fn new(gear: impl Into<String>, instance_id: Uuid) -> Self {
114        Self {
115            gear: gear.into(),
116            instance_id,
117            control: None,
118            grpc_services: HashMap::new(),
119            version: None,
120            inner: Arc::new(parking_lot::RwLock::new(InstanceRuntimeState {
121                last_heartbeat: Instant::now(),
122                state: InstanceState::Registered,
123            })),
124        }
125    }
126
127    pub fn with_control(mut self, ep: Endpoint) -> Self {
128        self.control = Some(ep);
129        self
130    }
131
132    pub fn with_version(mut self, v: impl Into<String>) -> Self {
133        self.version = Some(v.into());
134        self
135    }
136
137    pub fn with_grpc_service(mut self, name: impl Into<String>, ep: Endpoint) -> Self {
138        self.grpc_services.insert(name.into(), ep);
139        self
140    }
141
142    /// Get the current state of this instance
143    #[must_use]
144    pub fn state(&self) -> InstanceState {
145        self.inner.read().state
146    }
147
148    /// Get the last heartbeat timestamp
149    #[must_use]
150    pub fn last_heartbeat(&self) -> Instant {
151        self.inner.read().last_heartbeat
152    }
153}
154
155/// Central registry that tracks all running gear instances in the system.
156/// Provides discovery, health tracking, and round-robin load balancing.
157#[derive(Clone)]
158#[must_use]
159pub struct GearManager {
160    inner: DashMap<String, Vec<Arc<GearInstance>>>,
161    rr_counters: DashMap<String, usize>,
162    hb_ttl: Duration,
163    hb_grace: Duration,
164}
165
166impl std::fmt::Debug for GearManager {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        let gears: Vec<String> = self.inner.iter().map(|e| e.key().clone()).collect();
169        f.debug_struct("GearManager")
170            .field("instances_count", &self.inner.len())
171            .field("gears", &gears)
172            .field("heartbeat_ttl", &self.hb_ttl)
173            .field("heartbeat_grace", &self.hb_grace)
174            .finish_non_exhaustive()
175    }
176}
177
178impl GearManager {
179    pub fn new() -> Self {
180        Self {
181            inner: DashMap::new(),
182            rr_counters: DashMap::new(),
183            hb_ttl: Duration::from_secs(15),
184            hb_grace: Duration::from_secs(30),
185        }
186    }
187
188    pub fn with_heartbeat_policy(mut self, ttl: Duration, grace: Duration) -> Self {
189        self.hb_ttl = ttl;
190        self.hb_grace = grace;
191        self
192    }
193
194    /// Register or update a gear instance
195    pub fn register_instance(&self, instance: Arc<GearInstance>) {
196        let gear = instance.gear.clone();
197        let mut vec = self.inner.entry(gear).or_default();
198        // replace by instance_id if it already exists
199        if let Some(pos) = vec
200            .iter()
201            .position(|i| i.instance_id == instance.instance_id)
202        {
203            vec[pos] = instance;
204        } else {
205            vec.push(instance);
206        }
207    }
208
209    /// Mark an instance as ready
210    pub fn mark_ready(&self, gear: &str, instance_id: Uuid) {
211        if let Some(mut vec) = self.inner.get_mut(gear)
212            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
213        {
214            let mut state = inst.inner.write();
215            state.state = InstanceState::Ready;
216        }
217    }
218
219    /// Update the heartbeat timestamp for an instance
220    pub fn update_heartbeat(&self, gear: &str, instance_id: Uuid, at: Instant) {
221        if let Some(mut vec) = self.inner.get_mut(gear)
222            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
223        {
224            let mut state = inst.inner.write();
225            state.last_heartbeat = at;
226            // Transition Registered -> Healthy on first heartbeat
227            if state.state == InstanceState::Registered {
228                state.state = InstanceState::Healthy;
229            }
230        }
231    }
232
233    /// Mark an instance as quarantined
234    pub fn mark_quarantined(&self, gear: &str, instance_id: Uuid) {
235        if let Some(mut vec) = self.inner.get_mut(gear)
236            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
237        {
238            inst.inner.write().state = InstanceState::Quarantined;
239        }
240    }
241
242    /// Mark an instance as draining (graceful shutdown in progress)
243    pub fn mark_draining(&self, gear: &str, instance_id: Uuid) {
244        if let Some(mut vec) = self.inner.get_mut(gear)
245            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
246        {
247            inst.inner.write().state = InstanceState::Draining;
248        }
249    }
250
251    /// Remove an instance from the directory
252    pub fn deregister(&self, gear: &str, instance_id: Uuid) {
253        let mut remove_gear = false;
254        {
255            if let Some(mut vec) = self.inner.get_mut(gear) {
256                let list = vec.value_mut();
257                list.retain(|inst| inst.instance_id != instance_id);
258                if list.is_empty() {
259                    remove_gear = true;
260                }
261            }
262        }
263
264        if remove_gear {
265            self.inner.remove(gear);
266            self.rr_counters.remove(gear);
267        }
268    }
269
270    /// Get all instances of a specific gear
271    #[must_use]
272    pub fn instances_of(&self, gear: &str) -> Vec<Arc<GearInstance>> {
273        self.inner.get(gear).map(|v| v.clone()).unwrap_or_default()
274    }
275
276    /// Get all instances across all gears
277    #[must_use]
278    pub fn all_instances(&self) -> Vec<Arc<GearInstance>> {
279        self.inner
280            .iter()
281            .flat_map(|entry| entry.value().clone())
282            .collect()
283    }
284
285    /// Quarantine or evict stale instances based on heartbeat policy
286    pub fn evict_stale(&self, now: Instant) {
287        use InstanceState::{Draining, Quarantined};
288        let mut empty_gears = Vec::new();
289
290        for mut entry in self.inner.iter_mut() {
291            let gear = entry.key().clone();
292            let vec = entry.value_mut();
293            vec.retain(|inst| {
294                let state = inst.inner.read();
295                let age = now.saturating_duration_since(state.last_heartbeat);
296
297                // Quarantine instances that have exceeded TTL
298                if age >= self.hb_ttl && !matches!(state.state, Quarantined | Draining) {
299                    drop(state); // Release read lock before write
300                    inst.inner.write().state = Quarantined;
301                    return true; // Keep quarantined instances for now
302                }
303
304                // Evict quarantined instances that exceed grace period
305                if state.state == Quarantined && age >= self.hb_ttl + self.hb_grace {
306                    return false; // Remove from directory
307                }
308
309                true
310            });
311
312            if vec.is_empty() {
313                empty_gears.push(gear);
314            }
315        }
316
317        for gear in empty_gears {
318            self.inner.remove(&gear);
319            self.rr_counters.remove(&gear);
320        }
321    }
322
323    /// Pick an instance using round-robin selection, preferring healthy instances
324    #[must_use]
325    pub fn pick_instance_round_robin(&self, gear: &str) -> Option<Arc<GearInstance>> {
326        let instances_entry = self.inner.get(gear)?;
327        let instances = instances_entry.value();
328
329        if instances.is_empty() {
330            return None;
331        }
332
333        // Prefer healthy or ready instances
334        let healthy: Vec<_> = instances
335            .iter()
336            .filter(|inst| matches!(inst.state(), InstanceState::Healthy | InstanceState::Ready))
337            .cloned()
338            .collect();
339
340        let candidates: Vec<_> = if healthy.is_empty() {
341            instances.clone()
342        } else {
343            healthy
344        };
345
346        if candidates.is_empty() {
347            return None;
348        }
349
350        let len = candidates.len();
351        let mut counter = self.rr_counters.entry(gear.to_owned()).or_insert(0);
352        let idx = *counter % len;
353        *counter = (*counter + 1) % len;
354
355        candidates.get(idx).cloned()
356    }
357
358    /// Pick a service endpoint using round-robin, returning (gear, instance, endpoint).
359    /// Prefers healthy/ready instances and automatically rotates among them.
360    #[must_use]
361    pub fn pick_service_round_robin(
362        &self,
363        service_name: &str,
364    ) -> Option<(String, Arc<GearInstance>, Endpoint)> {
365        // Collect all instances that provide this service
366        let mut candidates = Vec::new();
367        for entry in &self.inner {
368            let gear = entry.key().clone();
369            for inst in entry.value() {
370                if let Some(ep) = inst.grpc_services.get(service_name) {
371                    let state = inst.state();
372                    if matches!(state, InstanceState::Healthy | InstanceState::Ready) {
373                        candidates.push((gear.clone(), inst.clone(), ep.clone()));
374                    }
375                }
376            }
377        }
378
379        if candidates.is_empty() {
380            return None;
381        }
382
383        // Use a counter keyed by service name for round-robin
384        let len = candidates.len();
385        let service_key = service_name.to_owned();
386        let mut counter = self.rr_counters.entry(service_key).or_insert(0);
387        let idx = *counter % len;
388        *counter = (*counter + 1) % len;
389
390        candidates.get(idx).cloned()
391    }
392}
393
394impl Default for GearManager {
395    fn default() -> Self {
396        Self::new()
397    }
398}
399
400#[cfg(test)]
401#[cfg_attr(coverage_nightly, coverage(off))]
402mod tests {
403    use super::*;
404    use std::thread::sleep;
405    use std::time::Duration;
406
407    #[test]
408    fn test_register_and_retrieve_instances() {
409        let dir = GearManager::new();
410        let instance_id = Uuid::new_v4();
411        let instance = Arc::new(
412            GearInstance::new("test_gear", instance_id)
413                .with_control(Endpoint::http("localhost", 8080))
414                .with_version("1.0.0"),
415        );
416
417        dir.register_instance(instance);
418
419        let instances = dir.instances_of("test_gear");
420        assert_eq!(instances.len(), 1);
421        assert_eq!(instances[0].instance_id, instance_id);
422        assert_eq!(instances[0].gear, "test_gear");
423        assert_eq!(instances[0].version, Some("1.0.0".to_owned()));
424    }
425
426    #[test]
427    fn test_register_multiple_instances() {
428        let dir = GearManager::new();
429
430        let id1 = Uuid::new_v4();
431        let id2 = Uuid::new_v4();
432        let instance1 = Arc::new(GearInstance::new("test_gear", id1));
433        let instance2 = Arc::new(GearInstance::new("test_gear", id2));
434
435        dir.register_instance(instance1);
436        dir.register_instance(instance2);
437
438        let registered = dir.instances_of("test_gear");
439        assert_eq!(registered.len(), 2);
440
441        let ids: Vec<_> = registered.iter().map(|i| i.instance_id).collect();
442        assert!(ids.contains(&id1));
443        assert!(ids.contains(&id2));
444    }
445
446    #[test]
447    fn test_update_existing_instance() {
448        let dir = GearManager::new();
449        let instance_id = Uuid::new_v4();
450
451        let initial_instance =
452            Arc::new(GearInstance::new("test_gear", instance_id).with_version("1.0.0"));
453        dir.register_instance(initial_instance);
454
455        let updated_instance =
456            Arc::new(GearInstance::new("test_gear", instance_id).with_version("2.0.0"));
457        dir.register_instance(updated_instance);
458
459        let registered = dir.instances_of("test_gear");
460        assert_eq!(registered.len(), 1, "Should not duplicate instance");
461        assert_eq!(registered[0].version, Some("2.0.0".to_owned()));
462    }
463
464    #[test]
465    fn test_mark_ready() {
466        let dir = GearManager::new();
467        let instance_id = Uuid::new_v4();
468        let instance = Arc::new(GearInstance::new("test_gear", instance_id));
469
470        dir.register_instance(instance);
471
472        dir.mark_ready("test_gear", instance_id);
473
474        let instances = dir.instances_of("test_gear");
475        assert_eq!(instances.len(), 1);
476        assert!(matches!(instances[0].state(), InstanceState::Ready));
477    }
478
479    #[test]
480    fn test_update_heartbeat() {
481        let dir = GearManager::new();
482        let instance_id = Uuid::new_v4();
483        let instance = Arc::new(GearInstance::new("test_gear", instance_id));
484        let initial_heartbeat = instance.last_heartbeat();
485
486        dir.register_instance(instance);
487
488        // Sleep to ensure time difference
489        sleep(Duration::from_millis(10));
490
491        let new_heartbeat = Instant::now();
492        dir.update_heartbeat("test_gear", instance_id, new_heartbeat);
493
494        let instances = dir.instances_of("test_gear");
495        assert!(instances[0].last_heartbeat() > initial_heartbeat);
496        assert!(matches!(instances[0].state(), InstanceState::Healthy));
497    }
498
499    #[test]
500    fn test_all_instances() {
501        let dir = GearManager::new();
502
503        let instance1 = Arc::new(GearInstance::new("gear_a", Uuid::new_v4()));
504        let instance2 = Arc::new(GearInstance::new("gear_b", Uuid::new_v4()));
505        let instance3 = Arc::new(GearInstance::new("gear_a", Uuid::new_v4()));
506
507        dir.register_instance(instance1);
508        dir.register_instance(instance2);
509        dir.register_instance(instance3);
510
511        let all = dir.all_instances();
512        assert_eq!(all.len(), 3);
513
514        let gears: Vec<_> = all.iter().map(|i| i.gear.as_str()).collect();
515        assert_eq!(gears.iter().filter(|&m| *m == "gear_a").count(), 2);
516        assert_eq!(gears.iter().filter(|&m| *m == "gear_b").count(), 1);
517    }
518
519    #[test]
520    fn test_pick_instance_round_robin() {
521        let dir = GearManager::new();
522
523        let id1 = Uuid::new_v4();
524        let id2 = Uuid::new_v4();
525        let instance1 = Arc::new(GearInstance::new("test_gear", id1));
526        let instance2 = Arc::new(GearInstance::new("test_gear", id2));
527
528        dir.register_instance(instance1);
529        dir.register_instance(instance2);
530
531        // Pick three times to verify round-robin behavior
532        let picked1 = dir.pick_instance_round_robin("test_gear").unwrap();
533        let picked2 = dir.pick_instance_round_robin("test_gear").unwrap();
534        let picked3 = dir.pick_instance_round_robin("test_gear").unwrap();
535
536        let ids = [
537            picked1.instance_id,
538            picked2.instance_id,
539            picked3.instance_id,
540        ];
541
542        // With 2 instances, we expect round-robin pattern like A, B, A
543        // Check that both instance IDs appear and that at least one repeats
544        assert!(ids.contains(&id1));
545        assert!(ids.contains(&id2));
546        // First and third pick should be the same (round-robin wraps)
547        assert_eq!(picked1.instance_id, picked3.instance_id);
548        // Second pick should be different from the first
549        assert_ne!(picked1.instance_id, picked2.instance_id);
550    }
551
552    #[test]
553    fn test_pick_instance_none_available() {
554        let dir = GearManager::new();
555        let picked = dir.pick_instance_round_robin("nonexistent_gear");
556        assert!(picked.is_none());
557    }
558
559    #[test]
560    fn test_endpoint_creation() {
561        let plain_ep = Endpoint::http("localhost", 8080);
562        assert_eq!(plain_ep.uri, "http://localhost:8080");
563
564        let secure_ep = Endpoint::https("localhost", 8443);
565        assert_eq!(secure_ep.uri, "https://localhost:8443");
566
567        let uds_ep = Endpoint::uds("/tmp/socket.sock");
568        assert!(uds_ep.uri.starts_with("unix://"));
569        assert!(uds_ep.uri.contains("socket.sock"));
570
571        let custom_ep = Endpoint::from_uri("http://example.com");
572        assert_eq!(custom_ep.uri, "http://example.com");
573    }
574
575    #[test]
576    fn test_endpoint_kind() {
577        let plain_ep = Endpoint::http("127.0.0.1", 8080);
578        match plain_ep.kind() {
579            EndpointKind::Tcp(addr) => {
580                assert_eq!(addr.ip().to_string(), "127.0.0.1");
581                assert_eq!(addr.port(), 8080);
582            }
583            _ => panic!("Expected TCP endpoint for http"),
584        }
585
586        let secure_ep = Endpoint::https("127.0.0.1", 8443);
587        match secure_ep.kind() {
588            EndpointKind::Tcp(addr) => {
589                assert_eq!(addr.ip().to_string(), "127.0.0.1");
590                assert_eq!(addr.port(), 8443);
591            }
592            _ => panic!("Expected TCP endpoint for https"),
593        }
594
595        let uds_ep = Endpoint::uds("/tmp/test.sock");
596        match uds_ep.kind() {
597            EndpointKind::Uds(path) => {
598                assert!(path.to_string_lossy().contains("test.sock"));
599            }
600            _ => panic!("Expected UDS endpoint"),
601        }
602
603        let other_ep = Endpoint::from_uri("grpc://example.com");
604        match other_ep.kind() {
605            EndpointKind::Other(uri) => {
606                assert_eq!(uri, "grpc://example.com");
607            }
608            _ => panic!("Expected Other endpoint"),
609        }
610    }
611
612    #[test]
613    fn test_gear_instance_builder() {
614        let instance_id = Uuid::new_v4();
615        let instance = GearInstance::new("test_gear", instance_id)
616            .with_control(Endpoint::http("localhost", 8080))
617            .with_version("1.2.3")
618            .with_grpc_service("service1", Endpoint::http("localhost", 8082))
619            .with_grpc_service("service2", Endpoint::http("localhost", 8083));
620
621        assert_eq!(instance.gear, "test_gear");
622        assert_eq!(instance.instance_id, instance_id);
623        assert!(instance.control.is_some());
624        assert_eq!(instance.version, Some("1.2.3".to_owned()));
625        assert_eq!(instance.grpc_services.len(), 2);
626        assert!(instance.grpc_services.contains_key("service1"));
627        assert!(instance.grpc_services.contains_key("service2"));
628        assert!(matches!(instance.state(), InstanceState::Registered));
629    }
630
631    #[test]
632    fn test_quarantine_and_evict() {
633        let ttl = Duration::from_millis(50);
634        let grace = Duration::from_millis(50);
635        let dir = GearManager::new().with_heartbeat_policy(ttl, grace);
636
637        let now = Instant::now();
638        let instance = GearInstance::new("test_gear", Uuid::new_v4());
639        // Set the last heartbeat to be stale
640        instance.inner.write().last_heartbeat = now
641            .checked_sub(ttl)
642            .and_then(|t| t.checked_sub(Duration::from_millis(10)))
643            .expect("test duration subtraction should not underflow");
644
645        dir.register_instance(Arc::new(instance));
646
647        dir.evict_stale(now);
648        let instances = dir.instances_of("test_gear");
649        assert_eq!(instances.len(), 1);
650        assert!(matches!(instances[0].state(), InstanceState::Quarantined));
651
652        let later = now + grace + Duration::from_millis(10);
653        dir.evict_stale(later);
654
655        let instances_after = dir.instances_of("test_gear");
656        assert!(instances_after.is_empty());
657    }
658
659    #[test]
660    fn test_instances_of_empty() {
661        let dir = GearManager::new();
662        let instances = dir.instances_of("nonexistent");
663        assert!(instances.is_empty());
664    }
665
666    #[test]
667    fn test_rr_prefers_healthy() {
668        let dir = GearManager::new();
669
670        // Create two instances: one healthy, one quarantined
671        let healthy_id = Uuid::new_v4();
672        let healthy = Arc::new(GearInstance::new("test_gear", healthy_id));
673        dir.register_instance(healthy);
674        dir.update_heartbeat("test_gear", healthy_id, Instant::now());
675
676        let quarantined_id = Uuid::new_v4();
677        let quarantined = Arc::new(GearInstance::new("test_gear", quarantined_id));
678        dir.register_instance(quarantined);
679        dir.mark_quarantined("test_gear", quarantined_id);
680
681        // RR should only pick the healthy instance
682        for _ in 0..5 {
683            let picked = dir.pick_instance_round_robin("test_gear").unwrap();
684            assert_eq!(picked.instance_id, healthy_id);
685        }
686    }
687
688    #[test]
689    fn test_pick_service_round_robin() {
690        let dir = GearManager::new();
691
692        let id1 = Uuid::new_v4();
693        let id2 = Uuid::new_v4();
694        // Register two instances providing the same service
695        let inst1 = Arc::new(
696            GearInstance::new("test_gear", id1)
697                .with_grpc_service("test.Service", Endpoint::http("127.0.0.1", 8001)),
698        );
699        let inst2 = Arc::new(
700            GearInstance::new("test_gear", id2)
701                .with_grpc_service("test.Service", Endpoint::http("127.0.0.1", 8002)),
702        );
703
704        dir.register_instance(inst1);
705        dir.register_instance(inst2);
706
707        // Mark both as healthy
708        dir.update_heartbeat("test_gear", id1, Instant::now());
709        dir.update_heartbeat("test_gear", id2, Instant::now());
710
711        // Pick should rotate between instances
712        let pick1 = dir.pick_service_round_robin("test.Service");
713        let pick2 = dir.pick_service_round_robin("test.Service");
714        let pick3 = dir.pick_service_round_robin("test.Service");
715
716        assert!(pick1.is_some());
717        assert!(pick2.is_some());
718        assert!(pick3.is_some());
719
720        let (_, inst1, ep1) = pick1.unwrap();
721        let (_, inst2, ep2) = pick2.unwrap();
722        let (_, inst3, _) = pick3.unwrap();
723
724        // First and third should be the same (round-robin)
725        assert_eq!(inst1.instance_id, inst3.instance_id);
726        // First and second should be different
727        assert_ne!(inst1.instance_id, inst2.instance_id);
728        // Endpoints should differ
729        assert_ne!(ep1, ep2);
730    }
731}