Skip to main content

toolkit/
directory.rs

1//! Directory API - contract for service discovery and instance resolution
2
3use anyhow::Result;
4use async_trait::async_trait;
5use std::sync::Arc;
6use uuid::Uuid;
7
8use crate::runtime::{Endpoint, GearInstance, GearManager};
9
10// Re-export all types from contracts - this is the single source of truth
11pub use cf_system_sdks::directory::{
12    DirectoryClient, RegisterInstanceInfo, ServiceEndpoint, ServiceInstanceInfo,
13};
14
15/// Local implementation of `DirectoryClient` that delegates to `GearManager`
16///
17/// This is the in-process implementation used by gears running in the same
18/// process as the gear orchestrator.
19pub 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        // Parse instance_id from string to Uuid
77        let instance_id = Uuid::parse_str(&info.instance_id)
78            .map_err(|e| anyhow::anyhow!("Invalid instance_id '{}': {}", info.instance_id, e))?;
79
80        // Build a GearInstance from RegisterInstanceInfo
81        let mut instance = GearInstance::new(info.gear.clone(), instance_id);
82
83        // Apply version if provided
84        if let Some(version) = info.version {
85            instance = instance.with_version(version);
86        }
87
88        // Add all gRPC services
89        for (service_name, endpoint) in info.grpc_services {
90            instance = instance.with_grpc_service(service_name, Endpoint::from_uri(endpoint.uri));
91        }
92
93        // Apply REST endpoint if provided
94        if let Some(rest) = info.rest_endpoint {
95            instance = instance.with_rest_endpoint(Endpoint::from_uri(rest.uri));
96        }
97
98        // Apply OpenAPI spec if provided
99        if let Some(spec) = info.openapi_spec {
100            instance = instance.with_openapi_spec(spec);
101        }
102
103        // Register the instance with the manager
104        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        // Register an instance through the API
146        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        // Verify the instance was registered
161        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        // REST endpoint resolves to the registered base URL.
186        let resolved = api.resolve_rest_service("billing").await.unwrap();
187        assert_eq!(resolved.uri, concat!("http", "://billing:8080"));
188
189        // OpenAPI spec can be retrieved.
190        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        // Register an instance first
210        let inst = Arc::new(GearInstance::new("test_gear", instance_id));
211        dir.register_instance(inst);
212
213        // Verify it exists
214        assert_eq!(dir.instances_of("test_gear").len(), 1);
215
216        // Deregister via API
217        api.deregister_instance("test_gear", &instance_id.to_string())
218            .await
219            .unwrap();
220
221        // Verify it's gone
222        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        // Register an instance first
234        let inst = Arc::new(GearInstance::new("test_gear", instance_id));
235        dir.register_instance(inst);
236
237        // Verify initial state is Registered
238        let instances = dir.instances_of("test_gear");
239        assert_eq!(instances[0].state(), InstanceState::Registered);
240
241        // Send heartbeat via API
242        api.send_heartbeat("test_gear", &instance_id.to_string())
243            .await
244            .unwrap();
245
246        // Verify state transitioned to Healthy
247        let instances = dir.instances_of("test_gear");
248        assert_eq!(instances[0].state(), InstanceState::Healthy);
249    }
250}