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