/build/source/nativelink-service/src/health_server.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink 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 core::convert::Infallible; |
16 | | use core::future::Future; |
17 | | use core::pin::Pin; |
18 | | use core::task::{Context, Poll}; |
19 | | |
20 | | use axum::body::Body; |
21 | | use bytes::Bytes; |
22 | | use futures::StreamExt; |
23 | | use http_body_util::Full; |
24 | | use hyper::header::{CONTENT_TYPE, HeaderValue}; |
25 | | use hyper::{Request, Response, StatusCode}; |
26 | | use nativelink_util::health_utils::{ |
27 | | HealthRegistry, HealthStatus, HealthStatusDescription, HealthStatusReporter, |
28 | | }; |
29 | | use tower::Service; |
30 | | use tracing::error_span; |
31 | | |
32 | | /// Content type header value for JSON. |
33 | | const JSON_CONTENT_TYPE: &str = "application/json; charset=utf-8"; |
34 | | |
35 | | #[derive(Debug, Clone)] |
36 | | pub struct HealthServer { |
37 | | health_registry: HealthRegistry, |
38 | | } |
39 | | |
40 | | impl HealthServer { |
41 | 0 | pub const fn new(health_registry: HealthRegistry) -> Self { |
42 | 0 | Self { health_registry } |
43 | 0 | } |
44 | | } |
45 | | |
46 | | impl Service<Request<Body>> for HealthServer { |
47 | | type Response = Response<Full<Bytes>>; |
48 | | type Error = Infallible; |
49 | | type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>; |
50 | | |
51 | 0 | fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
52 | 0 | Poll::Ready(Ok(())) |
53 | 0 | } |
54 | | |
55 | 0 | fn call(&mut self, _req: Request<Body>) -> Self::Future { |
56 | 0 | let health_registry = self.health_registry.clone(); |
57 | 0 | Box::pin(error_span!("health_server_call").in_scope(|| async move { |
58 | 0 | let health_status_descriptions: Vec<HealthStatusDescription> = |
59 | 0 | health_registry.health_status_report().collect().await; |
60 | | |
61 | 0 | match serde_json5::to_string(&health_status_descriptions) { |
62 | 0 | Ok(body) => { |
63 | 0 | let contains_failed_report = |
64 | 0 | health_status_descriptions.iter().any(|description| { |
65 | 0 | matches!(description.status, HealthStatus::Failed { .. }) |
66 | 0 | }); |
67 | 0 | let status_code = if contains_failed_report { Branch (67:42): [True: 0, False: 0]
Branch (67:42): [Folded - Ignored]
|
68 | 0 | StatusCode::SERVICE_UNAVAILABLE |
69 | | } else { |
70 | 0 | StatusCode::OK |
71 | | }; |
72 | | |
73 | 0 | Ok(Response::builder() |
74 | 0 | .status(status_code) |
75 | 0 | .header(CONTENT_TYPE, HeaderValue::from_static(JSON_CONTENT_TYPE)) |
76 | 0 | .body(Full::new(Bytes::from(body))) |
77 | 0 | .unwrap()) |
78 | | } |
79 | | |
80 | 0 | Err(e) => Ok(Response::builder() |
81 | 0 | .status(StatusCode::INTERNAL_SERVER_ERROR) |
82 | 0 | .header(CONTENT_TYPE, HeaderValue::from_static(JSON_CONTENT_TYPE)) |
83 | 0 | .body(Full::new(Bytes::from(format!("Internal Failure: {e:?}")))) |
84 | 0 | .unwrap()), |
85 | | } |
86 | 0 | })) |
87 | 0 | } |
88 | | } |