1use anyhow::Result;
4use async_trait::async_trait;
5use std::sync::Arc;
6use uuid::Uuid;
7
8use crate::runtime::{Endpoint, GearInstance, GearManager};
9
10pub use cf_system_sdks::directory::{
12 DirectoryClient, RegisterInstanceInfo, ServiceEndpoint, ServiceInstanceInfo,
13};
14
15pub struct LocalDirectoryClient {
20 mgr: Arc<GearManager>,
21}
22
23impl LocalDirectoryClient {
24 #[must_use]
25 pub fn new(mgr: Arc<GearManager>) -> Self {
26 Self { mgr }
27 }
28}
29
30#[async_trait]
31impl DirectoryClient for LocalDirectoryClient {
32 async fn resolve_grpc_service(&self, service_name: &str) -> Result<ServiceEndpoint> {
33 if let Some((_gear, _inst, ep)) = self.mgr.pick_service_round_robin(service_name) {
34 return Ok(ServiceEndpoint::new(ep.uri));
35 }
36
37 anyhow::bail!("Service not found or no healthy instances: {service_name}")
38 }
39
40 async fn resolve_rest_service(&self, gear_name: &str) -> Result<ServiceEndpoint> {
41 if let Some(ep) = self.mgr.pick_rest_endpoint_round_robin(gear_name) {
42 return Ok(ServiceEndpoint::new(ep.uri));
43 }
44
45 anyhow::bail!("No REST endpoint registered for gear: {gear_name}")
46 }
47
48 async fn get_openapi_spec(&self, gear_name: &str) -> Result<String> {
49 self.mgr
50 .openapi_spec_of(gear_name)
51 .ok_or_else(|| anyhow::anyhow!("No OpenAPI spec registered for gear: {gear_name}"))
52 }
53
54 async fn list_instances(&self, gear: &str) -> Result<Vec<ServiceInstanceInfo>> {
55 let mut result = Vec::new();
56
57 for inst in self.mgr.instances_of(gear) {
58 if let Some((_, ep)) = inst.grpc_services.iter().next() {
59 result.push(ServiceInstanceInfo {
60 gear: gear.to_owned(),
61 instance_id: inst.instance_id.to_string(),
62 endpoint: ServiceEndpoint::new(ep.uri.clone()),
63 version: inst.version.clone(),
64 rest_endpoint: inst
65 .rest_endpoint
66 .as_ref()
67 .map(|ep| ServiceEndpoint::new(ep.uri.clone())),
68 });
69 }
70 }
71
72 Ok(result)
73 }
74
75 async fn register_instance(&self, info: RegisterInstanceInfo) -> Result<()> {
76 let instance_id = Uuid::parse_str(&info.instance_id)
78 .map_err(|e| anyhow::anyhow!("Invalid instance_id '{}': {}", info.instance_id, e))?;
79
80 let mut instance = GearInstance::new(info.gear.clone(), instance_id);
82
83 if let Some(version) = info.version {
85 instance = instance.with_version(version);
86 }
87
88 for (service_name, endpoint) in info.grpc_services {
90 instance = instance.with_grpc_service(service_name, Endpoint::from_uri(endpoint.uri));
91 }
92
93 if let Some(rest) = info.rest_endpoint {
95 instance = instance.with_rest_endpoint(Endpoint::from_uri(rest.uri));
96 }
97
98 if let Some(spec) = info.openapi_spec {
100 instance = instance.with_openapi_spec(spec);
101 }
102
103 self.mgr.register_instance(Arc::new(instance));
105
106 Ok(())
107 }
108
109 async fn deregister_instance(&self, gear: &str, instance_id: &str) -> Result<()> {
110 let instance_id = Uuid::parse_str(instance_id)
111 .map_err(|e| anyhow::anyhow!("Invalid instance_id '{instance_id}': {e}"))?;
112 self.mgr.deregister(gear, instance_id);
113 Ok(())
114 }
115
116 async fn send_heartbeat(&self, gear: &str, instance_id: &str) -> Result<()> {
117 let instance_id = Uuid::parse_str(instance_id)
118 .map_err(|e| anyhow::anyhow!("Invalid instance_id '{instance_id}': {e}"))?;
119 self.mgr
120 .update_heartbeat(gear, instance_id, std::time::Instant::now());
121 Ok(())
122 }
123}
124
125#[cfg(test)]
126#[cfg_attr(coverage_nightly, coverage(off))]
127mod tests {
128 use super::*;
129
130 #[tokio::test]
131 async fn test_resolve_grpc_service_not_found() {
132 let dir = Arc::new(GearManager::new());
133 let api = LocalDirectoryClient::new(dir);
134
135 let result = api.resolve_grpc_service("nonexistent.Service").await;
136 assert!(result.is_err());
137 }
138
139 #[tokio::test]
140 async fn test_register_instance_via_api() {
141 let dir = Arc::new(GearManager::new());
142 let api = LocalDirectoryClient::new(dir.clone());
143
144 let instance_id = Uuid::new_v4();
145 let register_info = RegisterInstanceInfo {
147 gear: "test_gear".to_owned(),
148 instance_id: instance_id.to_string(),
149 grpc_services: vec![(
150 "test.Service".to_owned(),
151 ServiceEndpoint::http("127.0.0.1", 8001),
152 )],
153 version: Some("1.0.0".to_owned()),
154 rest_endpoint: None,
155 openapi_spec: None,
156 };
157
158 api.register_instance(register_info).await.unwrap();
159
160 let instances = dir.instances_of("test_gear");
162 assert_eq!(instances.len(), 1);
163 assert_eq!(instances[0].instance_id, instance_id);
164 assert_eq!(instances[0].version, Some("1.0.0".to_owned()));
165 assert!(instances[0].grpc_services.contains_key("test.Service"));
166 }
167
168 #[tokio::test]
169 async fn test_register_and_resolve_rest_and_openapi() {
170 let dir = Arc::new(GearManager::new());
171 let api = LocalDirectoryClient::new(dir.clone());
172
173 let instance_id = Uuid::new_v4();
174 let register_info = RegisterInstanceInfo {
175 gear: "billing".to_owned(),
176 instance_id: instance_id.to_string(),
177 grpc_services: vec![],
178 version: Some("1.0.0".to_owned()),
179 rest_endpoint: Some(ServiceEndpoint::http("billing", 8080)),
180 openapi_spec: Some("{\"openapi\":\"3.1.0\"}".to_owned()),
181 };
182
183 api.register_instance(register_info).await.unwrap();
184
185 let resolved = api.resolve_rest_service("billing").await.unwrap();
187 assert_eq!(resolved.uri, concat!("http", "://billing:8080"));
188
189 let spec = api.get_openapi_spec("billing").await.unwrap();
191 assert!(spec.contains("openapi"));
192 }
193
194 #[tokio::test]
195 async fn test_resolve_rest_and_openapi_not_found() {
196 let dir = Arc::new(GearManager::new());
197 let api = LocalDirectoryClient::new(dir);
198
199 assert!(api.resolve_rest_service("missing").await.is_err());
200 assert!(api.get_openapi_spec("missing").await.is_err());
201 }
202
203 #[tokio::test]
204 async fn test_deregister_instance_via_api() {
205 let dir = Arc::new(GearManager::new());
206 let api = LocalDirectoryClient::new(dir.clone());
207
208 let instance_id = Uuid::new_v4();
209 let inst = Arc::new(GearInstance::new("test_gear", instance_id));
211 dir.register_instance(inst);
212
213 assert_eq!(dir.instances_of("test_gear").len(), 1);
215
216 api.deregister_instance("test_gear", &instance_id.to_string())
218 .await
219 .unwrap();
220
221 assert_eq!(dir.instances_of("test_gear").len(), 0);
223 }
224
225 #[tokio::test]
226 async fn test_send_heartbeat_via_api() {
227 use crate::runtime::InstanceState;
228
229 let dir = Arc::new(GearManager::new());
230 let api = LocalDirectoryClient::new(dir.clone());
231
232 let instance_id = Uuid::new_v4();
233 let inst = Arc::new(GearInstance::new("test_gear", instance_id));
235 dir.register_instance(inst);
236
237 let instances = dir.instances_of("test_gear");
239 assert_eq!(instances[0].state(), InstanceState::Registered);
240
241 api.send_heartbeat("test_gear", &instance_id.to_string())
243 .await
244 .unwrap();
245
246 let instances = dir.instances_of("test_gear");
248 assert_eq!(instances[0].state(), InstanceState::Healthy);
249 }
250}