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]
326 async fn presence_and_hub_registrations_keep_each_others_half() {
327 for hub_first in [true, false] {
328 let api = LocalDirectoryClient::new(Arc::new(GearManager::new()));
330 let id = Uuid::new_v4().to_string();
331 let presence = RegisterInstanceInfo::new("worker", &id)
333 .with_rest_endpoint(ServiceEndpoint::http("127.0.0.1", 8080))
334 .with_openapi_spec(r#"{"openapi":"3.1.0"}"#);
335 let hub = RegisterInstanceInfo::new("worker", &id)
337 .with_grpc_services(vec![(
338 "worker.Svc".to_owned(),
339 ServiceEndpoint::http("127.0.0.1", 50051),
340 )])
341 .with_version("1.0.0");
342 let order = if hub_first {
343 [&hub, &presence, &presence, &hub]
344 } else {
345 [&presence, &hub, &hub, &presence]
346 };
347 for info in order {
348 api.register_instance(info.clone()).await.unwrap();
349 }
350 let instances = api.list_instances("worker").await.unwrap();
351 assert_eq!(instances.len(), 1);
352 assert_eq!(instances[0].grpc_services.len(), 1, "gRPC services kept");
353 assert!(instances[0].rest_endpoint.is_some(), "REST endpoint kept");
354 assert_eq!(
355 api.get_openapi_spec("worker").await.unwrap(),
356 r#"{"openapi":"3.1.0"}"#,
357 "OpenAPI spec kept"
358 );
359 }
360 }
361
362 #[tokio::test]
363 async fn register_rejects_invalid_labels_at_store_boundary() {
364 let dir = Arc::new(GearManager::new());
368 let api = LocalDirectoryClient::new(Arc::clone(&dir));
369
370 let err = api
371 .register_instance(
372 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
373 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
374 .with_labels(labels(&[("bad key", "7")])),
375 )
376 .await
377 .expect_err("an invalid label must be rejected at the store boundary");
378 assert!(
379 err.downcast_ref::<DirectoryInvalidArgument>().is_some(),
380 "store-boundary label rejection must be typed InvalidArgument, got: {err}"
381 );
382 assert!(
383 dir.instances_of("worker").is_empty(),
384 "a rejected registration must not be stored"
385 );
386 }
387
388 #[tokio::test]
389 async fn test_register_and_resolve_rest_and_openapi() {
390 let dir = Arc::new(GearManager::new());
391 let api = LocalDirectoryClient::new(dir.clone());
392
393 let instance_id = Uuid::new_v4();
394 let register_info = RegisterInstanceInfo::new("billing", instance_id.to_string())
395 .with_version("1.0.0")
396 .with_rest_endpoint(ServiceEndpoint::http("billing", 8080))
397 .with_openapi_spec("{\"openapi\":\"3.1.0\"}");
398
399 api.register_instance(register_info).await.unwrap();
400
401 let resolved = api.resolve_rest_service("billing").await.unwrap();
403 assert_eq!(resolved.uri, concat!("http", "://billing:8080"));
404
405 let spec = api.get_openapi_spec("billing").await.unwrap();
407 assert!(spec.contains("openapi"));
408 }
409
410 #[tokio::test]
411 async fn test_resolve_rest_and_openapi_not_found() {
412 let dir = Arc::new(GearManager::new());
413 let api = LocalDirectoryClient::new(dir);
414
415 let rest_err = api.resolve_rest_service("missing").await.unwrap_err();
418 assert!(
419 rest_err.downcast_ref::<DirectoryNotFound>().is_some(),
420 "expected the typed not-found sentinel, got: {rest_err:?}"
421 );
422
423 let spec_err = api.get_openapi_spec("missing").await.unwrap_err();
424 assert!(
425 spec_err.downcast_ref::<DirectoryNotFound>().is_some(),
426 "expected the typed not-found sentinel, got: {spec_err:?}"
427 );
428 }
429
430 #[tokio::test]
431 async fn test_deregister_instance_via_api() {
432 let dir = Arc::new(GearManager::new());
433 let api = LocalDirectoryClient::new(dir.clone());
434
435 let instance_id = Uuid::new_v4();
436 let inst = Arc::new(GearInstance::new("test_gear", instance_id));
438 dir.register_instance(inst).unwrap();
439
440 assert_eq!(dir.instances_of("test_gear").len(), 1);
442
443 api.deregister_instance("test_gear", &instance_id.to_string())
445 .await
446 .unwrap();
447
448 assert_eq!(dir.instances_of("test_gear").len(), 0);
450 }
451
452 #[tokio::test]
453 async fn test_send_heartbeat_via_api() {
454 use crate::runtime::InstanceState;
455
456 let dir = Arc::new(GearManager::new());
457 let api = LocalDirectoryClient::new(dir.clone());
458
459 let instance_id = Uuid::new_v4();
460 let inst = Arc::new(GearInstance::new("test_gear", instance_id));
462 dir.register_instance(inst).unwrap();
463
464 let instances = dir.instances_of("test_gear");
466 assert_eq!(instances[0].state(), InstanceState::Registered);
467
468 api.send_heartbeat("test_gear", &instance_id.to_string())
470 .await
471 .unwrap();
472
473 let instances = dir.instances_of("test_gear");
475 assert_eq!(instances[0].state(), InstanceState::Healthy);
476 }
477
478 #[tokio::test]
479 async fn test_list_all_instances_across_gears() {
480 let dir = Arc::new(GearManager::new());
481 let api = LocalDirectoryClient::new(Arc::clone(&dir));
482
483 for (gear, port) in [("billing", 8080u16), ("catalog", 8081u16)] {
485 api.register_instance(
486 RegisterInstanceInfo::new(gear, Uuid::new_v4().to_string())
487 .with_version("1.0.0")
488 .with_rest_endpoint(ServiceEndpoint::http(gear, port))
489 .with_openapi_spec(format!("{{\"openapi\":\"3.1.0\",\"x\":\"{gear}\"}}")),
490 )
491 .await
492 .unwrap();
493 }
494
495 api.register_instance(
499 RegisterInstanceInfo::new("reporting", Uuid::new_v4().to_string())
500 .with_grpc_services(vec![(
501 "reporting.Service".to_owned(),
502 ServiceEndpoint::new("http://reporting:7000"),
503 )])
504 .with_version("1.0.0"),
505 )
506 .await
507 .unwrap();
508
509 let all = api.list_all_instances().await.unwrap();
510 assert_eq!(all.len(), 3);
511
512 let billing = all.iter().find(|i| i.gear == "billing").expect("billing");
513 assert_eq!(
514 billing.rest_endpoint.as_ref().map(|e| e.uri.as_str()),
515 Some("http://billing:8080")
516 );
517 assert!(billing.openapi_spec_hash.is_some());
521 assert!(
522 api.get_openapi_spec("billing")
523 .await
524 .expect("billing spec")
525 .contains("billing")
526 );
527
528 let reporting = all
531 .iter()
532 .find(|i| i.gear == "reporting")
533 .expect("reporting");
534 assert_eq!(
535 reporting.endpoint.as_ref().map(|e| e.uri.as_str()),
536 Some("http://reporting:7000")
537 );
538 assert!(reporting.rest_endpoint.is_none());
539 assert!(reporting.openapi_spec_hash.is_none());
540 }
541
542 fn labels(pairs: &[(&str, &str)]) -> std::collections::BTreeMap<String, String> {
543 pairs
544 .iter()
545 .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
546 .collect()
547 }
548
549 #[tokio::test]
550 async fn register_carries_labels_into_list_instances() {
551 let dir = Arc::new(GearManager::new());
552 let api = LocalDirectoryClient::new(Arc::clone(&dir));
553
554 api.register_instance(
555 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
556 .with_grpc_services(vec![(
557 "worker.Svc".to_owned(),
558 ServiceEndpoint::new("http://worker:7000"),
559 )])
560 .with_labels(labels(&[("shard", "7")])),
561 )
562 .await
563 .unwrap();
564
565 let listed = api.list_instances("worker").await.unwrap();
566 assert_eq!(listed.len(), 1);
567 assert_eq!(listed[0].labels.get("shard"), Some(&"7".to_owned()));
568 }
569
570 #[tokio::test]
571 async fn resolve_by_labels_selects_matching_instances() {
572 let dir = Arc::new(GearManager::new());
573 let api = LocalDirectoryClient::new(Arc::clone(&dir));
574
575 for (id, shard) in [("a", "7"), ("b", "8"), ("c", "7")] {
576 api.register_instance(
577 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
578 .with_grpc_services(vec![(
579 format!("worker.{id}"),
580 ServiceEndpoint::new(format!("http://worker-{id}:7000")),
581 )])
582 .with_labels(labels(&[("shard", shard)])),
583 )
584 .await
585 .unwrap();
586 }
587
588 let matched = api
589 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
590 .await
591 .unwrap();
592 assert_eq!(matched.len(), 2, "two instances carry shard=7");
593 assert!(
594 matched
595 .iter()
596 .all(|i| i.labels.get("shard") == Some(&"7".to_owned()))
597 );
598 }
599
600 #[tokio::test]
601 async fn labelless_reregister_preserves_labels_through_api() {
602 let dir = Arc::new(GearManager::new());
603 let api = LocalDirectoryClient::new(Arc::clone(&dir));
604
605 let instance_id = Uuid::new_v4().to_string();
606 api.register_instance(
607 RegisterInstanceInfo::new("worker", instance_id.clone())
608 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
609 .with_labels(labels(&[("shard", "7")])),
610 )
611 .await
612 .unwrap();
613
614 api.register_instance(
616 RegisterInstanceInfo::new("worker", instance_id)
617 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
618 .with_version("2.0.0"),
619 )
620 .await
621 .unwrap();
622
623 let listed = api.list_instances("worker").await.unwrap();
624 assert_eq!(listed.len(), 1);
625 assert_eq!(
626 listed[0].labels.get("shard"),
627 Some(&"7".to_owned()),
628 "label-less re-registration must not wipe stored labels"
629 );
630
631 let matched = api
633 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
634 .await
635 .unwrap();
636 assert_eq!(matched.len(), 1);
637 }
638
639 #[tokio::test]
640 async fn resolve_by_labels_carries_live_serving_state() {
641 let dir = Arc::new(GearManager::new());
642 let api = LocalDirectoryClient::new(Arc::clone(&dir));
643
644 let instance_id = Uuid::new_v4();
645 api.register_instance(
646 RegisterInstanceInfo::new("worker", instance_id.to_string())
647 .with_grpc_services(vec![(
648 "worker.Svc".to_owned(),
649 ServiceEndpoint::new("http://worker:7000"),
650 )])
651 .with_labels(labels(&[("shard", "7")])),
652 )
653 .await
654 .unwrap();
655
656 let matched = api
658 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
659 .await
660 .unwrap();
661 assert_eq!(matched.len(), 1);
662 assert_eq!(matched[0].state, InstanceState::Registered);
663 assert!(!matched[0].state.is_serving());
664
665 dir.update_heartbeat("worker", instance_id, std::time::Instant::now());
668 let matched = api
669 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
670 .await
671 .unwrap();
672 assert_eq!(matched[0].state, InstanceState::Healthy);
673 assert!(matched[0].state.is_serving());
674 }
675
676 #[tokio::test]
677 async fn resolve_by_labels_omits_openapi_spec() {
678 let dir = Arc::new(GearManager::new());
679 let api = LocalDirectoryClient::new(Arc::clone(&dir));
680
681 api.register_instance(
682 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
683 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
684 .with_openapi_spec("{\"openapi\":\"3.1.0\"}")
685 .with_labels(labels(&[("shard", "7")])),
686 )
687 .await
688 .unwrap();
689
690 let matched = api
692 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
693 .await
694 .unwrap();
695 assert_eq!(matched.len(), 1);
696 assert!(
697 matched[0].openapi_spec_hash.is_some(),
698 "in-process resolve_by_labels carries only the spec hash, never the document"
699 );
700
701 let listed = api.list_instances("worker").await.unwrap();
704 assert!(
705 listed[0].openapi_spec_hash.is_some(),
706 "list_instances must carry the spec hash, never the document"
707 );
708 }
709
710 #[tokio::test]
711 async fn list_all_instances_skips_endpoint_less() {
712 let dir = Arc::new(GearManager::new());
713 let api = LocalDirectoryClient::new(Arc::clone(&dir));
714
715 api.register_instance(
718 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
719 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080")),
720 )
721 .await
722 .unwrap();
723 api.register_instance(RegisterInstanceInfo::new(
724 "placeholder",
725 Uuid::new_v4().to_string(),
726 ))
727 .await
728 .unwrap();
729
730 let all = api.list_all_instances().await.unwrap();
731 assert_eq!(all.len(), 1, "the endpoint-less instance is skipped");
732 assert_eq!(all[0].gear, "worker");
733 assert!(
734 all.iter()
735 .all(|i| i.endpoint.as_ref().is_some_and(|e| !e.uri.is_empty())),
736 "no instance may carry an absent or empty-URI endpoint"
737 );
738 }
739
740 #[tokio::test]
741 async fn label_path_returns_endpoint_less_match() {
742 let dir = Arc::new(GearManager::new());
743 let api = LocalDirectoryClient::new(Arc::clone(&dir));
744
745 api.register_instance(
750 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
751 .with_labels(labels(&[("shard", "7")])),
752 )
753 .await
754 .unwrap();
755
756 let matched = api
757 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
758 .await
759 .unwrap();
760 assert_eq!(
761 matched.len(),
762 1,
763 "resolve_by_labels must return a matched instance even with no endpoint"
764 );
765 assert!(
766 matched[0].endpoint.is_none(),
767 "a no-endpoint match carries endpoint = None, not an empty-URI sentinel"
768 );
769 assert!(matched[0].rest_endpoint.is_none());
770 assert_eq!(matched[0].labels.get("shard"), Some(&"7".to_owned()));
771
772 let listed = api.list_instances("worker").await.unwrap();
774 assert_eq!(
775 listed.len(),
776 1,
777 "list_instances must not drop endpoint-less instances"
778 );
779 assert!(listed[0].endpoint.is_none());
780 }
781}