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

Source (jump to first uncovered line)
// Copyright 2024 The Native Link Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// 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 std::borrow::Cow;
use std::collections::HashMap;
use std::fmt::Debug;
use std::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use futures::{Stream, StreamExt};
use parking_lot::Mutex;
use serde::Serialize;

/// 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>,
    },
}

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>,
    ) -> HealthStatus {
        Self::Initializing {
            struct_name: component.struct_name(),
            message,
        }
    }

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

    pub fn new_failed(
        component: &(impl HealthStatusIndicator + ?Sized),
        message: Cow<'static, str>,
    ) -> HealthStatus {
        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 {
        std::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,
}

/// 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.
    pub fn sub_builder(&mut self, namespace: &str) -> HealthRegistryBuilder {
        HealthRegistryBuilder {
            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>)>,
}

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

/// Health status reporter implementation for the health registry that provides a stream
/// of health status descriptions.
impl HealthStatusReporter for HealthRegistry {
    fn health_status_report(
        &self,
    ) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>> {
        Box::pin(futures::stream::iter(self.indicators.iter()).then(
            |(namespace, indicator)| async move {
                HealthStatusDescription {
                    namespace: namespace.clone(),
                    status: indicator.check_health(namespace.clone()).await0,
                }
            },
        ))
    }
}

/// 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: 2024-10-22 12:33

Line	Count	Source (jump to first uncovered line)
1		// Copyright 2024 The Native Link Authors. All rights reserved.
2		//
3		// Licensed under the Apache License, Version 2.0 (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		// http://www.apache.org/licenses/LICENSE-2.0
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 std::borrow::Cow;
16		use std::collections::HashMap;
17		use std::fmt::Debug;
18		use std::pin::Pin;
19		use std::sync::Arc;
20
21		use async_trait::async_trait;
22		use futures::{Stream, StreamExt};
23		use parking_lot::Mutex;
24		use serde::Serialize;
25
26		/// Struct name health indicator component.
27		type StructName = str;
28		/// Readable message status of the health indicator.
29		type Message = str;
30
31		#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
32		pub enum HealthStatus {
33		Ok {
34		struct_name: &'static StructName,
35		message: Cow<'static, Message>,
36		},
37		Initializing {
38		struct_name: &'static StructName,
39		message: Cow<'static, Message>,
40		},
41		/// This status is used to indicate a non-fatal issue with the component.
42		Warning {
43		struct_name: &'static StructName,
44		message: Cow<'static, Message>,
45		},
46		Failed {
47		struct_name: &'static StructName,
48		message: Cow<'static, Message>,
49		},
50		}
51
52		impl HealthStatus {
53	0	pub fn new_ok(
54	0	component: &(impl HealthStatusIndicator + ?Sized),
55	0	message: Cow<'static, str>,
56	0	) -> Self {
57	0	Self::Ok {
58	0	struct_name: component.struct_name(),
59	0	message,
60	0	}
61	0	}
62
63	0	pub fn new_initializing(
64	0	component: &(impl HealthStatusIndicator + ?Sized),
65	0	message: Cow<'static, str>,
66	0	) -> HealthStatus {
67	0	Self::Initializing {
68	0	struct_name: component.struct_name(),
69	0	message,
70	0	}
71	0	}
72
73	0	pub fn new_warning(
74	0	component: &(impl HealthStatusIndicator + ?Sized),
75	0	message: Cow<'static, str>,
76	0	) -> HealthStatus {
77	0	Self::Warning {
78	0	struct_name: component.struct_name(),
79	0	message,
80	0	}
81	0	}
82
83	0	pub fn new_failed(
84	0	component: &(impl HealthStatusIndicator + ?Sized),
85	0	message: Cow<'static, str>,
86	0	) -> HealthStatus {
87	0	Self::Failed {
88	0	struct_name: component.struct_name(),
89	0	message,
90	0	}
91	0	}
92		}
93
94		#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
95		pub struct HealthStatusDescription {
96		pub namespace: Cow<'static, str>,
97		pub status: HealthStatus,
98		}
99
100		/// Health status indicator trait. This trait is used to define
101		/// a health status indicator by implementing the `check_health` function.
102		/// A default implementation is provided for the `check_health` function
103		/// that returns healthy component.
104		#[async_trait]
105		pub trait HealthStatusIndicator: Sync + Send + Unpin {
106		fn get_name(&self) -> &'static str;
107
108		/// Returns the name of the struct implementing the trait.
109	0	fn struct_name(&self) -> &'static str {
110	0	std::any::type_name::<Self>()
111	0	}
112
113		/// Check the health status of the component. This function should be
114		/// implemented by the component to check the health status of the component.
115		async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus;
116		}
117
118		type HealthRegistryBuilderState =
119		Arc<Mutex<HashMap<Cow<'static, str>, Arc<dyn HealthStatusIndicator>>>>;
120		pub struct HealthRegistryBuilder {
121		namespace: Cow<'static, str>,
122		state: HealthRegistryBuilderState,
123		}
124
125		/// Health registry builder that is used to build a health registry.
126		/// The builder provides creation, registering of health status indicators,
127		/// sub building scoped health registries and building the health registry.
128		/// `build()` should be called once for finalizing the production of a health registry.
129		impl HealthRegistryBuilder {
130	5	pub fn new(namespace: &str) -> Self {
131	5	Self {
132	5	namespace: format!("/{namespace}").into(),
133	5	state: Arc::new(Mutex::new(HashMap::new())),
134	5	}
135	5	}
136
137		/// Register a health status indicator at current namespace.
138	9	pub fn register_indicator(&mut self, indicator: Arc<dyn HealthStatusIndicator>) {
139	9	let name = format!("{}/{}", self.namespace, indicator.get_name());
140	9	self.state.lock().insert(name.into(), indicator);
141	9	}
142
143		/// Create a sub builder for a namespace.
144	4	pub fn sub_builder(&mut self, namespace: &str) -> HealthRegistryBuilder {
145	4	HealthRegistryBuilder {
146	4	namespace: format!("{}/{}", self.namespace, namespace).into(),
147	4	state: self.state.clone(),
148	4	}
149	4	}
150
151		/// Finalize the production of the health registry.
152	5	pub fn build(&mut self) -> HealthRegistry {
153	5	HealthRegistry {
154	5	indicators: self.state.lock().clone().into_iter().collect(),
155	5	}
156	5	}
157		}
158
159		#[derive(Default, Clone)]
160		pub struct HealthRegistry {
161		indicators: Vec<(Cow<'static, str>, Arc<dyn HealthStatusIndicator>)>,
162		}
163
164		pub trait HealthStatusReporter {
165		fn health_status_report(
166		&self,
167		) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>>;
168		}
169
170		/// Health status reporter implementation for the health registry that provides a stream
171		/// of health status descriptions.
172		impl HealthStatusReporter for HealthRegistry {
173	5	fn health_status_report(
174	5	&self,
175	5	) -> Pin<Box<dyn Stream<Item = HealthStatusDescription> + Send + '_>> {
176	5	Box::pin(futures::stream::iter(self.indicators.iter()).then(
177	9	\|(namespace, indicator)\| async move {
178	9	HealthStatusDescription {
179	9	namespace: namespace.clone(),
180	9	status: indicator.check_health(namespace.clone()). await0 ,
181		}
182	18	},
183	5	))
184	5	}
185		}
186
187		/// Default health status indicator implementation for a component.
188		/// Generally used for components that don't need custom implementations
189		/// of the `check_health` function.
190		#[macro_export]
191		macro_rules! default_health_status_indicator {
192		($type:ty) => {
193		#[async_trait::async_trait]
194		impl HealthStatusIndicator for $type {
195	0	fn get_name(&self) -> &'static str {
196	0	stringify!($type)
197	0	}
198
199	0	async fn check_health(
200	0	&self,
201	0	namespace: std::borrow::Cow<'static, str>,
202	0	) -> nativelink_util::health_utils::HealthStatus {
203	0	StoreDriver::check_health(Pin::new(self), namespace).await
204		}
205		}
206		};
207		}
208
209		// Re-scoped for the health_utils module.
210		pub use crate::default_health_status_indicator;