Skip to content

Commit

Permalink
fix(health): Correctly implement spec for overall health (#897)
Browse files Browse the repository at this point in the history
Co-authored-by: snowpoke <hello@snowpoke.ink>
  • Loading branch information
snowpoke and snowpoke authored Jan 24, 2022
1 parent d574cfd commit 2b0ffee
Showing 1 changed file with 43 additions and 15 deletions.
58 changes: 43 additions & 15 deletions tonic-health/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,10 @@ pub struct HealthReporter {

impl HealthReporter {
fn new() -> Self {
let statuses = Arc::new(RwLock::new(HashMap::new()));
// According to the gRPC Health Check specification, the empty service "" corresponds to the overall server health
let server_status = ("".to_string(), watch::channel(ServingStatus::Serving));

let statuses = Arc::new(RwLock::new(HashMap::from([server_status])));

HealthReporter { statuses }
}
Expand Down Expand Up @@ -166,9 +169,7 @@ mod tests {
use crate::proto::HealthCheckRequest;
use crate::server::{HealthReporter, HealthService};
use crate::ServingStatus;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{watch, RwLock};
use tokio::sync::watch;
use tokio_stream::StreamExt;
use tonic::{Code, Request, Status};

Expand All @@ -183,23 +184,35 @@ mod tests {
}

async fn make_test_service() -> (HealthReporter, HealthService) {
let state = Arc::new(RwLock::new(HashMap::new()));
state.write().await.insert(
"TestService".to_string(),
watch::channel(ServingStatus::Unknown),
);
(
HealthReporter {
statuses: state.clone(),
},
HealthService::new(state.clone()),
)
let health_reporter = HealthReporter::new();

// insert test value
{
let mut statuses = health_reporter.statuses.write().await;
statuses.insert(
"TestService".to_string(),
watch::channel(ServingStatus::Unknown),
);
}

let health_service = HealthService::new(health_reporter.statuses.clone());
(health_reporter, health_service)
}

#[tokio::test]
async fn test_service_check() {
let (mut reporter, service) = make_test_service().await;

// Overall server health
let resp = service
.check(Request::new(HealthCheckRequest {
service: "".to_string(),
}))
.await;
assert!(resp.is_ok());
let resp = resp.unwrap().into_inner();
assert_serving_status(resp.status, ServingStatus::Serving);

// Unregistered service
let resp = service
.check(Request::new(HealthCheckRequest {
Expand Down Expand Up @@ -237,6 +250,21 @@ mod tests {
async fn test_service_watch() {
let (mut reporter, service) = make_test_service().await;

// Overall server health
let resp = service
.watch(Request::new(HealthCheckRequest {
service: "".to_string(),
}))
.await;
assert!(resp.is_ok());
let mut resp = resp.unwrap().into_inner();
let item = resp
.next()
.await
.expect("streamed response is Some")
.expect("response is ok");
assert_serving_status(item.status, ServingStatus::Serving);

// Unregistered service
let resp = service
.watch(Request::new(HealthCheckRequest {
Expand Down

0 comments on commit 2b0ffee

Please sign in to comment.