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    pub rest_endpoint: Option<Endpoint>,
97    pub openapi_spec: Option<String>,
98    inner: Arc<parking_lot::RwLock<InstanceRuntimeState>>,
99}
100
101impl Clone for GearInstance {
102    fn clone(&self) -> Self {
103        Self {
104            gear: self.gear.clone(),
105            instance_id: self.instance_id,
106            control: self.control.clone(),
107            grpc_services: self.grpc_services.clone(),
108            version: self.version.clone(),
109            rest_endpoint: self.rest_endpoint.clone(),
110            openapi_spec: self.openapi_spec.clone(),
111            inner: Arc::clone(&self.inner),
112        }
113    }
114}
115
116impl GearInstance {
117    pub fn new(gear: impl Into<String>, instance_id: Uuid) -> Self {
118        Self {
119            gear: gear.into(),
120            instance_id,
121            control: None,
122            grpc_services: HashMap::new(),
123            version: None,
124            rest_endpoint: None,
125            openapi_spec: None,
126            inner: Arc::new(parking_lot::RwLock::new(InstanceRuntimeState {
127                last_heartbeat: Instant::now(),
128                state: InstanceState::Registered,
129            })),
130        }
131    }
132
133    pub fn with_control(mut self, ep: Endpoint) -> Self {
134        self.control = Some(ep);
135        self
136    }
137
138    pub fn with_version(mut self, v: impl Into<String>) -> Self {
139        self.version = Some(v.into());
140        self
141    }
142
143    pub fn with_grpc_service(mut self, name: impl Into<String>, ep: Endpoint) -> Self {
144        self.grpc_services.insert(name.into(), ep);
145        self
146    }
147
148    pub fn with_rest_endpoint(mut self, ep: Endpoint) -> Self {
149        self.rest_endpoint = Some(ep);
150        self
151    }
152
153    pub fn with_openapi_spec(mut self, spec: impl Into<String>) -> Self {
154        self.openapi_spec = Some(spec.into());
155        self
156    }
157
158    /// Get the current state of this instance
159    #[must_use]
160    pub fn state(&self) -> InstanceState {
161        self.inner.read().state
162    }
163
164    /// Get the last heartbeat timestamp
165    #[must_use]
166    pub fn last_heartbeat(&self) -> Instant {
167        self.inner.read().last_heartbeat
168    }
169}
170
171/// Central registry that tracks all running gear instances in the system.
172/// Provides discovery, health tracking, and round-robin load balancing.
173#[derive(Clone)]
174#[must_use]
175pub struct GearManager {
176    inner: DashMap<String, Vec<Arc<GearInstance>>>,
177    rr_counters: DashMap<String, usize>,
178    hb_ttl: Duration,
179    hb_grace: Duration,
180}
181
182impl std::fmt::Debug for GearManager {
183    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
184        let gears: Vec<String> = self.inner.iter().map(|e| e.key().clone()).collect();
185        f.debug_struct("GearManager")
186            .field("instances_count", &self.inner.len())
187            .field("gears", &gears)
188            .field("heartbeat_ttl", &self.hb_ttl)
189            .field("heartbeat_grace", &self.hb_grace)
190            .finish_non_exhaustive()
191    }
192}
193
194impl GearManager {
195    pub fn new() -> Self {
196        Self {
197            inner: DashMap::new(),
198            rr_counters: DashMap::new(),
199            hb_ttl: Duration::from_secs(15),
200            hb_grace: Duration::from_secs(30),
201        }
202    }
203
204    pub fn with_heartbeat_policy(mut self, ttl: Duration, grace: Duration) -> Self {
205        self.hb_ttl = ttl;
206        self.hb_grace = grace;
207        self
208    }
209
210    /// Register or update a gear instance
211    pub fn register_instance(&self, instance: Arc<GearInstance>) {
212        let gear = instance.gear.clone();
213        let mut vec = self.inner.entry(gear).or_default();
214        // replace by instance_id if it already exists
215        if let Some(pos) = vec
216            .iter()
217            .position(|i| i.instance_id == instance.instance_id)
218        {
219            vec[pos] = instance;
220        } else {
221            vec.push(instance);
222        }
223    }
224
225    /// Mark an instance as ready
226    pub fn mark_ready(&self, gear: &str, instance_id: Uuid) {
227        if let Some(mut vec) = self.inner.get_mut(gear)
228            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
229        {
230            let mut state = inst.inner.write();
231            state.state = InstanceState::Ready;
232        }
233    }
234
235    /// Update the heartbeat timestamp for an instance
236    pub fn update_heartbeat(&self, gear: &str, instance_id: Uuid, at: Instant) {
237        if let Some(mut vec) = self.inner.get_mut(gear)
238            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
239        {
240            let mut state = inst.inner.write();
241            state.last_heartbeat = at;
242            // Transition Registered -> Healthy on first heartbeat
243            if state.state == InstanceState::Registered {
244                state.state = InstanceState::Healthy;
245            }
246        }
247    }
248
249    /// Mark an instance as quarantined
250    pub fn mark_quarantined(&self, gear: &str, instance_id: Uuid) {
251        if let Some(mut vec) = self.inner.get_mut(gear)
252            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
253        {
254            inst.inner.write().state = InstanceState::Quarantined;
255        }
256    }
257
258    /// Mark an instance as draining (graceful shutdown in progress)
259    pub fn mark_draining(&self, gear: &str, instance_id: Uuid) {
260        if let Some(mut vec) = self.inner.get_mut(gear)
261            && let Some(inst) = vec.iter_mut().find(|i| i.instance_id == instance_id)
262        {
263            inst.inner.write().state = InstanceState::Draining;
264        }
265    }
266
267    /// Remove an instance from the directory
268    pub fn deregister(&self, gear: &str, instance_id: Uuid) {
269        let mut remove_gear = false;
270        {
271            if let Some(mut vec) = self.inner.get_mut(gear) {
272                let list = vec.value_mut();
273                list.retain(|inst| inst.instance_id != instance_id);
274                if list.is_empty() {
275                    remove_gear = true;
276                }
277            }
278        }
279
280        if remove_gear {
281            self.inner.remove(gear);
282            self.rr_counters.remove(gear);
283            self.rr_counters.remove(&format!("rest:{gear}"));
284        }
285    }
286
287    /// Get all instances of a specific gear
288    #[must_use]
289    pub fn instances_of(&self, gear: &str) -> Vec<Arc<GearInstance>> {
290        self.inner.get(gear).map(|v| v.clone()).unwrap_or_default()
291    }
292
293    /// Get all instances across all gears
294    #[must_use]
295    pub fn all_instances(&self) -> Vec<Arc<GearInstance>> {
296        self.inner
297            .iter()
298            .flat_map(|entry| entry.value().clone())
299            .collect()
300    }
301
302    /// Quarantine or evict stale instances based on heartbeat policy
303    pub fn evict_stale(&self, now: Instant) {
304        use InstanceState::{Draining, Quarantined};
305        let mut empty_gears = Vec::new();
306
307        for mut entry in self.inner.iter_mut() {
308            let gear = entry.key().clone();
309            let vec = entry.value_mut();
310            vec.retain(|inst| {
311                let state = inst.inner.read();
312                let age = now.saturating_duration_since(state.last_heartbeat);
313
314                // Quarantine instances that have exceeded TTL
315                if age >= self.hb_ttl && !matches!(state.state, Quarantined | Draining) {
316                    drop(state); // Release read lock before write
317                    inst.inner.write().state = Quarantined;
318                    return true; // Keep quarantined instances for now
319                }
320
321                // Evict quarantined instances that exceed grace period
322                if state.state == Quarantined && age >= self.hb_ttl + self.hb_grace {
323                    return false; // Remove from directory
324                }
325
326                true
327            });
328
329            if vec.is_empty() {
330                empty_gears.push(gear);
331            }
332        }
333
334        for gear in empty_gears {
335            self.inner.remove(&gear);
336            self.rr_counters.remove(&gear);
337            self.rr_counters.remove(&format!("rest:{gear}"));
338        }
339    }
340
341    /// Pick an instance using round-robin selection, preferring healthy instances
342    #[must_use]
343    pub fn pick_instance_round_robin(&self, gear: &str) -> Option<Arc<GearInstance>> {
344        let instances_entry = self.inner.get(gear)?;
345        let instances = instances_entry.value();
346
347        if instances.is_empty() {
348            return None;
349        }
350
351        // Prefer healthy or ready instances
352        let healthy: Vec<_> = instances
353            .iter()
354            .filter(|inst| matches!(inst.state(), InstanceState::Healthy | InstanceState::Ready))
355            .cloned()
356            .collect();
357
358        let candidates: Vec<_> = if healthy.is_empty() {
359            instances.clone()
360        } else {
361            healthy
362        };
363
364        if candidates.is_empty() {
365            return None;
366        }
367
368        let len = candidates.len();
369        let mut counter = self.rr_counters.entry(gear.to_owned()).or_insert(0);
370        let idx = *counter % len;
371        *counter = (*counter + 1) % len;
372
373        candidates.get(idx).cloned()
374    }
375
376    /// Pick a service endpoint using round-robin, returning (gear, instance, endpoint).
377    /// Prefers healthy/ready instances and automatically rotates among them.
378    #[must_use]
379    pub fn pick_service_round_robin(
380        &self,
381        service_name: &str,
382    ) -> Option<(String, Arc<GearInstance>, Endpoint)> {
383        // Collect all instances that provide this service
384        let mut candidates = Vec::new();
385        for entry in &self.inner {
386            let gear = entry.key().clone();
387            for inst in entry.value() {
388                if let Some(ep) = inst.grpc_services.get(service_name) {
389                    let state = inst.state();
390                    if matches!(state, InstanceState::Healthy | InstanceState::Ready) {
391                        candidates.push((gear.clone(), inst.clone(), ep.clone()));
392                    }
393                }
394            }
395        }
396
397        if candidates.is_empty() {
398            return None;
399        }
400
401        // Use a counter keyed by service name for round-robin
402        let len = candidates.len();
403        let service_key = service_name.to_owned();
404        let mut counter = self.rr_counters.entry(service_key).or_insert(0);
405        let idx = *counter % len;
406        *counter = (*counter + 1) % len;
407
408        candidates.get(idx).cloned()
409    }
410
411    /// Resolve a REST endpoint for a gear using round-robin over instances that
412    /// expose one, preferring healthy/ready instances.
413    #[must_use]
414    pub fn pick_rest_endpoint_round_robin(&self, gear: &str) -> Option<Endpoint> {
415        let instances_entry = self.inner.get(gear)?;
416        let instances = instances_entry.value();
417
418        // Only instances that actually expose a REST endpoint are candidates.
419        let with_rest: Vec<_> = instances
420            .iter()
421            .filter(|inst| inst.rest_endpoint.is_some())
422            .cloned()
423            .collect();
424
425        if with_rest.is_empty() {
426            return None;
427        }
428
429        // Prefer healthy/ready instances, otherwise fall back to any with REST.
430        let healthy: Vec<_> = with_rest
431            .iter()
432            .filter(|inst| matches!(inst.state(), InstanceState::Healthy | InstanceState::Ready))
433            .cloned()
434            .collect();
435
436        let candidates = if healthy.is_empty() {
437            with_rest
438        } else {
439            healthy
440        };
441
442        let len = candidates.len();
443        let rr_key = format!("rest:{gear}");
444        let mut counter = self.rr_counters.entry(rr_key).or_insert(0);
445        let idx = *counter % len;
446        *counter = (*counter + 1) % len;
447
448        candidates
449            .get(idx)
450            .and_then(|inst| inst.rest_endpoint.clone())
451    }
452
453    /// Retrieve the `OpenAPI` spec of a gear, taken from the first registered
454    /// instance that published one.
455    #[must_use]
456    pub fn openapi_spec_of(&self, gear: &str) -> Option<String> {
457        let instances_entry = self.inner.get(gear)?;
458        instances_entry
459            .value()
460            .iter()
461            .find_map(|inst| inst.openapi_spec.clone())
462    }
463}
464
465impl Default for GearManager {
466    fn default() -> Self {
467        Self::new()
468    }
469}
470
471#[cfg(test)]
472#[cfg_attr(coverage_nightly, coverage(off))]
473mod tests {
474    use super::*;
475    use std::thread::sleep;
476    use std::time::Duration;
477
478    #[test]
479    fn test_register_and_retrieve_instances() {
480        let dir = GearManager::new();
481        let instance_id = Uuid::new_v4();
482        let instance = Arc::new(
483            GearInstance::new("test_gear", instance_id)
484                .with_control(Endpoint::http("localhost", 8080))
485                .with_version("1.0.0"),
486        );
487
488        dir.register_instance(instance);
489
490        let instances = dir.instances_of("test_gear");
491        assert_eq!(instances.len(), 1);
492        assert_eq!(instances[0].instance_id, instance_id);
493        assert_eq!(instances[0].gear, "test_gear");
494        assert_eq!(instances[0].version, Some("1.0.0".to_owned()));
495    }
496
497    #[test]
498    fn test_register_multiple_instances() {
499        let dir = GearManager::new();
500
501        let id1 = Uuid::new_v4();
502        let id2 = Uuid::new_v4();
503        let instance1 = Arc::new(GearInstance::new("test_gear", id1));
504        let instance2 = Arc::new(GearInstance::new("test_gear", id2));
505
506        dir.register_instance(instance1);
507        dir.register_instance(instance2);
508
509        let registered = dir.instances_of("test_gear");
510        assert_eq!(registered.len(), 2);
511
512        let ids: Vec<_> = registered.iter().map(|i| i.instance_id).collect();
513        assert!(ids.contains(&id1));
514        assert!(ids.contains(&id2));
515    }
516
517    #[test]
518    fn test_update_existing_instance() {
519        let dir = GearManager::new();
520        let instance_id = Uuid::new_v4();
521
522        let initial_instance =
523            Arc::new(GearInstance::new("test_gear", instance_id).with_version("1.0.0"));
524        dir.register_instance(initial_instance);
525
526        let updated_instance =
527            Arc::new(GearInstance::new("test_gear", instance_id).with_version("2.0.0"));
528        dir.register_instance(updated_instance);
529
530        let registered = dir.instances_of("test_gear");
531        assert_eq!(registered.len(), 1, "Should not duplicate instance");
532        assert_eq!(registered[0].version, Some("2.0.0".to_owned()));
533    }
534
535    #[test]
536    fn test_mark_ready() {
537        let dir = GearManager::new();
538        let instance_id = Uuid::new_v4();
539        let instance = Arc::new(GearInstance::new("test_gear", instance_id));
540
541        dir.register_instance(instance);
542
543        dir.mark_ready("test_gear", instance_id);
544
545        let instances = dir.instances_of("test_gear");
546        assert_eq!(instances.len(), 1);
547        assert!(matches!(instances[0].state(), InstanceState::Ready));
548    }
549
550    #[test]
551    fn test_update_heartbeat() {
552        let dir = GearManager::new();
553        let instance_id = Uuid::new_v4();
554        let instance = Arc::new(GearInstance::new("test_gear", instance_id));
555        let initial_heartbeat = instance.last_heartbeat();
556
557        dir.register_instance(instance);
558
559        // Sleep to ensure time difference
560        sleep(Duration::from_millis(10));
561
562        let new_heartbeat = Instant::now();
563        dir.update_heartbeat("test_gear", instance_id, new_heartbeat);
564
565        let instances = dir.instances_of("test_gear");
566        assert!(instances[0].last_heartbeat() > initial_heartbeat);
567        assert!(matches!(instances[0].state(), InstanceState::Healthy));
568    }
569
570    #[test]
571    fn test_all_instances() {
572        let dir = GearManager::new();
573
574        let instance1 = Arc::new(GearInstance::new("gear_a", Uuid::new_v4()));
575        let instance2 = Arc::new(GearInstance::new("gear_b", Uuid::new_v4()));
576        let instance3 = Arc::new(GearInstance::new("gear_a", Uuid::new_v4()));
577
578        dir.register_instance(instance1);
579        dir.register_instance(instance2);
580        dir.register_instance(instance3);
581
582        let all = dir.all_instances();
583        assert_eq!(all.len(), 3);
584
585        let gears: Vec<_> = all.iter().map(|i| i.gear.as_str()).collect();
586        assert_eq!(gears.iter().filter(|&m| *m == "gear_a").count(), 2);
587        assert_eq!(gears.iter().filter(|&m| *m == "gear_b").count(), 1);
588    }
589
590    #[test]
591    fn test_pick_instance_round_robin() {
592        let dir = GearManager::new();
593
594        let id1 = Uuid::new_v4();
595        let id2 = Uuid::new_v4();
596        let instance1 = Arc::new(GearInstance::new("test_gear", id1));
597        let instance2 = Arc::new(GearInstance::new("test_gear", id2));
598
599        dir.register_instance(instance1);
600        dir.register_instance(instance2);
601
602        // Pick three times to verify round-robin behavior
603        let picked1 = dir.pick_instance_round_robin("test_gear").unwrap();
604        let picked2 = dir.pick_instance_round_robin("test_gear").unwrap();
605        let picked3 = dir.pick_instance_round_robin("test_gear").unwrap();
606
607        let ids = [
608            picked1.instance_id,
609            picked2.instance_id,
610            picked3.instance_id,
611        ];
612
613        // With 2 instances, we expect round-robin pattern like A, B, A
614        // Check that both instance IDs appear and that at least one repeats
615        assert!(ids.contains(&id1));
616        assert!(ids.contains(&id2));
617        // First and third pick should be the same (round-robin wraps)
618        assert_eq!(picked1.instance_id, picked3.instance_id);
619        // Second pick should be different from the first
620        assert_ne!(picked1.instance_id, picked2.instance_id);
621    }
622
623    #[test]
624    fn test_pick_instance_none_available() {
625        let dir = GearManager::new();
626        let picked = dir.pick_instance_round_robin("nonexistent_gear");
627        assert!(picked.is_none());
628    }
629
630    #[test]
631    fn test_endpoint_creation() {
632        let plain_ep = Endpoint::http("localhost", 8080);
633        assert_eq!(plain_ep.uri, "http://localhost:8080");
634
635        let secure_ep = Endpoint::https("localhost", 8443);
636        assert_eq!(secure_ep.uri, "https://localhost:8443");
637
638        let uds_ep = Endpoint::uds("/tmp/socket.sock");
639        assert!(uds_ep.uri.starts_with("unix://"));
640        assert!(uds_ep.uri.contains("socket.sock"));
641
642        let custom_ep = Endpoint::from_uri("http://example.com");
643        assert_eq!(custom_ep.uri, "http://example.com");
644    }
645
646    #[test]
647    fn test_endpoint_kind() {
648        let plain_ep = Endpoint::http("127.0.0.1", 8080);
649        match plain_ep.kind() {
650            EndpointKind::Tcp(addr) => {
651                assert_eq!(addr.ip().to_string(), "127.0.0.1");
652                assert_eq!(addr.port(), 8080);
653            }
654            _ => panic!("Expected TCP endpoint for http"),
655        }
656
657        let secure_ep = Endpoint::https("127.0.0.1", 8443);
658        match secure_ep.kind() {
659            EndpointKind::Tcp(addr) => {
660                assert_eq!(addr.ip().to_string(), "127.0.0.1");
661                assert_eq!(addr.port(), 8443);
662            }
663            _ => panic!("Expected TCP endpoint for https"),
664        }
665
666        let uds_ep = Endpoint::uds("/tmp/test.sock");
667        match uds_ep.kind() {
668            EndpointKind::Uds(path) => {
669                assert!(path.to_string_lossy().contains("test.sock"));
670            }
671            _ => panic!("Expected UDS endpoint"),
672        }
673
674        let other_ep = Endpoint::from_uri("grpc://example.com");
675        match other_ep.kind() {
676            EndpointKind::Other(uri) => {
677                assert_eq!(uri, "grpc://example.com");
678            }
679            _ => panic!("Expected Other endpoint"),
680        }
681    }
682
683    #[test]
684    fn test_gear_instance_builder() {
685        let instance_id = Uuid::new_v4();
686        let instance = GearInstance::new("test_gear", instance_id)
687            .with_control(Endpoint::http("localhost", 8080))
688            .with_version("1.2.3")
689            .with_grpc_service("service1", Endpoint::http("localhost", 8082))
690            .with_grpc_service("service2", Endpoint::http("localhost", 8083));
691
692        assert_eq!(instance.gear, "test_gear");
693        assert_eq!(instance.instance_id, instance_id);
694        assert!(instance.control.is_some());
695        assert_eq!(instance.version, Some("1.2.3".to_owned()));
696        assert_eq!(instance.grpc_services.len(), 2);
697        assert!(instance.grpc_services.contains_key("service1"));
698        assert!(instance.grpc_services.contains_key("service2"));
699        assert!(matches!(instance.state(), InstanceState::Registered));
700    }
701
702    #[test]
703    fn test_quarantine_and_evict() {
704        let ttl = Duration::from_millis(50);
705        let grace = Duration::from_millis(50);
706        let dir = GearManager::new().with_heartbeat_policy(ttl, grace);
707
708        let now = Instant::now();
709        let instance = GearInstance::new("test_gear", Uuid::new_v4());
710        // Set the last heartbeat to be stale
711        instance.inner.write().last_heartbeat = now
712            .checked_sub(ttl)
713            .and_then(|t| t.checked_sub(Duration::from_millis(10)))
714            .expect("test duration subtraction should not underflow");
715
716        dir.register_instance(Arc::new(instance));
717
718        dir.evict_stale(now);
719        let instances = dir.instances_of("test_gear");
720        assert_eq!(instances.len(), 1);
721        assert!(matches!(instances[0].state(), InstanceState::Quarantined));
722
723        let later = now + grace + Duration::from_millis(10);
724        dir.evict_stale(later);
725
726        let instances_after = dir.instances_of("test_gear");
727        assert!(instances_after.is_empty());
728    }
729
730    #[test]
731    fn test_instances_of_empty() {
732        let dir = GearManager::new();
733        let instances = dir.instances_of("nonexistent");
734        assert!(instances.is_empty());
735    }
736
737    #[test]
738    fn test_rr_prefers_healthy() {
739        let dir = GearManager::new();
740
741        // Create two instances: one healthy, one quarantined
742        let healthy_id = Uuid::new_v4();
743        let healthy = Arc::new(GearInstance::new("test_gear", healthy_id));
744        dir.register_instance(healthy);
745        dir.update_heartbeat("test_gear", healthy_id, Instant::now());
746
747        let quarantined_id = Uuid::new_v4();
748        let quarantined = Arc::new(GearInstance::new("test_gear", quarantined_id));
749        dir.register_instance(quarantined);
750        dir.mark_quarantined("test_gear", quarantined_id);
751
752        // RR should only pick the healthy instance
753        for _ in 0..5 {
754            let picked = dir.pick_instance_round_robin("test_gear").unwrap();
755            assert_eq!(picked.instance_id, healthy_id);
756        }
757    }
758
759    #[test]
760    fn test_pick_rest_endpoint_and_openapi() {
761        let dir = GearManager::new();
762        let id = Uuid::new_v4();
763        let instance = Arc::new(
764            GearInstance::new("billing", id)
765                .with_rest_endpoint(Endpoint::http("billing", 8080))
766                .with_openapi_spec("{\"openapi\":\"3.1.0\"}"),
767        );
768        dir.register_instance(instance);
769
770        let rest = dir.pick_rest_endpoint_round_robin("billing").unwrap();
771        assert_eq!(rest.uri, "http://billing:8080");
772
773        let spec = dir.openapi_spec_of("billing").unwrap();
774        assert!(spec.contains("openapi"));
775    }
776
777    #[test]
778    fn test_pick_rest_endpoint_none_when_absent() {
779        let dir = GearManager::new();
780        let id = Uuid::new_v4();
781        // Instance exposes only a gRPC service, no REST endpoint.
782        let instance = Arc::new(
783            GearInstance::new("grpc_only", id)
784                .with_grpc_service("some.Service", Endpoint::http("127.0.0.1", 9000)),
785        );
786        dir.register_instance(instance);
787
788        assert!(dir.pick_rest_endpoint_round_robin("grpc_only").is_none());
789        assert!(dir.openapi_spec_of("grpc_only").is_none());
790        assert!(dir.pick_rest_endpoint_round_robin("missing").is_none());
791    }
792
793    #[test]
794    fn test_pick_rest_endpoint_round_robin_rotates() {
795        let dir = GearManager::new();
796        let id1 = Uuid::new_v4();
797        let id2 = Uuid::new_v4();
798        let inst1 = Arc::new(
799            GearInstance::new("web", id1).with_rest_endpoint(Endpoint::http("127.0.0.1", 8001)),
800        );
801        let inst2 = Arc::new(
802            GearInstance::new("web", id2).with_rest_endpoint(Endpoint::http("127.0.0.1", 8002)),
803        );
804        dir.register_instance(inst1);
805        dir.register_instance(inst2);
806        dir.update_heartbeat("web", id1, Instant::now());
807        dir.update_heartbeat("web", id2, Instant::now());
808
809        let ep1 = dir.pick_rest_endpoint_round_robin("web").unwrap();
810        let ep2 = dir.pick_rest_endpoint_round_robin("web").unwrap();
811        let ep3 = dir.pick_rest_endpoint_round_robin("web").unwrap();
812
813        assert_ne!(ep1, ep2);
814        assert_eq!(ep1, ep3);
815    }
816
817    #[test]
818    fn test_pick_service_round_robin() {
819        let dir = GearManager::new();
820
821        let id1 = Uuid::new_v4();
822        let id2 = Uuid::new_v4();
823        // Register two instances providing the same service
824        let inst1 = Arc::new(
825            GearInstance::new("test_gear", id1)
826                .with_grpc_service("test.Service", Endpoint::http("127.0.0.1", 8001)),
827        );
828        let inst2 = Arc::new(
829            GearInstance::new("test_gear", id2)
830                .with_grpc_service("test.Service", Endpoint::http("127.0.0.1", 8002)),
831        );
832
833        dir.register_instance(inst1);
834        dir.register_instance(inst2);
835
836        // Mark both as healthy
837        dir.update_heartbeat("test_gear", id1, Instant::now());
838        dir.update_heartbeat("test_gear", id2, Instant::now());
839
840        // Pick should rotate between instances
841        let pick1 = dir.pick_service_round_robin("test.Service");
842        let pick2 = dir.pick_service_round_robin("test.Service");
843        let pick3 = dir.pick_service_round_robin("test.Service");
844
845        assert!(pick1.is_some());
846        assert!(pick2.is_some());
847        assert!(pick3.is_some());
848
849        let (_, inst1, ep1) = pick1.unwrap();
850        let (_, inst2, ep2) = pick2.unwrap();
851        let (_, inst3, _) = pick3.unwrap();
852
853        // First and third should be the same (round-robin)
854        assert_eq!(inst1.instance_id, inst3.instance_id);
855        // First and second should be different
856        assert_ne!(inst1.instance_id, inst2.instance_id);
857        // Endpoints should differ
858        assert_ne!(ep1, ep2);
859    }
860
861    #[test]
862    fn test_deregister_clears_rr_counters() {
863        let dir = GearManager::new();
864        let id = Uuid::new_v4();
865        let instance = Arc::new(
866            GearInstance::new("web", id).with_rest_endpoint(Endpoint::http("127.0.0.1", 8001)),
867        );
868        dir.register_instance(instance);
869        dir.update_heartbeat("web", id, Instant::now());
870
871        // Exercise both round-robin counters so the keys are created.
872        assert!(dir.pick_instance_round_robin("web").is_some());
873        assert!(dir.pick_rest_endpoint_round_robin("web").is_some());
874
875        assert!(dir.rr_counters.contains_key("web"));
876        assert!(dir.rr_counters.contains_key("rest:web"));
877
878        dir.deregister("web", id);
879
880        assert!(!dir.rr_counters.contains_key("web"));
881        assert!(!dir.rr_counters.contains_key("rest:web"));
882    }
883
884    #[test]
885    fn test_evict_stale_clears_rr_counters() {
886        let ttl = Duration::from_millis(50);
887        let grace = Duration::from_millis(50);
888        let dir = GearManager::new().with_heartbeat_policy(ttl, grace);
889
890        let now = Instant::now();
891        let id = Uuid::new_v4();
892        let instance = Arc::new(
893            GearInstance::new("web", id).with_rest_endpoint(Endpoint::http("127.0.0.1", 8001)),
894        );
895        // Set the last heartbeat to be stale so the instance is quarantined then evicted.
896        instance.inner.write().last_heartbeat = now
897            .checked_sub(ttl)
898            .and_then(|t| t.checked_sub(Duration::from_millis(10)))
899            .expect("test duration subtraction should not underflow");
900
901        dir.register_instance(instance);
902        assert!(dir.pick_rest_endpoint_round_robin("web").is_some());
903
904        assert!(dir.rr_counters.contains_key("rest:web"));
905
906        // First eviction pass quarantines the stale instance.
907        dir.evict_stale(now);
908        let instances = dir.instances_of("web");
909        assert_eq!(instances.len(), 1);
910        assert!(matches!(instances[0].state(), InstanceState::Quarantined));
911
912        // Second pass evicts quarantined instances that exceeded the grace period.
913        let later = now + grace + Duration::from_millis(10);
914        dir.evict_stale(later);
915
916        assert!(dir.instances_of("web").is_empty());
917        assert!(!dir.rr_counters.contains_key("web"));
918        assert!(!dir.rr_counters.contains_key("rest:web"));
919    }
920}
toolkit/runtime/gear_manager.rs

toolkit/runtime/
gear_manager.rs
