Track successful db roundtrip with vss_service_up

We add `is_healthy` to the `KvStore` API to allow callers to check
whether the backing store passes a simple health check.

The `vss_service_up` gauge metric now tracks `is_healthy`.

Commit written with codex.

Author: tankyleo

rust/api/src/kv_store.rs

@@ -33,4 +33,7 @@ pub trait KvStore: Send + Sync {
 	async fn list_key_versions(
 		&self, user_token: String, request: ListKeyVersionsRequest,
 	) -> Result<ListKeyVersionsResponse, VssError>;
+
+	/// Returns whether the backing store can complete a lightweight health-check round trip.
+	async fn is_healthy(&self) -> bool;
 }

rust/api/src/kv_store_tests.rs

@@ -37,6 +37,7 @@ macro_rules! define_kv_store_tests {
 			};
 		}
 
+		create_test!(is_healthy_should_return_true);
 		create_test!(put_should_succeed_when_single_object_put_operation);
 		create_test!(put_should_succeed_when_multi_object_put_operation);
 		create_test!(put_should_fail_when_key_version_mismatched);
@@ -68,6 +69,13 @@ pub trait KvStoreTestSuite {
 	/// Creates and returns a new instance of the store to be tested.
 	async fn create_store() -> Self::Store;
 
+	async fn is_healthy_should_return_true() -> Result<(), VssError> {
+		let kv_store = Self::create_store().await;
+		assert!(kv_store.is_healthy().await);
+
+		Ok(())
+	}
+
 	async fn put_should_succeed_when_single_object_put_operation() -> Result<(), VssError> {
 		let kv_store = Self::create_store().await;
 		let ctx = TestContext::new(&kv_store);

rust/impls/src/postgres_store.rs

@@ -34,6 +34,7 @@ const KEY_COLUMN: &str = "key";
 const VALUE_COLUMN: &str = "value";
 const VERSION_COLUMN: &str = "version";
 const SORT_ORDER_COLUMN: &str = "sort_order";
+const HEALTH_CHECK_STMT: &str = "SELECT 1";
 
 /// Page token is the `sort_order` value of the last item in the previous page,
 /// encoded as a decimal string.
@@ -739,6 +740,24 @@ where
 
 		Ok(ListKeyVersionsResponse { key_versions, next_page_token, global_version })
 	}
+
+	async fn is_healthy(&self) -> bool {
+		let conn = match self.pool.get().await {
+			Ok(conn) => conn,
+			Err(e) => {
+				debug!("Postgres health check failed while getting connection: {}", e);
+				return false;
+			},
+		};
+
+		match conn.query_one(HEALTH_CHECK_STMT, &[]).await {
+			Ok(row) => row.get::<_, i32>(0) == 1,
+			Err(e) => {
+				debug!("Postgres health check query failed: {}", e);
+				false
+			},
+		}
+	}
 }
 
 #[cfg(test)]

rust/server/src/vss_service.rs

@@ -17,6 +17,7 @@ use api::types::{
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::time::Duration;
 
 use log::{debug, trace};
 
@@ -25,6 +26,7 @@ use crate::util::KeyValueVecKeyPrinter;
 const MAXIMUM_REQUEST_BODY_SIZE: usize = 1024 * 1024 * 1024;
 const PROTOCOL_VERSION_HEADER: &str = "vss-protocol-version";
 const PROTOCOL_VERSION: &str = "0";
+const HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(2);
 
 #[derive(Clone, Copy)]
 pub(crate) struct VssServiceConfig {
@@ -80,14 +82,8 @@ impl Service<Request<Incoming>> for VssService {
 
 		Box::pin(async move {
 			if path == "/metrics" {
-				let response = b"# HELP vss_service_up Is the vss service up?\n# TYPE vss_service_up gauge\nvss_service_up 1\n";
-				return Ok(Response::builder()
-					.status(StatusCode::OK)
-					.header("Content-Type", "text/plain; version=0.0.4")
-					.header(PROTOCOL_VERSION_HEADER, PROTOCOL_VERSION.as_bytes())
-					.body(Full::new(Bytes::from_static(response)))
-					// unwrap safety: body only errors when previous chained calls failed.
-					.unwrap());
+				let is_healthy = check_store_health(&store).await;
+				return Ok(build_metrics_response(is_healthy));
 			}
 
 			let prefix_stripped_path = path.strip_prefix(BASE_PATH_PREFIX).unwrap_or_default();
@@ -257,6 +253,33 @@ fn build_response(status_code: StatusCode, body: Bytes) -> Response<Full<Bytes>>
 		.unwrap()
 }
 
+async fn check_store_health(store: &Arc<dyn KvStore>) -> bool {
+	match tokio::time::timeout(HEALTH_CHECK_TIMEOUT, store.is_healthy()).await {
+		Ok(is_healthy) => is_healthy,
+		Err(_) => false,
+	}
+}
+
+fn build_metrics_response(is_healthy: bool) -> Response<Full<Bytes>> {
+	let response = if is_healthy {
+		Bytes::from_static(
+			b"# HELP vss_service_up Is the vss service up?\n# TYPE vss_service_up gauge\nvss_service_up 1\n",
+		)
+	} else {
+		Bytes::from_static(
+			b"# HELP vss_service_up Is the vss service up?\n# TYPE vss_service_up gauge\nvss_service_up 0\n",
+		)
+	};
+
+	Response::builder()
+		.status(StatusCode::OK)
+		.header("Content-Type", "text/plain; version=0.0.4")
+		.header(PROTOCOL_VERSION_HEADER, PROTOCOL_VERSION.as_bytes())
+		.body(Full::new(response))
+		// unwrap safety: body only errors when previous chained calls failed.
+		.unwrap()
+}
+
 fn build_error_response(e: VssError) -> Response<Full<Bytes>> {
 	let (status_code, error_response) = match e {
 		VssError::NoSuchKeyError(msg) => {
@@ -326,4 +349,36 @@ mod tests {
 			PROTOCOL_VERSION,
 		);
 	}
+
+	#[tokio::test]
+	async fn build_metrics_response_reports_healthy_store() {
+		let response = build_metrics_response(true);
+		assert_eq!(response.status(), StatusCode::OK);
+		assert_eq!(
+			response.headers().get(PROTOCOL_VERSION_HEADER).unwrap().to_str().unwrap(),
+			PROTOCOL_VERSION,
+		);
+
+		let body = response.into_body().collect().await.unwrap().to_bytes();
+		assert_eq!(
+			body.as_ref(),
+			b"# HELP vss_service_up Is the vss service up?\n# TYPE vss_service_up gauge\nvss_service_up 1\n",
+		);
+	}
+
+	#[tokio::test]
+	async fn build_metrics_response_reports_unhealthy_store() {
+		let response = build_metrics_response(false);
+		assert_eq!(response.status(), StatusCode::OK);
+		assert_eq!(
+			response.headers().get(PROTOCOL_VERSION_HEADER).unwrap().to_str().unwrap(),
+			PROTOCOL_VERSION,
+		);
+
+		let body = response.into_body().collect().await.unwrap().to_bytes();
+		assert_eq!(
+			body.as_ref(),
+			b"# HELP vss_service_up Is the vss service up?\n# TYPE vss_service_up gauge\nvss_service_up 0\n",
+		);
+	}
 }