/build/source/nativelink-util/src/health_utils.rs

Source
// Copyright 2024 The Native Link Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::fmt::Debug;
use core::pin::Pin;
use core::time::Duration;
use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use futures::{Stream, StreamExt};
use parking_lot::Mutex;
use serde::Serialize;
use tokio::time::timeout;
use tracing::warn;

use crate::metrics::record_health_check;

/// Struct name health indicator component.
type StructName = str;
/// Readable message status of the health indicator.
type Message = str;

#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
pub enum HealthStatus {
    Ok {
        struct_name: &'static StructName,
        message: Cow<'static, Message>,
    },
    Initializing {
        struct_name: &'static StructName,
        message: Cow<'static, Message>,
    },
    /// This status is used to indicate a non-fatal issue with the component.
    Warning {
        struct_name: &'static StructName,
        message: Cow<'static, Message>,
    },
    Failed {
        struct_name: &'static StructName,
        message: Cow<'static, Message>,
    },
    Timeout {
        struct_name: &'static StructName,
    },
}

impl HealthStatus {
    /// Stable label for metrics. Part of the wire contract, so dashboards
    /// break if these strings change.
    #[must_use]
    pub const fn metric_label(&self) -> &'static str {
        match self {
            Self::Ok { .. } => "ok",
            Self::Initializing { .. } => "initializing",
            Self::Warning { .. } => "warning",
            Self::Failed { .. } => "failed",
            Self::Timeout { .. } => "timeout",
        }
    }

    pub fn new_ok(
        component: &(impl HealthStatusIndicator + ?Sized),
        message: Cow<'static, str>,
    ) -> Self {
        Self::Ok {
            struct_name: component.struct_name(),
            message,
        }
    }

    pub fn new_initializing(
        component: &(impl HealthStatusIndicator + ?Sized),
        message: Cow<'static, str>,
    ) -> Self {
        Self::Initializing {
            struct_name: component.struct_name(),
            message,
        }
    }

    pub fn new_warning(
        component: &(impl HealthStatusIndicator + ?Sized),
        message: Cow<'static, str>,
    ) -> Self {
        Self::Warning {
            struct_name: component.struct_name(),
            message,
        }
    }

    pub fn new_failed(
        component: &(impl HealthStatusIndicator + ?Sized),
        message: Cow<'static, str>,
    ) -> Self {
        Self::Failed {
            struct_name: component.struct_name(),
            message,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
pub struct HealthStatusDescription {
    pub namespace: Cow<'static, str>,
    pub status: HealthStatus,
}

/// Health status indicator trait. This trait is used to define
/// a health status indicator by implementing the `check_health` function.
/// A default implementation is provided for the `check_health` function
/// that returns healthy component.
#[async_trait]
pub trait HealthStatusIndicator: Sync + Send + Unpin {
    fn get_name(&self) -> &'static str;

    /// Returns the name of the struct implementing the trait.
    fn struct_name(&self) -> &'static str {
        core::any::type_name::<Self>()
    }

    /// Check the health status of the component. This function should be
    /// implemented by the component to check the health status of the component.
    async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus;
}

type HealthRegistryBuilderState =
    Arc<Mutex<HashMap<Cow<'static, str>, Arc<dyn HealthStatusIndicator>>>>;

pub struct HealthRegistryBuilder {
    namespace: Cow<'static, str>,
    state: HealthRegistryBuilderState,
}

impl Debug for HealthRegistryBuilder {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("HealthRegistryBuilder")
            .field("namespace", &self.namespace)
            .finish_non_exhaustive()
    }
}

/// Health registry builder that is used to build a health registry.
/// The builder provides creation, registering of health status indicators,
/// sub-building scoped health registries and building the health registry.
/// `build()` should be called once for finalizing the production of a health registry.
impl HealthRegistryBuilder {
    pub fn new(namespace: &str) -> Self {
        Self {
            namespace: format!("/{namespace}").into(),
            state: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Register a health status indicator at current namespace.
    pub fn register_indicator(&mut self, indicator: Arc<dyn HealthStatusIndicator>) {
        let name = format!("{}/{}", self.namespace, indicator.get_name());
        self.state.lock().insert(name.into(), indicator);
    }

    /// Create a sub builder for a namespace.
    #[must_use]
    pub fn sub_builder(&mut self, namespace: &str) -> Self {
        Self {
            namespace: format!("{}/{}", self.namespace, namespace).into(),
            state: self.state.clone(),
        }
    }

    /// Finalize the production of the health registry.
    pub fn build(&mut self) -> HealthRegistry {
        HealthRegistry {
            indicators: self.state.lock().clone().into_iter().collect(),
        }
    }
}

#[derive(Default, Clone)]
pub struct HealthRegistry {
    indicators: Vec<(Cow<'static, str>, Arc<dyn HealthStatusIndicator>)>,
}

impl Debug for HealthRegistry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("HealthRegistry")
            .field(
                "indicators",
                &self
                    .indicators
                    .iter()
                    .map(|(name, _)| name)
                    .collect::<Vec<_>>(),
            )
            .finish()
    }
}

pub trait HealthStatusReporter {
    fn health_status_report(
        &self,
        timeout: &Duration,
    ) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>>;
}

/// Health status reporter implementation for the health registry that provides a stream
/// of health status descriptions.
///
/// Indicator checks run **in parallel**: each indicator's
/// `check_health` does real I/O against its store (write/has/read
/// roundtrip per [`nativelink_util::store_trait::StoreDriver::check_health`]), so iterating
/// indicators serially makes the total response time `N * timeout`.
/// Under load that easily exceeds Kubernetes liveness probe budgets,
/// causing kubelet to kill an otherwise-healthy pod whose only sin
/// was that one congested store made the probe handler queue behind
/// it. Parallel execution caps the total at ~`timeout` regardless of
/// how many stores are registered.
impl HealthStatusReporter for HealthRegistry {
    fn health_status_report(
        &self,
        timeout_limit: &Duration,
    ) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>> {
        let local_timeout_limit = *timeout_limit;
        Box::pin(
            futures::stream::iter(self.indicators.iter().map(
                move |(namespace, indicator)| async move {
                    let status_res = timeout(
                        local_timeout_limit,
                        indicator.check_health(namespace.clone()),
                    )
                    .await;
                    let status = status_res.unwrap_or_else(|_| {4
                        let struct_name = indicator.struct_name();
                        warn!(struct_name, "Timeout during health check");
                        HealthStatus::Timeout { struct_name }
                    });
                    record_health_check(namespace, status.metric_label());
                    HealthStatusDescription {
                        namespace: namespace.clone(),
                        status,
                    }
                },
            ))
            // Drive every indicator's check concurrently rather than
            // serially. The order of the resulting descriptions is
            // not part of the API contract; collect-into-Vec callers
            // already ignore order.
            .buffer_unordered(usize::MAX),
        )
    }
}

/// Default health status indicator implementation for a component.
/// Generally used for components that don't need custom implementations
/// of the `check_health` function.
#[macro_export]
macro_rules! default_health_status_indicator {
    ($type:ty) => {
        #[async_trait::async_trait]
        impl HealthStatusIndicator for $type {
            fn get_name(&self) -> &'static str {
                stringify!($type)
            }

            async fn check_health(
                &self,
                namespace: std::borrow::Cow<'static, str>,
            ) -> nativelink_util::health_utils::HealthStatus {
                StoreDriver::check_health(Pin::new(self), namespace).await
            }
        }
    };
}

// Re-scoped for the health_utils module.
pub use crate::default_health_status_indicator;

Coverage Report

Created: 2026-09-18 20:40

Line	Count	Source
1		// Copyright 2024 The Native Link Authors. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		use core::fmt::Debug;
16		use core::pin::Pin;
17		use core::time::Duration;
18		use std::borrow::Cow;
19		use std::collections::HashMap;
20		use std::sync::Arc;
21
22		use async_trait::async_trait;
23		use futures::{Stream, StreamExt};
24		use parking_lot::Mutex;
25		use serde::Serialize;
26		use tokio::time::timeout;
27		use tracing::warn;
28
29		use crate::metrics::record_health_check;
30
31		/// Struct name health indicator component.
32		type StructName = str;
33		/// Readable message status of the health indicator.
34		type Message = str;
35
36		#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
37		pub enum HealthStatus {
38		Ok {
39		struct_name: &'static StructName,
40		message: Cow<'static, Message>,
41		},
42		Initializing {
43		struct_name: &'static StructName,
44		message: Cow<'static, Message>,
45		},
46		/// This status is used to indicate a non-fatal issue with the component.
47		Warning {
48		struct_name: &'static StructName,
49		message: Cow<'static, Message>,
50		},
51		Failed {
52		struct_name: &'static StructName,
53		message: Cow<'static, Message>,
54		},
55		Timeout {
56		struct_name: &'static StructName,
57		},
58		}
59
60		impl HealthStatus {
61		/// Stable label for metrics. Part of the wire contract, so dashboards
62		/// break if these strings change.
63		#[must_use]
64	20	pub const fn metric_label(&self) -> &'static str {
65	20	match self {
66	14	Self::Ok { .. } => "ok",
67	0	Self::Initializing { .. } => "initializing",
68	0	Self::Warning { .. } => "warning",
69	1	Self::Failed { .. } => "failed",
70	5	Self::Timeout { .. } => "timeout",
71		}
72	20	}
73
74	6	pub fn new_ok(
75	6	component: &(impl HealthStatusIndicator + ?Sized),
76	6	message: Cow<'static, str>,
77	6	) -> Self {
78	6	Self::Ok {
79	6	struct_name: component.struct_name(),
80	6	message,
81	6	}
82	6	}
83
84	0	pub fn new_initializing(
85	0	component: &(impl HealthStatusIndicator + ?Sized),
86	0	message: Cow<'static, str>,
87	0	) -> Self {
88	0	Self::Initializing {
89	0	struct_name: component.struct_name(),
90	0	message,
91	0	}
92	0	}
93
94	0	pub fn new_warning(
95	0	component: &(impl HealthStatusIndicator + ?Sized),
96	0	message: Cow<'static, str>,
97	0	) -> Self {
98	0	Self::Warning {
99	0	struct_name: component.struct_name(),
100	0	message,
101	0	}
102	0	}
103
104	3	pub fn new_failed(
105	3	component: &(impl HealthStatusIndicator + ?Sized),
106	3	message: Cow<'static, str>,
107	3	) -> Self {
108	3	Self::Failed {
109	3	struct_name: component.struct_name(),
110	3	message,
111	3	}
112	3	}
113		}
114
115		#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
116		pub struct HealthStatusDescription {
117		pub namespace: Cow<'static, str>,
118		pub status: HealthStatus,
119		}
120
121		/// Health status indicator trait. This trait is used to define
122		/// a health status indicator by implementing the `check_health` function.
123		/// A default implementation is provided for the `check_health` function
124		/// that returns healthy component.
125		#[async_trait]
126		pub trait HealthStatusIndicator: Sync + Send + Unpin {
127		fn get_name(&self) -> &'static str;
128
129		/// Returns the name of the struct implementing the trait.
130	12	fn struct_name(&self) -> &'static str {
131	12	core::any::type_name::<Self>()
132	12	}
133
134		/// Check the health status of the component. This function should be
135		/// implemented by the component to check the health status of the component.
136		async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus;
137		}
138
139		type HealthRegistryBuilderState =
140		Arc<Mutex<HashMap<Cow<'static, str>, Arc<dyn HealthStatusIndicator>>>>;
141
142		pub struct HealthRegistryBuilder {
143		namespace: Cow<'static, str>,
144		state: HealthRegistryBuilderState,
145		}
146
147		impl Debug for HealthRegistryBuilder {
148	0	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
149	0	f.debug_struct("HealthRegistryBuilder")
150	0	.field("namespace", &self.namespace)
151	0	.finish_non_exhaustive()
152	0	}
153		}
154
155		/// Health registry builder that is used to build a health registry.
156		/// The builder provides creation, registering of health status indicators,
157		/// sub-building scoped health registries and building the health registry.
158		/// `build()` should be called once for finalizing the production of a health registry.
159		impl HealthRegistryBuilder {
160	10	pub fn new(namespace: &str) -> Self {
161	10	Self {
162	10	namespace: format!("/{namespace}").into(),
163	10	state: Arc::new(Mutex::new(HashMap::new())),
164	10	}
165	10	}
166
167		/// Register a health status indicator at current namespace.
168	17	pub fn register_indicator(&mut self, indicator: Arc<dyn HealthStatusIndicator>) {
169	17	let name = format!("{}/{}", self.namespace, indicator.get_name());
170	17	self.state.lock().insert(name.into(), indicator);
171	17	}
172
173		/// Create a sub builder for a namespace.
174		#[must_use]
175	4	pub fn sub_builder(&mut self, namespace: &str) -> Self {
176	4	Self {
177	4	namespace: format!("{}/{}", self.namespace, namespace).into(),
178	4	state: self.state.clone(),
179	4	}
180	4	}
181
182		/// Finalize the production of the health registry.
183	10	pub fn build(&mut self) -> HealthRegistry {
184	10	HealthRegistry {
185	10	indicators: self.state.lock().clone().into_iter().collect(),
186	10	}
187	10	}
188		}
189
190		#[derive(Default, Clone)]
191		pub struct HealthRegistry {
192		indicators: Vec<(Cow<'static, str>, Arc<dyn HealthStatusIndicator>)>,
193		}
194
195		impl Debug for HealthRegistry {
196	0	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
197	0	f.debug_struct("HealthRegistry")
198	0	.field(
199	0	"indicators",
200	0	&self
201	0	.indicators
202	0	.iter()
203	0	.map(\|(name, _)\| name)
204	0	.collect::<Vec<_>>(),
205		)
206	0	.finish()
207	0	}
208		}
209
210		pub trait HealthStatusReporter {
211		fn health_status_report(
212		&self,
213		timeout: &Duration,
214		) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>>;
215		}
216
217		/// Health status reporter implementation for the health registry that provides a stream
218		/// of health status descriptions.
219		///
220		/// Indicator checks run in parallel: each indicator's
221		/// `check_health` does real I/O against its store (write/has/read
222		/// roundtrip per [`nativelink_util::store_trait::StoreDriver::check_health`]), so iterating
223		/// indicators serially makes the total response time `N * timeout`.
224		/// Under load that easily exceeds Kubernetes liveness probe budgets,
225		/// causing kubelet to kill an otherwise-healthy pod whose only sin
226		/// was that one congested store made the probe handler queue behind
227		/// it. Parallel execution caps the total at ~`timeout` regardless of
228		/// how many stores are registered.
229		impl HealthStatusReporter for HealthRegistry {
230	10	fn health_status_report(
231	10	&self,
232	10	timeout_limit: &Duration,
233	10	) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>> {
234	10	let local_timeout_limit = *timeout_limit;
235	10	Box::pin(
236	10	futures::stream::iter(self.indicators.iter().map(
237	17	move \|(namespace, indicator)\| async move {
238	17	let status_res = timeout(
239	17	local_timeout_limit,
240	17	indicator.check_health(namespace.clone()),
241	17	)
242	17	.await;
243	17	let status = status_res.unwrap_or_else(\|_\| {4
244	4	let struct_name = indicator.struct_name();
245	4	warn!(struct_name, "Timeout during health check");
246	4	HealthStatus::Timeout { struct_name }
247	4	});
248	17	record_health_check(namespace, status.metric_label());
249	17	HealthStatusDescription {
250	17	namespace: namespace.clone(),
251	17	status,
252	17	}
253	34	},
254		))
255		// Drive every indicator's check concurrently rather than
256		// serially. The order of the resulting descriptions is
257		// not part of the API contract; collect-into-Vec callers
258		// already ignore order.
259	10	.buffer_unordered(usize::MAX),
260		)
261	10	}
262		}
263
264		/// Default health status indicator implementation for a component.
265		/// Generally used for components that don't need custom implementations
266		/// of the `check_health` function.
267		#[macro_export]
268		macro_rules! default_health_status_indicator {
269		($type:ty) => {
270		#[async_trait::async_trait]
271		impl HealthStatusIndicator for $type {
272	2	fn get_name(&self) -> &'static str {
273	2	stringify!($type)
274	2	}
275
276		async fn check_health(
277		&self,
278		namespace: std::borrow::Cow<'static, str>,
279		) -> nativelink_util::health_utils::HealthStatus {
280		StoreDriver::check_health(Pin::new(self), namespace).await
281	0	}
282		}
283		};
284		}
285
286		// Re-scoped for the health_utils module.
287		pub use crate::default_health_status_indicator;