/build/source/nativelink-service/src/health_server.rs

Source
// Copyright 2024 The NativeLink 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::convert::Infallible;
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use core::time::Duration;

use axum::body::Body;
use bytes::Bytes;
use futures::StreamExt;
use http_body_util::Full;
use hyper::header::{CONTENT_TYPE, HeaderValue};
use hyper::{Request, Response, StatusCode};
use nativelink_config::cas_server::HealthConfig;
use nativelink_util::health_utils::{
    HealthRegistry, HealthStatus, HealthStatusDescription, HealthStatusReporter,
};
use tower::Service;
use tracing::error_span;

/// Content type header value for JSON.
const JSON_CONTENT_TYPE: &str = "application/json; charset=utf-8";

// Note: This must be kept in sync with the documentation in
// `HealthConfig::timeout_seconds`.
const DEFAULT_HEALTH_CHECK_TIMEOUT_SECONDS: u64 = 5;

#[derive(Debug, Clone)]
pub struct HealthServer {
    health_registry: HealthRegistry,
    timeout: Duration,
}

impl HealthServer {
    pub const fn new(health_registry: HealthRegistry, health_cfg: &HealthConfig) -> Self {
        let timeout = Duration::from_secs(if health_cfg.timeout_seconds == 0 {
            DEFAULT_HEALTH_CHECK_TIMEOUT_SECONDS
        } else {
            health_cfg.timeout_seconds
        });
        Self {
            health_registry,
            timeout,
        }
    }
}

impl Service<Request<Body>> for HealthServer {
    type Response = Response<Full<Bytes>>;
    type Error = Infallible;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, _req: Request<Body>) -> Self::Future {
        let health_registry = self.health_registry.clone();
        let local_timeout = self.timeout;
        Box::pin(error_span!("health_server_call").in_scope(|| async move {
            let health_status_descriptions: Vec<HealthStatusDescription> = health_registry
                .health_status_report(&local_timeout)
                .collect()
                .await;

            match serde_json5::to_string(&health_status_descriptions) {
                Ok(body) => {
                    let contains_failed_report =
                        health_status_descriptions.iter().any(|description| {2
                            matches!(description.status, HealthStatus::Failed { .. })
                                | matches!1(description.status, HealthStatus::Timeout { .. })
                        });
                    let status_code = if contains_failed_report {
                        StatusCode::SERVICE_UNAVAILABLE
                    } else {
                        StatusCode::OK
                    };

                    Ok(Response::builder()
                        .status(status_code)
                        .header(CONTENT_TYPE, HeaderValue::from_static(JSON_CONTENT_TYPE))
                        .body(Full::new(Bytes::from(body)))
                        .unwrap())
                }

                Err(e) => Ok(Response::builder()
                    .status(StatusCode::INTERNAL_SERVER_ERROR)
                    .header(CONTENT_TYPE, HeaderValue::from_static(JSON_CONTENT_TYPE))
                    .body(Full::new(Bytes::from(format!("Internal Failure: {e:?}"))))
                    .unwrap()),
            }
        }))
    }
}

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

Coverage Report

Created: 2025-10-30 00:14