Skip to main content

saluki_core/health/
worker.rs

1use async_trait::async_trait;
2use saluki_api::{DynamicRoute, EndpointType};
3use saluki_common::sync::shutdown::ShutdownHandle;
4use saluki_error::generic_error;
5
6use super::HealthRegistry;
7use crate::{
8    diagnostic::DiagnosticHandle,
9    runtime::{state::DataspaceRegistry, InitializationError, Supervisable, SupervisorFuture},
10};
11
12/// A worker that runs the health registry.
13///
14/// This is the only way to run the health registry's liveness probing event loop. The worker
15/// implements [`Supervisable`], so it should be added to a [`Supervisor`][crate::runtime::Supervisor]
16/// to be managed as part of a supervision tree.
17pub struct HealthRegistryWorker {
18    health_registry: HealthRegistry,
19}
20
21impl HealthRegistryWorker {
22    pub(super) fn new(health_registry: HealthRegistry) -> Self {
23        Self { health_registry }
24    }
25}
26
27#[async_trait]
28impl Supervisable for HealthRegistryWorker {
29    fn name(&self) -> &str {
30        "health-registry"
31    }
32
33    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
34        let runner = self.health_registry.clone().into_runner()?;
35
36        let health_routes = DynamicRoute::http(EndpointType::Unprivileged, self.health_registry.api_handler());
37
38        let health_registry = self.health_registry.clone();
39        let flare_handle = DiagnosticHandle::new("health.json", move || health_registry.snapshot_json().into_bytes());
40
41        Ok(Box::pin(async move {
42            let dataspace =
43                DataspaceRegistry::try_current().ok_or_else(|| generic_error!("Dataspace not available."))?;
44
45            // Register our API routes and diagnostic artifact handle before we actually start running.
46            dataspace.assert(health_routes, "health-registry-api");
47            dataspace.assert(flare_handle, "diag-health");
48
49            // We pass the shutdown handle into the runner here, instead of our usual `select! { shutdown => ...,
50            // main_loop_future => ... }` pattern because we try to ensure that we give back the liveness receiver
51            // before the runner completes.
52            //
53            // TODO: We should actually use something like a proper mutex guard so that returning the receiver happens
54            // automatically when the runner future goes out of scope and is dropped, since right now we wouldn't be
55            // able to ensure the current behavior (returning the receiver before the runner completes) happens in the
56            // face of an exceptional error.
57            runner.run(process_shutdown).await;
58
59            Ok(())
60        }))
61    }
62}