1use std::collections::HashMap;
4use std::fmt;
5use std::sync::Arc;
6
7use tokio::sync::Mutex;
8use tokio::task::JoinHandle;
9use torrust_net_primitives::service_binding::ServiceBinding;
10
11pub type ServiceHeathCheckResult = Result<String, String>;
13
14#[derive(Debug)]
18pub struct ServiceHealthCheckJob {
19 pub info: String,
20 pub job: JoinHandle<ServiceHeathCheckResult>,
21}
22
23impl ServiceHealthCheckJob {
24 #[must_use]
25 pub fn new(info: String, job: JoinHandle<ServiceHeathCheckResult>) -> Self {
26 Self { info, job }
27 }
28}
29
30pub type FnSpawnServiceHeathCheck = fn(&ServiceBinding) -> ServiceHealthCheckJob;
34
35#[derive(Clone, Debug)]
40pub struct ServiceRegistration<M> {
41 service_binding: ServiceBinding,
42 metadata: M,
43 health_check: Option<FnSpawnServiceHeathCheck>,
44}
45
46impl<M> ServiceRegistration<M> {
47 #[must_use]
48 pub fn new(service_binding: ServiceBinding, metadata: M, health_check: Option<FnSpawnServiceHeathCheck>) -> Self {
49 Self {
50 service_binding,
51 metadata,
52 health_check,
53 }
54 }
55
56 #[must_use]
57 pub fn service_binding(&self) -> &ServiceBinding {
58 &self.service_binding
59 }
60
61 #[must_use]
62 pub fn metadata(&self) -> &M {
63 &self.metadata
64 }
65
66 #[must_use]
67 pub fn spawn_check(&self) -> Option<ServiceHealthCheckJob> {
68 self.health_check.map(|health_check| health_check(&self.service_binding))
69 }
70}
71
72#[derive(Clone, Debug)]
74pub struct RegisteredService<M> {
75 registration: ServiceRegistration<M>,
76}
77
78impl<M> RegisteredService<M> {
79 #[must_use]
80 pub fn service_binding(&self) -> &ServiceBinding {
81 self.registration.service_binding()
82 }
83
84 #[must_use]
85 pub fn metadata(&self) -> &M {
86 self.registration.metadata()
87 }
88
89 #[must_use]
90 pub fn spawn_check(&self) -> Option<ServiceHealthCheckJob> {
91 self.registration.spawn_check()
92 }
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
97pub enum RegistrationError {
98 DuplicateBinding(ServiceBinding),
100}
101
102impl fmt::Display for RegistrationError {
103 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
104 match self {
105 Self::DuplicateBinding(service_binding) => {
106 write!(formatter, "a service is already registered for binding {service_binding}")
107 }
108 }
109 }
110}
111
112impl std::error::Error for RegistrationError {}
113
114#[derive(Debug)]
120pub struct ServiceRegistrationForm<M> {
121 registar: Registar<M>,
122}
123
124impl<M> ServiceRegistrationForm<M> {
125 pub async fn register(self, registration: ServiceRegistration<M>) -> Result<(), RegistrationError> {
135 self.registar.insert(registration).await
136 }
137}
138
139#[derive(Debug)]
141pub struct Registar<M = ()> {
142 registry: Arc<Mutex<HashMap<ServiceBinding, ServiceRegistration<M>>>>,
143}
144
145impl<M> Clone for Registar<M> {
146 fn clone(&self) -> Self {
147 Self {
148 registry: self.registry.clone(),
149 }
150 }
151}
152
153impl<M> Default for Registar<M> {
154 fn default() -> Self {
155 Self {
156 registry: Arc::default(),
157 }
158 }
159}
160
161impl<M> Registar<M> {
162 #[must_use]
164 pub fn give_form(&self) -> ServiceRegistrationForm<M> {
165 ServiceRegistrationForm { registar: self.clone() }
166 }
167
168 async fn insert(&self, service_registration: ServiceRegistration<M>) -> Result<(), RegistrationError> {
169 let mut mutex = self.registry.lock().await;
170
171 if mutex.contains_key(service_registration.service_binding()) {
172 return Err(RegistrationError::DuplicateBinding(
173 service_registration.service_binding().clone(),
174 ));
175 }
176
177 mutex.insert(service_registration.service_binding.clone(), service_registration);
178
179 Ok(())
180 }
181
182 pub async fn deregister(&self, service_binding: &ServiceBinding) -> Option<RegisteredService<M>> {
193 self.registry
194 .lock()
195 .await
196 .remove(service_binding)
197 .map(|registration| RegisteredService { registration })
198 }
199
200 pub async fn services(&self) -> Vec<RegisteredService<M>>
205 where
206 M: Clone,
207 {
208 let mutex = self.registry.lock().await;
209 let mut services: Vec<_> = mutex
210 .values()
211 .cloned()
212 .map(|registration| RegisteredService { registration })
213 .collect();
214 services.sort_by(|left, right| {
215 protocol_sort_key(&left.service_binding().protocol())
216 .cmp(&protocol_sort_key(&right.service_binding().protocol()))
217 .then_with(|| {
218 left.service_binding()
219 .bind_address()
220 .cmp(&right.service_binding().bind_address())
221 })
222 });
223 services
224 }
225
226 pub async fn services_matching<F>(&self, predicate: F) -> Vec<RegisteredService<M>>
228 where
229 M: Clone,
230 F: Fn(&M) -> bool,
231 {
232 self.services()
233 .await
234 .into_iter()
235 .filter(|service| predicate(service.metadata()))
236 .collect()
237 }
238}
239
240fn protocol_sort_key(protocol: &torrust_net_primitives::service_binding::Protocol) -> u8 {
241 match protocol {
242 torrust_net_primitives::service_binding::Protocol::UDP => 0,
243 torrust_net_primitives::service_binding::Protocol::HTTP => 1,
244 torrust_net_primitives::service_binding::Protocol::HTTPS => 2,
245 }
246}
247
248#[cfg(test)]
249mod tests {
250 use std::net::{Ipv4Addr, SocketAddr};
251
252 use torrust_net_primitives::service_binding::Protocol;
253
254 use super::{Registar, RegistrationError, ServiceRegistration};
255
256 fn binding(protocol: Protocol, port: u16) -> torrust_net_primitives::service_binding::ServiceBinding {
257 torrust_net_primitives::service_binding::ServiceBinding::new(protocol, SocketAddr::from((Ipv4Addr::LOCALHOST, port)))
258 .expect("test binding should be valid")
259 }
260
261 #[tokio::test]
262 async fn it_should_make_a_registration_visible_after_acknowledgement() {
263 let registar = Registar::default();
264
265 registar
266 .give_form()
267 .register(ServiceRegistration::new(binding(Protocol::HTTP, 8000), "first", None))
268 .await
269 .expect("registration should succeed");
270
271 assert_eq!(registar.services().await[0].metadata(), &"first");
272 }
273
274 #[tokio::test]
275 async fn it_should_return_services_in_deterministic_binding_order() {
276 let registar = Registar::default();
277
278 registar
279 .give_form()
280 .register(ServiceRegistration::new(binding(Protocol::HTTP, 9000), "second", None))
281 .await
282 .expect("registration should succeed");
283 registar
284 .give_form()
285 .register(ServiceRegistration::new(binding(Protocol::HTTP, 8000), "first", None))
286 .await
287 .expect("registration should succeed");
288
289 let metadata: Vec<_> = registar
290 .services()
291 .await
292 .into_iter()
293 .map(|service| *service.metadata())
294 .collect();
295
296 assert_eq!(metadata, ["first", "second"]);
297 }
298
299 #[tokio::test]
300 async fn it_should_order_services_by_protocol_then_final_binding() {
301 let registar = Registar::default();
302
303 for (protocol, port, metadata) in [
304 (Protocol::HTTPS, 8000, "https"),
305 (Protocol::HTTP, 9000, "http-second"),
306 (Protocol::UDP, 9000, "udp-second"),
307 (Protocol::HTTP, 8000, "http-first"),
308 (Protocol::UDP, 8000, "udp-first"),
309 ] {
310 registar
311 .give_form()
312 .register(ServiceRegistration::new(binding(protocol, port), metadata, None))
313 .await
314 .expect("registration should succeed");
315 }
316
317 let metadata: Vec<_> = registar
318 .services()
319 .await
320 .into_iter()
321 .map(|service| *service.metadata())
322 .collect();
323
324 assert_eq!(metadata, ["udp-first", "udp-second", "http-first", "http-second", "https"]);
325 }
326
327 #[tokio::test]
328 async fn it_should_reject_duplicate_final_bindings() {
329 let registar = Registar::default();
330 let service_binding = binding(Protocol::HTTP, 8000);
331
332 registar
333 .give_form()
334 .register(ServiceRegistration::new(service_binding.clone(), (), None))
335 .await
336 .expect("initial registration should succeed");
337
338 let error = registar
339 .give_form()
340 .register(ServiceRegistration::new(service_binding.clone(), (), None))
341 .await
342 .expect_err("duplicate registration should fail");
343
344 assert_eq!(error, RegistrationError::DuplicateBinding(service_binding));
345 }
346
347 #[tokio::test]
348 async fn it_should_remove_a_registration_and_return_it_when_deregistered() {
349 let registar = Registar::default();
350 let service_binding = binding(Protocol::HTTP, 8000);
351
352 registar
353 .give_form()
354 .register(ServiceRegistration::new(service_binding.clone(), "first", None))
355 .await
356 .expect("registration should succeed");
357
358 let removed = registar.deregister(&service_binding).await;
359
360 assert_eq!(removed.map(|service| *service.metadata()), Some("first"));
361 assert!(registar.services().await.is_empty());
362 }
363
364 #[tokio::test]
365 async fn it_should_allow_registering_a_binding_again_after_it_was_deregistered() {
366 let registar = Registar::default();
367 let service_binding = binding(Protocol::HTTP, 8000);
368
369 registar
370 .give_form()
371 .register(ServiceRegistration::new(service_binding.clone(), "first", None))
372 .await
373 .expect("initial registration should succeed");
374 registar.deregister(&service_binding).await;
375
376 registar
377 .give_form()
378 .register(ServiceRegistration::new(service_binding, "second", None))
379 .await
380 .expect("registration after deregistration should succeed");
381
382 assert_eq!(registar.services().await[0].metadata(), &"second");
383 }
384
385 #[tokio::test]
386 async fn it_should_return_none_when_the_binding_to_deregister_is_not_registered() {
387 let registar = Registar::<()>::default();
388
389 let removed = registar.deregister(&binding(Protocol::HTTP, 8000)).await;
390
391 assert!(removed.is_none());
392 }
393
394 #[tokio::test]
395 async fn it_should_keep_the_other_registrations_when_one_is_deregistered() {
396 let registar = Registar::default();
397 let removed_binding = binding(Protocol::HTTP, 8000);
398
399 for (port, metadata) in [(8000, "removed"), (9000, "kept")] {
400 registar
401 .give_form()
402 .register(ServiceRegistration::new(binding(Protocol::HTTP, port), metadata, None))
403 .await
404 .expect("registration should succeed");
405 }
406
407 registar.deregister(&removed_binding).await;
408
409 let metadata: Vec<_> = registar
410 .services()
411 .await
412 .into_iter()
413 .map(|service| *service.metadata())
414 .collect();
415 assert_eq!(metadata, ["kept"]);
416 }
417}