Skip to main content

torrust_server_lib/
registar.rs

1//! Runtime service registry.
2
3use 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
11/// A [`ServiceHeathCheckResult`] is returned by a completed health check.
12pub type ServiceHeathCheckResult = Result<String, String>;
13
14/// The [`ServiceHealthCheckJob`] has a health check job with it's metadata
15///
16/// The `job` awaits a [`ServiceHeathCheckResult`].
17#[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
30/// The function specification [`FnSpawnServiceHeathCheck`].
31///
32/// A function fulfilling this specification will spawn a new [`ServiceHealthCheckJob`].
33pub type FnSpawnServiceHeathCheck = fn(&ServiceBinding) -> ServiceHealthCheckJob;
34
35/// Immutable data reported by a started local service.
36///
37/// Metadata belongs to the application that uses the registry. The registry
38/// does not assign semantics to it.
39#[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/// A cloneable, immutable view of a registered service.
73#[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/// Registration failure.
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub enum RegistrationError {
98    /// A service already owns this final local binding.
99    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/// A single-use registration capability for one started service.
115///
116/// Obtain one form per successfully bound service with
117/// [`Registar::give_form`]. Consuming [`Self::register`] acknowledges that the
118/// registration is visible in registry snapshots.
119#[derive(Debug)]
120pub struct ServiceRegistrationForm<M> {
121    registar: Registar<M>,
122}
123
124impl<M> ServiceRegistrationForm<M> {
125    /// Inserts a registration and returns only after it is visible to queries.
126    ///
127    /// A caller may treat successful completion as its registry-readiness
128    /// acknowledgement after binding its listener.
129    ///
130    /// # Errors
131    ///
132    /// Returns [`RegistrationError::DuplicateBinding`] when another service is
133    /// already registered with the same final binding.
134    pub async fn register(self, registration: ServiceRegistration<M>) -> Result<(), RegistrationError> {
135        self.registar.insert(registration).await
136    }
137}
138
139/// The [`Registar`] manages immutable runtime service registrations.
140#[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    /// Returns a capability to register one service.
163    #[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    /// Removes the registration for a final binding and returns it.
183    ///
184    /// Returns `None` when no service is registered with that binding. A
185    /// service that stops should remove its own registration, so the same
186    /// binding can be registered again, for example when a stopped test
187    /// environment restarts on a fixed port.
188    ///
189    /// The registry is keyed by binding only, so this removes whichever
190    /// registration holds that binding. Callers should pass only a binding
191    /// they registered themselves.
192    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    /// Returns a deterministic, side-effect-free snapshot of all services.
201    ///
202    /// Results are ordered by protocol, then final socket address, never by
203    /// insertion or hash-map iteration order.
204    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    /// Returns a deterministic, side-effect-free metadata query result.
227    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}