1use anyhow::Result;
4use async_trait::async_trait;
5use std::hash::{Hash, Hasher};
6use std::sync::Arc;
7use uuid::Uuid;
8
9use crate::runtime::{Endpoint, GearInstance, GearManager, GrpcServiceNameConflict};
10
11fn openapi_spec_hash(spec: &str) -> String {
28 let mut hasher = std::collections::hash_map::DefaultHasher::new();
29 spec.hash(&mut hasher);
30 format!("{:016x}", hasher.finish())
31}
32
33pub use cf_system_sdks::directory::{
35 DirectoryClient, DirectoryInvalidArgument, DirectoryNotFound, DirectoryPermissionDenied,
36 DirectoryServiceNameConflict, GrpcServiceInfo, InstanceState, LabelSelector,
37 RegisterInstanceInfo, ServiceEndpoint, ServiceInstanceInfo,
38};
39
40fn runtime_state_to_domain(state: crate::runtime::InstanceState) -> InstanceState {
44 use crate::runtime::InstanceState as Rt;
45 match state {
46 Rt::Registered => InstanceState::Registered,
47 Rt::Ready => InstanceState::Ready,
48 Rt::Healthy => InstanceState::Healthy,
49 Rt::Quarantined => InstanceState::Quarantined,
50 Rt::Draining => InstanceState::Draining,
51 }
52}
53
54fn project_instance(gear: &str, inst: &GearInstance) -> ServiceInstanceInfo {
62 let endpoint = inst
63 .grpc_services
64 .values()
65 .next()
66 .or(inst.rest_endpoint.as_ref())
67 .map(|ep| ServiceEndpoint::new(ep.uri.clone()));
68
69 ServiceInstanceInfo::new(gear, inst.instance_id.to_string())
70 .with_endpoint(endpoint)
71 .with_version(inst.version.clone())
72 .with_rest_endpoint(
73 inst.rest_endpoint
74 .as_ref()
75 .map(|ep| ServiceEndpoint::new(ep.uri.clone())),
76 )
77 .with_openapi_spec_hash(inst.openapi_spec.as_deref().map(openapi_spec_hash))
78 .with_grpc_services(
82 inst.grpc_services
83 .iter()
84 .map(|(name, e)| (name.clone(), ServiceEndpoint::new(e.uri.clone())))
85 .collect(),
86 )
87 .with_labels(inst.labels.clone())
88 .with_state(runtime_state_to_domain(inst.state()))
89}
90
91pub struct LocalDirectoryClient {
96 mgr: Arc<GearManager>,
97}
98
99impl LocalDirectoryClient {
100 #[must_use]
101 pub fn new(mgr: Arc<GearManager>) -> Self {
102 Self { mgr }
103 }
104}
105
106#[async_trait]
107impl DirectoryClient for LocalDirectoryClient {
108 async fn resolve_grpc_service(&self, service_name: &str) -> Result<ServiceEndpoint> {
116 if let Some((_gear, _inst, ep)) = self.mgr.pick_service_round_robin(service_name) {
117 return Ok(ServiceEndpoint::new(ep.uri));
118 }
119
120 Err(DirectoryNotFound::new(format!("service {service_name}")).into())
121 }
122
123 async fn resolve_rest_service(&self, gear_name: &str) -> Result<ServiceEndpoint> {
124 if let Some(ep) = self.mgr.pick_rest_endpoint_round_robin(gear_name) {
125 return Ok(ServiceEndpoint::new(ep.uri));
126 }
127
128 Err(DirectoryNotFound::new(format!("gear {gear_name}")).into())
129 }
130
131 async fn get_openapi_spec(&self, gear_name: &str) -> Result<String> {
132 self.mgr.openapi_spec_of(gear_name).ok_or_else(|| {
133 DirectoryNotFound::new(format!("openapi spec for gear {gear_name}")).into()
134 })
135 }
136
137 async fn list_instances(&self, gear: &str) -> Result<Vec<ServiceInstanceInfo>> {
138 Ok(self
141 .mgr
142 .instances_of(gear)
143 .iter()
144 .map(|inst| project_instance(gear, inst))
145 .collect())
146 }
147
148 async fn resolve_by_labels(
149 &self,
150 gear: &str,
151 selector: &LabelSelector,
152 ) -> Result<Vec<ServiceInstanceInfo>> {
153 Ok(self
156 .mgr
157 .instances_of(gear)
158 .iter()
159 .map(|inst| project_instance(gear, inst))
160 .filter(|i| selector.matches(&i.labels))
161 .collect())
162 }
163
164 async fn list_all_instances(&self) -> Result<Vec<ServiceInstanceInfo>> {
165 Ok(self
169 .mgr
170 .all_instances()
171 .iter()
172 .filter_map(|inst| {
173 let info = project_instance(&inst.gear, inst);
174 if info.endpoint.is_none() {
175 tracing::debug!(
176 gear = %inst.gear,
177 instance_id = %inst.instance_id,
178 "skipping instance with no gRPC or REST endpoint from cross-gear snapshot"
179 );
180 return None;
181 }
182 Some(info.without_labels())
183 })
184 .collect())
185 }
186
187 async fn register_instance(&self, info: RegisterInstanceInfo) -> Result<()> {
188 let instance_id = Uuid::parse_str(&info.instance_id).map_err(|e| {
190 anyhow::Error::from(DirectoryInvalidArgument::new(format!(
195 "invalid instance_id '{}': {e}",
196 info.instance_id
197 )))
198 })?;
199
200 cf_system_sdks::directory::validate_labels(&info.labels)
207 .map_err(|e| anyhow::Error::from(DirectoryInvalidArgument::new(e.to_string())))?;
208
209 let mut instance = GearInstance::new(info.gear.clone(), instance_id);
211
212 if let Some(version) = info.version {
214 instance = instance.with_version(version);
215 }
216
217 for (service_name, endpoint) in info.grpc_services {
219 instance = instance.with_grpc_service(service_name, Endpoint::from_uri(endpoint.uri));
220 }
221
222 if let Some(rest) = info.rest_endpoint {
224 instance = instance.with_rest_endpoint(Endpoint::from_uri(rest.uri));
225 }
226
227 if let Some(spec) = info.openapi_spec {
229 instance = instance.with_openapi_spec(spec);
230 }
231
232 if !info.labels.is_empty() {
234 instance = instance.with_labels(info.labels);
235 }
236
237 self.mgr.register_instance(Arc::new(instance)).map_err(
243 |GrpcServiceNameConflict {
244 service_name,
245 owner,
246 recoverable,
247 }| {
248 anyhow::Error::from(DirectoryServiceNameConflict {
249 service_name,
250 owner,
251 recoverable,
252 })
253 },
254 )?;
255
256 Ok(())
257 }
258
259 async fn deregister_instance(&self, gear: &str, instance_id: &str) -> Result<()> {
260 let instance_id = Uuid::parse_str(instance_id)
261 .map_err(|e| anyhow::anyhow!("Invalid instance_id '{instance_id}': {e}"))?;
262 self.mgr.deregister(gear, instance_id);
263 Ok(())
264 }
265
266 async fn send_heartbeat(&self, gear: &str, instance_id: &str) -> Result<()> {
267 let instance_id = Uuid::parse_str(instance_id)
268 .map_err(|e| anyhow::anyhow!("Invalid instance_id '{instance_id}': {e}"))?;
269 self.mgr
270 .update_heartbeat(gear, instance_id, std::time::Instant::now());
271 Ok(())
272 }
273}
274
275#[cfg(test)]
276#[cfg_attr(coverage_nightly, coverage(off))]
277mod tests {
278 use super::*;
279
280 #[tokio::test]
281 async fn test_resolve_grpc_service_not_found() {
282 let dir = Arc::new(GearManager::new());
283 let api = LocalDirectoryClient::new(dir);
284
285 let err = api
286 .resolve_grpc_service("nonexistent.Service")
287 .await
288 .unwrap_err();
289 assert!(
293 err.downcast_ref::<DirectoryNotFound>().is_some(),
294 "expected the typed not-found sentinel, got: {err:?}"
295 );
296 }
297
298 #[tokio::test]
299 async fn test_register_instance_via_api() {
300 let dir = Arc::new(GearManager::new());
301 let api = LocalDirectoryClient::new(dir.clone());
302
303 let instance_id = Uuid::new_v4();
304 let register_info = RegisterInstanceInfo::new("test_gear", instance_id.to_string())
306 .with_grpc_services(vec![(
307 "test.Service".to_owned(),
308 ServiceEndpoint::http("127.0.0.1", 8001),
309 )])
310 .with_version("1.0.0");
311
312 api.register_instance(register_info).await.unwrap();
313
314 let instances = dir.instances_of("test_gear");
316 assert_eq!(instances.len(), 1);
317 assert_eq!(instances[0].instance_id, instance_id);
318 assert_eq!(instances[0].version, Some("1.0.0".to_owned()));
319 assert!(instances[0].grpc_services.contains_key("test.Service"));
320 }
321
322 #[tokio::test]
323 async fn register_rejects_invalid_labels_at_store_boundary() {
324 let dir = Arc::new(GearManager::new());
328 let api = LocalDirectoryClient::new(Arc::clone(&dir));
329
330 let err = api
331 .register_instance(
332 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
333 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
334 .with_labels(labels(&[("bad key", "7")])),
335 )
336 .await
337 .expect_err("an invalid label must be rejected at the store boundary");
338 assert!(
339 err.downcast_ref::<DirectoryInvalidArgument>().is_some(),
340 "store-boundary label rejection must be typed InvalidArgument, got: {err}"
341 );
342 assert!(
343 dir.instances_of("worker").is_empty(),
344 "a rejected registration must not be stored"
345 );
346 }
347
348 #[tokio::test]
349 async fn test_register_and_resolve_rest_and_openapi() {
350 let dir = Arc::new(GearManager::new());
351 let api = LocalDirectoryClient::new(dir.clone());
352
353 let instance_id = Uuid::new_v4();
354 let register_info = RegisterInstanceInfo::new("billing", instance_id.to_string())
355 .with_version("1.0.0")
356 .with_rest_endpoint(ServiceEndpoint::http("billing", 8080))
357 .with_openapi_spec("{\"openapi\":\"3.1.0\"}");
358
359 api.register_instance(register_info).await.unwrap();
360
361 let resolved = api.resolve_rest_service("billing").await.unwrap();
363 assert_eq!(resolved.uri, concat!("http", "://billing:8080"));
364
365 let spec = api.get_openapi_spec("billing").await.unwrap();
367 assert!(spec.contains("openapi"));
368 }
369
370 #[tokio::test]
371 async fn test_resolve_rest_and_openapi_not_found() {
372 let dir = Arc::new(GearManager::new());
373 let api = LocalDirectoryClient::new(dir);
374
375 let rest_err = api.resolve_rest_service("missing").await.unwrap_err();
378 assert!(
379 rest_err.downcast_ref::<DirectoryNotFound>().is_some(),
380 "expected the typed not-found sentinel, got: {rest_err:?}"
381 );
382
383 let spec_err = api.get_openapi_spec("missing").await.unwrap_err();
384 assert!(
385 spec_err.downcast_ref::<DirectoryNotFound>().is_some(),
386 "expected the typed not-found sentinel, got: {spec_err:?}"
387 );
388 }
389
390 #[tokio::test]
391 async fn test_deregister_instance_via_api() {
392 let dir = Arc::new(GearManager::new());
393 let api = LocalDirectoryClient::new(dir.clone());
394
395 let instance_id = Uuid::new_v4();
396 let inst = Arc::new(GearInstance::new("test_gear", instance_id));
398 dir.register_instance(inst).unwrap();
399
400 assert_eq!(dir.instances_of("test_gear").len(), 1);
402
403 api.deregister_instance("test_gear", &instance_id.to_string())
405 .await
406 .unwrap();
407
408 assert_eq!(dir.instances_of("test_gear").len(), 0);
410 }
411
412 #[tokio::test]
413 async fn test_send_heartbeat_via_api() {
414 use crate::runtime::InstanceState;
415
416 let dir = Arc::new(GearManager::new());
417 let api = LocalDirectoryClient::new(dir.clone());
418
419 let instance_id = Uuid::new_v4();
420 let inst = Arc::new(GearInstance::new("test_gear", instance_id));
422 dir.register_instance(inst).unwrap();
423
424 let instances = dir.instances_of("test_gear");
426 assert_eq!(instances[0].state(), InstanceState::Registered);
427
428 api.send_heartbeat("test_gear", &instance_id.to_string())
430 .await
431 .unwrap();
432
433 let instances = dir.instances_of("test_gear");
435 assert_eq!(instances[0].state(), InstanceState::Healthy);
436 }
437
438 #[tokio::test]
439 async fn test_list_all_instances_across_gears() {
440 let dir = Arc::new(GearManager::new());
441 let api = LocalDirectoryClient::new(Arc::clone(&dir));
442
443 for (gear, port) in [("billing", 8080u16), ("catalog", 8081u16)] {
445 api.register_instance(
446 RegisterInstanceInfo::new(gear, Uuid::new_v4().to_string())
447 .with_version("1.0.0")
448 .with_rest_endpoint(ServiceEndpoint::http(gear, port))
449 .with_openapi_spec(format!("{{\"openapi\":\"3.1.0\",\"x\":\"{gear}\"}}")),
450 )
451 .await
452 .unwrap();
453 }
454
455 api.register_instance(
459 RegisterInstanceInfo::new("reporting", Uuid::new_v4().to_string())
460 .with_grpc_services(vec![(
461 "reporting.Service".to_owned(),
462 ServiceEndpoint::new("http://reporting:7000"),
463 )])
464 .with_version("1.0.0"),
465 )
466 .await
467 .unwrap();
468
469 let all = api.list_all_instances().await.unwrap();
470 assert_eq!(all.len(), 3);
471
472 let billing = all.iter().find(|i| i.gear == "billing").expect("billing");
473 assert_eq!(
474 billing.rest_endpoint.as_ref().map(|e| e.uri.as_str()),
475 Some("http://billing:8080")
476 );
477 assert!(billing.openapi_spec_hash.is_some());
481 assert!(
482 api.get_openapi_spec("billing")
483 .await
484 .expect("billing spec")
485 .contains("billing")
486 );
487
488 let reporting = all
491 .iter()
492 .find(|i| i.gear == "reporting")
493 .expect("reporting");
494 assert_eq!(
495 reporting.endpoint.as_ref().map(|e| e.uri.as_str()),
496 Some("http://reporting:7000")
497 );
498 assert!(reporting.rest_endpoint.is_none());
499 assert!(reporting.openapi_spec_hash.is_none());
500 }
501
502 fn labels(pairs: &[(&str, &str)]) -> std::collections::BTreeMap<String, String> {
503 pairs
504 .iter()
505 .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
506 .collect()
507 }
508
509 #[tokio::test]
510 async fn register_carries_labels_into_list_instances() {
511 let dir = Arc::new(GearManager::new());
512 let api = LocalDirectoryClient::new(Arc::clone(&dir));
513
514 api.register_instance(
515 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
516 .with_grpc_services(vec![(
517 "worker.Svc".to_owned(),
518 ServiceEndpoint::new("http://worker:7000"),
519 )])
520 .with_labels(labels(&[("shard", "7")])),
521 )
522 .await
523 .unwrap();
524
525 let listed = api.list_instances("worker").await.unwrap();
526 assert_eq!(listed.len(), 1);
527 assert_eq!(listed[0].labels.get("shard"), Some(&"7".to_owned()));
528 }
529
530 #[tokio::test]
531 async fn resolve_by_labels_selects_matching_instances() {
532 let dir = Arc::new(GearManager::new());
533 let api = LocalDirectoryClient::new(Arc::clone(&dir));
534
535 for (id, shard) in [("a", "7"), ("b", "8"), ("c", "7")] {
536 api.register_instance(
537 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
538 .with_grpc_services(vec![(
539 format!("worker.{id}"),
540 ServiceEndpoint::new(format!("http://worker-{id}:7000")),
541 )])
542 .with_labels(labels(&[("shard", shard)])),
543 )
544 .await
545 .unwrap();
546 }
547
548 let matched = api
549 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
550 .await
551 .unwrap();
552 assert_eq!(matched.len(), 2, "two instances carry shard=7");
553 assert!(
554 matched
555 .iter()
556 .all(|i| i.labels.get("shard") == Some(&"7".to_owned()))
557 );
558 }
559
560 #[tokio::test]
561 async fn labelless_reregister_preserves_labels_through_api() {
562 let dir = Arc::new(GearManager::new());
563 let api = LocalDirectoryClient::new(Arc::clone(&dir));
564
565 let instance_id = Uuid::new_v4().to_string();
566 api.register_instance(
567 RegisterInstanceInfo::new("worker", instance_id.clone())
568 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
569 .with_labels(labels(&[("shard", "7")])),
570 )
571 .await
572 .unwrap();
573
574 api.register_instance(
576 RegisterInstanceInfo::new("worker", instance_id)
577 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
578 .with_version("2.0.0"),
579 )
580 .await
581 .unwrap();
582
583 let listed = api.list_instances("worker").await.unwrap();
584 assert_eq!(listed.len(), 1);
585 assert_eq!(
586 listed[0].labels.get("shard"),
587 Some(&"7".to_owned()),
588 "label-less re-registration must not wipe stored labels"
589 );
590
591 let matched = api
593 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
594 .await
595 .unwrap();
596 assert_eq!(matched.len(), 1);
597 }
598
599 #[tokio::test]
600 async fn resolve_by_labels_carries_live_serving_state() {
601 let dir = Arc::new(GearManager::new());
602 let api = LocalDirectoryClient::new(Arc::clone(&dir));
603
604 let instance_id = Uuid::new_v4();
605 api.register_instance(
606 RegisterInstanceInfo::new("worker", instance_id.to_string())
607 .with_grpc_services(vec![(
608 "worker.Svc".to_owned(),
609 ServiceEndpoint::new("http://worker:7000"),
610 )])
611 .with_labels(labels(&[("shard", "7")])),
612 )
613 .await
614 .unwrap();
615
616 let matched = api
618 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
619 .await
620 .unwrap();
621 assert_eq!(matched.len(), 1);
622 assert_eq!(matched[0].state, InstanceState::Registered);
623 assert!(!matched[0].state.is_serving());
624
625 dir.update_heartbeat("worker", instance_id, std::time::Instant::now());
628 let matched = api
629 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
630 .await
631 .unwrap();
632 assert_eq!(matched[0].state, InstanceState::Healthy);
633 assert!(matched[0].state.is_serving());
634 }
635
636 #[tokio::test]
637 async fn resolve_by_labels_omits_openapi_spec() {
638 let dir = Arc::new(GearManager::new());
639 let api = LocalDirectoryClient::new(Arc::clone(&dir));
640
641 api.register_instance(
642 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
643 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
644 .with_openapi_spec("{\"openapi\":\"3.1.0\"}")
645 .with_labels(labels(&[("shard", "7")])),
646 )
647 .await
648 .unwrap();
649
650 let matched = api
652 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
653 .await
654 .unwrap();
655 assert_eq!(matched.len(), 1);
656 assert!(
657 matched[0].openapi_spec_hash.is_some(),
658 "in-process resolve_by_labels carries only the spec hash, never the document"
659 );
660
661 let listed = api.list_instances("worker").await.unwrap();
664 assert!(
665 listed[0].openapi_spec_hash.is_some(),
666 "list_instances must carry the spec hash, never the document"
667 );
668 }
669
670 #[tokio::test]
671 async fn list_all_instances_skips_endpoint_less() {
672 let dir = Arc::new(GearManager::new());
673 let api = LocalDirectoryClient::new(Arc::clone(&dir));
674
675 api.register_instance(
678 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
679 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080")),
680 )
681 .await
682 .unwrap();
683 api.register_instance(RegisterInstanceInfo::new(
684 "placeholder",
685 Uuid::new_v4().to_string(),
686 ))
687 .await
688 .unwrap();
689
690 let all = api.list_all_instances().await.unwrap();
691 assert_eq!(all.len(), 1, "the endpoint-less instance is skipped");
692 assert_eq!(all[0].gear, "worker");
693 assert!(
694 all.iter()
695 .all(|i| i.endpoint.as_ref().is_some_and(|e| !e.uri.is_empty())),
696 "no instance may carry an absent or empty-URI endpoint"
697 );
698 }
699
700 #[tokio::test]
701 async fn label_path_returns_endpoint_less_match() {
702 let dir = Arc::new(GearManager::new());
703 let api = LocalDirectoryClient::new(Arc::clone(&dir));
704
705 api.register_instance(
710 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
711 .with_labels(labels(&[("shard", "7")])),
712 )
713 .await
714 .unwrap();
715
716 let matched = api
717 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
718 .await
719 .unwrap();
720 assert_eq!(
721 matched.len(),
722 1,
723 "resolve_by_labels must return a matched instance even with no endpoint"
724 );
725 assert!(
726 matched[0].endpoint.is_none(),
727 "a no-endpoint match carries endpoint = None, not an empty-URI sentinel"
728 );
729 assert!(matched[0].rest_endpoint.is_none());
730 assert_eq!(matched[0].labels.get("shard"), Some(&"7".to_owned()));
731
732 let listed = api.list_instances("worker").await.unwrap();
734 assert_eq!(
735 listed.len(),
736 1,
737 "list_instances must not drop endpoint-less instances"
738 );
739 assert!(listed[0].endpoint.is_none());
740 }
741}