1use dashmap::DashMap;
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7use uuid::Uuid;
8
9#[derive(Clone, Debug, PartialEq, Eq, Hash)]
11pub struct Endpoint {
12 pub uri: String,
13}
14
15#[derive(Clone, Debug, PartialEq, Eq)]
17pub enum EndpointKind {
18 Tcp(std::net::SocketAddr),
20 Uds(std::path::PathBuf),
22 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 #[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#[derive(Clone, Debug)]
82pub struct InstanceRuntimeState {
83 pub last_heartbeat: Instant,
84 pub state: InstanceState,
85}
86
87#[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 #[must_use]
144 pub fn state(&self) -> InstanceState {
145 self.inner.read().state
146 }
147
148 #[must_use]
150 pub fn last_heartbeat(&self) -> Instant {
151 self.inner.read().last_heartbeat
152 }
153}
154
155#[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 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 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 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 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 if state.state == InstanceState::Registered {
228 state.state = InstanceState::Healthy;
229 }
230 }
231 }
232
233 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 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 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 #[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 #[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 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 if age >= self.hb_ttl && !matches!(state.state, Quarantined | Draining) {
299 drop(state); inst.inner.write().state = Quarantined;
301 return true; }
303
304 if state.state == Quarantined && age >= self.hb_ttl + self.hb_grace {
306 return false; }
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 #[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 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 #[must_use]
361 pub fn pick_service_round_robin(
362 &self,
363 service_name: &str,
364 ) -> Option<(String, Arc<GearInstance>, Endpoint)> {
365 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 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(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 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 assert!(ids.contains(&id1));
545 assert!(ids.contains(&id2));
546 assert_eq!(picked1.instance_id, picked3.instance_id);
548 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 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 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 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 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 dir.update_heartbeat("test_gear", id1, Instant::now());
709 dir.update_heartbeat("test_gear", id2, Instant::now());
710
711 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 assert_eq!(inst1.instance_id, inst3.instance_id);
726 assert_ne!(inst1.instance_id, inst2.instance_id);
728 assert_ne!(ep1, ep2);
730 }
731}