Fix unwrap function in the Prometheus server code (#446)

TraceMachina · Dec 7, 2023 · 406eab7 · 406eab7
1 parent ab0f1ac
commit 406eab7
Showing 1 changed file with 14 additions and 17 deletions.
diff --git a/src/bin/cas.rs b/src/bin/cas.rs
@@ -24,7 +24,7 @@ use error::{make_err, Code, Error, ResultExt};
 use futures::future::{select_all, BoxFuture, OptionFuture, TryFutureExt};
 use futures::FutureExt;
 use hyper::server::conn::Http;
-use hyper::{Body, Response};
+use hyper::{Response, StatusCode};
 use native_link_config::cas_server::{
     CasConfig, CompressionAlgorithm, ConfigDigestHashFunction, GlobalConfig, ServerConfig, WorkerConfig,
 };
@@ -349,11 +349,10 @@ async fn inner_main(cfg: CasConfig, server_start_timestamp: u64) -> Result<(), B
             .route_service("/status", axum::routing::get(move || async move { "Ok".to_string() }));
 
         if let Some(prometheus_cfg) = services.prometheus {
-            fn error_to_response<E: std::error::Error>(e: E) -> hyper::Response<Body> {
-                hyper::Response::builder()
-                    .status(500)
-                    .body(format!("Error: {e:?}").into())
-                    .unwrap()
+            fn error_to_response<E: std::error::Error>(e: E) -> Response<String> {
+                let mut response = Response::new(format!("Error: {e:?}"));
+                *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
+                response
             }
             let path = if prometheus_cfg.path.is_empty() {
                 DEFAULT_PROMETHEUS_METRICS_PATH
@@ -375,17 +374,15 @@ async fn inner_main(cfg: CasConfig, server_start_timestamp: u64) -> Result<(), B
                                 // This is a hack to get around this bug: https://github.com/prometheus/client_rust/issues/155
                                 buf = buf.replace("native_link_native_link_stores_", "");
                                 buf = buf.replace("native_link_native_link_workers_", "");
-                                let body = Body::from(buf);
-                                Response::builder()
-                                    .header(
-                                        hyper::header::CONTENT_TYPE,
-                                        // Per spec we should probably use `application/openmetrics-text; version=1.0.0; charset=utf-8`
-                                        // https://github.com/OpenObservability/OpenMetrics/blob/1386544931307dff279688f332890c31b6c5de36/specification/OpenMetrics.md#overall-structure
-                                        // However, this makes debugging more difficult, so we use the old text/plain instead.
-                                        "text/plain; version=0.0.4; charset=utf-8",
-                                    )
-                                    .body(body)
-                                    .unwrap()
+                                let mut response = Response::new(buf);
+                                // Per spec we should probably use `application/openmetrics-text; version=1.0.0; charset=utf-8`
+                                // https://github.com/OpenObservability/OpenMetrics/blob/1386544931307dff279688f332890c31b6c5de36/specification/OpenMetrics.md#overall-structure
+                                // However, this makes debugging more difficult, so we use the old text/plain instead.
+                                response.headers_mut().insert(
+                                    hyper::header::CONTENT_TYPE,
+                                    hyper::header::HeaderValue::from_static("text/plain; version=0.0.4; charset=utf-8"),
+                                );
+                                response
                             })
                             .unwrap_or_else(error_to_response)
                     })