/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;

/// 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 {
    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;
                    HealthStatusDescription {
                        namespace: namespace.clone(),
                        status: status_res.unwrap_or_else(|_| {4
                            let struct_name = indicator.struct_name();
                            warn!(struct_name, "Timeout during health check");
                            HealthStatus::Timeout { struct_name }
                        }),
                    }
                },
            ))
            // 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-07-21 15:28

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