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.is_none());
463 assert!(
464 api.get_openapi_spec("billing")
465 .await
466 .expect("billing spec")
467 .contains("billing")
468 );
469
470 let reporting = all
473 .iter()
474 .find(|i| i.gear == "reporting")
475 .expect("reporting");
476 assert_eq!(
477 reporting.endpoint.as_ref().map(|e| e.uri.as_str()),
478 Some("http://reporting:7000")
479 );
480 assert!(reporting.rest_endpoint.is_none());
481 assert!(reporting.openapi_spec.is_none());
482 }
483
484 fn labels(pairs: &[(&str, &str)]) -> std::collections::BTreeMap<String, String> {
485 pairs
486 .iter()
487 .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
488 .collect()
489 }
490
491 #[tokio::test]
492 async fn register_carries_labels_into_list_instances() {
493 let dir = Arc::new(GearManager::new());
494 let api = LocalDirectoryClient::new(Arc::clone(&dir));
495
496 api.register_instance(
497 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
498 .with_grpc_services(vec![(
499 "worker.Svc".to_owned(),
500 ServiceEndpoint::new("http://worker:7000"),
501 )])
502 .with_labels(labels(&[("shard", "7")])),
503 )
504 .await
505 .unwrap();
506
507 let listed = api.list_instances("worker").await.unwrap();
508 assert_eq!(listed.len(), 1);
509 assert_eq!(listed[0].labels.get("shard"), Some(&"7".to_owned()));
510 }
511
512 #[tokio::test]
513 async fn resolve_by_labels_selects_matching_instances() {
514 let dir = Arc::new(GearManager::new());
515 let api = LocalDirectoryClient::new(Arc::clone(&dir));
516
517 for (id, shard) in [("a", "7"), ("b", "8"), ("c", "7")] {
518 api.register_instance(
519 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
520 .with_grpc_services(vec![(
521 format!("worker.{id}"),
522 ServiceEndpoint::new(format!("http://worker-{id}:7000")),
523 )])
524 .with_labels(labels(&[("shard", shard)])),
525 )
526 .await
527 .unwrap();
528 }
529
530 let matched = api
531 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
532 .await
533 .unwrap();
534 assert_eq!(matched.len(), 2, "two instances carry shard=7");
535 assert!(
536 matched
537 .iter()
538 .all(|i| i.labels.get("shard") == Some(&"7".to_owned()))
539 );
540 }
541
542 #[tokio::test]
543 async fn labelless_reregister_preserves_labels_through_api() {
544 let dir = Arc::new(GearManager::new());
545 let api = LocalDirectoryClient::new(Arc::clone(&dir));
546
547 let instance_id = Uuid::new_v4().to_string();
548 api.register_instance(
549 RegisterInstanceInfo::new("worker", instance_id.clone())
550 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
551 .with_labels(labels(&[("shard", "7")])),
552 )
553 .await
554 .unwrap();
555
556 api.register_instance(
558 RegisterInstanceInfo::new("worker", instance_id)
559 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
560 .with_version("2.0.0"),
561 )
562 .await
563 .unwrap();
564
565 let listed = api.list_instances("worker").await.unwrap();
566 assert_eq!(listed.len(), 1);
567 assert_eq!(
568 listed[0].labels.get("shard"),
569 Some(&"7".to_owned()),
570 "label-less re-registration must not wipe stored labels"
571 );
572
573 let matched = api
575 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
576 .await
577 .unwrap();
578 assert_eq!(matched.len(), 1);
579 }
580
581 #[tokio::test]
582 async fn resolve_by_labels_carries_live_serving_state() {
583 let dir = Arc::new(GearManager::new());
584 let api = LocalDirectoryClient::new(Arc::clone(&dir));
585
586 let instance_id = Uuid::new_v4();
587 api.register_instance(
588 RegisterInstanceInfo::new("worker", instance_id.to_string())
589 .with_grpc_services(vec![(
590 "worker.Svc".to_owned(),
591 ServiceEndpoint::new("http://worker:7000"),
592 )])
593 .with_labels(labels(&[("shard", "7")])),
594 )
595 .await
596 .unwrap();
597
598 let matched = api
600 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
601 .await
602 .unwrap();
603 assert_eq!(matched.len(), 1);
604 assert_eq!(matched[0].state, InstanceState::Registered);
605 assert!(!matched[0].state.is_serving());
606
607 dir.update_heartbeat("worker", instance_id, std::time::Instant::now());
610 let matched = api
611 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
612 .await
613 .unwrap();
614 assert_eq!(matched[0].state, InstanceState::Healthy);
615 assert!(matched[0].state.is_serving());
616 }
617
618 #[tokio::test]
619 async fn resolve_by_labels_omits_openapi_spec() {
620 let dir = Arc::new(GearManager::new());
621 let api = LocalDirectoryClient::new(Arc::clone(&dir));
622
623 api.register_instance(
624 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
625 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080"))
626 .with_openapi_spec("{\"openapi\":\"3.1.0\"}")
627 .with_labels(labels(&[("shard", "7")])),
628 )
629 .await
630 .unwrap();
631
632 let matched = api
634 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
635 .await
636 .unwrap();
637 assert_eq!(matched.len(), 1);
638 assert!(
639 matched[0].openapi_spec.is_none(),
640 "in-process resolve_by_labels must not attach the OpenAPI document"
641 );
642
643 let listed = api.list_instances("worker").await.unwrap();
646 assert!(
647 listed[0].openapi_spec.is_none(),
648 "list_instances must not attach the OpenAPI document"
649 );
650 assert!(
651 listed[0].openapi_spec_hash.is_some(),
652 "list_instances must still carry the spec hash"
653 );
654 }
655
656 #[tokio::test]
657 async fn list_all_instances_skips_endpoint_less() {
658 let dir = Arc::new(GearManager::new());
659 let api = LocalDirectoryClient::new(Arc::clone(&dir));
660
661 api.register_instance(
664 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
665 .with_rest_endpoint(ServiceEndpoint::new("http://worker:8080")),
666 )
667 .await
668 .unwrap();
669 api.register_instance(RegisterInstanceInfo::new(
670 "placeholder",
671 Uuid::new_v4().to_string(),
672 ))
673 .await
674 .unwrap();
675
676 let all = api.list_all_instances().await.unwrap();
677 assert_eq!(all.len(), 1, "the endpoint-less instance is skipped");
678 assert_eq!(all[0].gear, "worker");
679 assert!(
680 all.iter()
681 .all(|i| i.endpoint.as_ref().is_some_and(|e| !e.uri.is_empty())),
682 "no instance may carry an absent or empty-URI endpoint"
683 );
684 }
685
686 #[tokio::test]
687 async fn label_path_returns_endpoint_less_match() {
688 let dir = Arc::new(GearManager::new());
689 let api = LocalDirectoryClient::new(Arc::clone(&dir));
690
691 api.register_instance(
696 RegisterInstanceInfo::new("worker", Uuid::new_v4().to_string())
697 .with_labels(labels(&[("shard", "7")])),
698 )
699 .await
700 .unwrap();
701
702 let matched = api
703 .resolve_by_labels("worker", &LabelSelector::new().with("shard", "7"))
704 .await
705 .unwrap();
706 assert_eq!(
707 matched.len(),
708 1,
709 "resolve_by_labels must return a matched instance even with no endpoint"
710 );
711 assert!(
712 matched[0].endpoint.is_none(),
713 "a no-endpoint match carries endpoint = None, not an empty-URI sentinel"
714 );
715 assert!(matched[0].rest_endpoint.is_none());
716 assert_eq!(matched[0].labels.get("shard"), Some(&"7".to_owned()));
717
718 let listed = api.list_instances("worker").await.unwrap();
720 assert_eq!(
721 listed.len(),
722 1,
723 "list_instances must not drop endpoint-less instances"
724 );
725 assert!(listed[0].endpoint.is_none());
726 }
727}