rskit_discovery/
factory.rs

1//! Explicit provider factory registry.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use rskit_errors::{AppError, AppResult, ErrorCode};
7
8use crate::config::DiscoveryConfig;
9use crate::traits::{Discovery, Registry};
10
11/// A pair of `(Registry, Discovery)` returned by a provider factory.
12pub type ProviderPair = (Arc<dyn Registry>, Arc<dyn Discovery>);
13
14/// Factory function type: creates a provider pair from a discovery config.
15pub type ProviderFactory = Box<dyn Fn(&DiscoveryConfig) -> AppResult<ProviderPair> + Send + Sync>;
16
17/// Explicit discovery provider registry.
18#[derive(Default)]
19pub struct DiscoveryRegistry {
20    factories: HashMap<String, ProviderFactory>,
21}
22
23impl DiscoveryRegistry {
24    /// Create an empty provider registry.
25    #[must_use]
26    pub fn new() -> Self {
27        Self::default()
28    }
29
30    /// Create a registry with built-in providers registered.
31    #[must_use]
32    pub fn builtins() -> Self {
33        let mut registry = Self::new();
34        registry.register(
35            "static",
36            Box::new(|config| {
37                let mem = Arc::new(crate::memory::InMemoryDiscovery::new());
38                for ep in &config.static_endpoints {
39                    let inst = crate::instance::ServiceInstance {
40                        id: format!("{}-{}:{}", ep.name, ep.address, ep.port),
41                        name: ep.name.clone(),
42                        address: ep.address.clone(),
43                        port: ep.port,
44                        healthy: ep.healthy,
45                        weight: ep.weight,
46                        tags: ep.tags.clone(),
47                        metadata: ep.metadata.clone(),
48                    };
49                    tokio::task::block_in_place(|| {
50                        tokio::runtime::Handle::current().block_on(mem.add(&ep.name, inst))
51                    });
52                }
53                let arc: Arc<crate::memory::InMemoryDiscovery> = mem;
54                Ok((arc.clone(), arc))
55            }),
56        );
57
58        #[cfg(feature = "consul")]
59        registry.register(
60            "consul",
61            Box::new(|config| {
62                let addr = if config.addr.is_empty() {
63                    "localhost:8500"
64                } else {
65                    &config.addr
66                };
67                let token = if config.token.is_empty() {
68                    None
69                } else {
70                    Some(config.token.clone())
71                };
72                let consul = Arc::new(crate::consul::ConsulDiscovery::new(addr, token)?);
73                Ok((consul.clone(), consul))
74            }),
75        );
76
77        registry
78    }
79
80    /// Register a provider factory under the given name.
81    pub fn register(&mut self, name: impl Into<String>, factory: ProviderFactory) {
82        let name = name.into();
83        tracing::debug!(provider = %name, "registered discovery provider factory");
84        self.factories.insert(name, factory);
85    }
86
87    /// Create a provider pair for the provider specified by the config.
88    pub fn create(&self, config: &DiscoveryConfig) -> AppResult<ProviderPair> {
89        let factory = self.factories.get(&config.provider).ok_or_else(|| {
90            AppError::new(
91                ErrorCode::InvalidInput,
92                format!(
93                    "unsupported discovery provider {:?} (registered: {:?})",
94                    config.provider,
95                    self.factories.keys().collect::<Vec<_>>()
96                ),
97            )
98        })?;
99        factory(config)
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn unknown_provider_returns_error() {
109        let registry = DiscoveryRegistry::builtins();
110        let cfg = DiscoveryConfig {
111            provider: "unknown-provider".to_string(),
112            ..Default::default()
113        };
114        let result = registry.create(&cfg);
115        assert!(result.is_err());
116    }
117
118    #[tokio::test]
119    async fn static_provider_creates_successfully() {
120        let registry = DiscoveryRegistry::builtins();
121        let cfg = DiscoveryConfig {
122            provider: "static".to_string(),
123            ..Default::default()
124        };
125        let result = registry.create(&cfg);
126        assert!(result.is_ok());
127    }
128
129    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
130    async fn static_provider_registers_configured_endpoints() {
131        let registry = DiscoveryRegistry::builtins();
132        let cfg = DiscoveryConfig {
133            provider: "static".to_string(),
134            static_endpoints: vec![crate::config::StaticEndpoint {
135                name: "users".to_string(),
136                address: "127.0.0.1".to_string(),
137                port: 8080,
138                protocol: "grpc".to_string(),
139                tags: vec!["blue".to_string()],
140                metadata: [("zone".to_string(), "a".to_string())]
141                    .into_iter()
142                    .collect(),
143                weight: 3,
144                healthy: false,
145            }],
146            ..Default::default()
147        };
148
149        let (_reg, disc) = registry.create(&cfg).unwrap();
150        let instances = disc.resolve("users").await.unwrap();
151
152        assert_eq!(instances.len(), 1);
153        assert_eq!(instances[0].id, "users-127.0.0.1:8080");
154        assert_eq!(instances[0].tags, vec!["blue"]);
155        assert_eq!(
156            instances[0].metadata.get("zone").map(String::as_str),
157            Some("a")
158        );
159        assert_eq!(instances[0].weight, 3);
160        assert!(!instances[0].healthy);
161    }
162
163    #[test]
164    fn explicit_factory_registration_can_override_provider_name() {
165        let mut registry = DiscoveryRegistry::new();
166        registry.register(
167            "custom",
168            Box::new(|_config| {
169                let mem = Arc::new(crate::memory::InMemoryDiscovery::new());
170                Ok((mem.clone(), mem))
171            }),
172        );
173
174        let cfg = DiscoveryConfig {
175            provider: "custom".to_string(),
176            ..Default::default()
177        };
178
179        assert!(registry.create(&cfg).is_ok());
180    }
181
182    #[cfg(feature = "consul")]
183    #[test]
184    fn builtin_consul_factory_uses_default_address_and_optional_token() {
185        let registry = DiscoveryRegistry::builtins();
186        let cfg = DiscoveryConfig {
187            provider: "consul".to_string(),
188            token: "secret".to_string(),
189            ..Default::default()
190        };
191
192        assert!(registry.create(&cfg).is_ok());
193    }
194}
rskit_discovery/factory.rs

rskit_discovery/
factory.rs
