optimize header name handling in Grpc::map_response (#1359)

* tonic: optimize MetadataMap::into_sanitized_headers * tonic: optimize Status::add_header --------- Co-authored-by: tottoto <[email protected]>
hyperium · Sep 20, 2024 · a09d453 · a09d453
1 parent 99b663e
commit a09d453
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 17 deletions.
diff --git a/tonic/src/metadata/map.rs b/tonic/src/metadata/map.rs
@@ -1,3 +1,5 @@
+use http::HeaderName;
+
 pub(crate) use self::as_encoding_agnostic_metadata_key::AsEncodingAgnosticMetadataKey;
 pub(crate) use self::as_metadata_key::AsMetadataKey;
 pub(crate) use self::into_metadata_key::IntoMetadataKey;
@@ -200,13 +202,13 @@ pub(crate) const GRPC_TIMEOUT_HEADER: &str = "grpc-timeout";
 
 impl MetadataMap {
     // Headers reserved by the gRPC protocol.
-    pub(crate) const GRPC_RESERVED_HEADERS: [&'static str; 6] = [
-        "te",
-        "user-agent",
-        "content-type",
-        "grpc-message",
-        "grpc-message-type",
-        "grpc-status",
+    pub(crate) const GRPC_RESERVED_HEADERS: [HeaderName; 6] = [
+        HeaderName::from_static("te"),
+        HeaderName::from_static("user-agent"),
+        HeaderName::from_static("content-type"),
+        HeaderName::from_static("grpc-message"),
+        HeaderName::from_static("grpc-message-type"),
+        HeaderName::from_static("grpc-status"),
     ];
 
     /// Create an empty `MetadataMap`.
@@ -251,7 +253,7 @@ impl MetadataMap {
 
     pub(crate) fn into_sanitized_headers(mut self) -> http::HeaderMap {
         for r in &Self::GRPC_RESERVED_HEADERS {
-            self.headers.remove(*r);
+            self.headers.remove(r);
         }
         self.headers
     }

diff --git a/tonic/src/request.rs b/tonic/src/request.rs
@@ -438,15 +438,19 @@ pub(crate) enum SanitizeHeaders {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::metadata::{MetadataKey, MetadataValue};
+
     use http::Uri;
 
     #[test]
     fn reserved_headers_are_excluded() {
         let mut r = Request::new(1);
 
         for header in &MetadataMap::GRPC_RESERVED_HEADERS {
-            r.metadata_mut()
-                .insert(*header, MetadataValue::from_static("invalid"));
+            r.metadata_mut().insert(
+                MetadataKey::unchecked_from_header_name(header.clone()),
+                MetadataValue::from_static("invalid"),
+            );
         }
 
         let http_request = r.into_http(

diff --git a/tonic/src/response.rs b/tonic/src/response.rs
@@ -130,15 +130,17 @@ impl<T> Response<T> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::metadata::MetadataValue;
+    use crate::metadata::{MetadataKey, MetadataValue};
 
     #[test]
     fn reserved_headers_are_excluded() {
         let mut r = Response::new(1);
 
         for header in &MetadataMap::GRPC_RESERVED_HEADERS {
-            r.metadata_mut()
-                .insert(*header, MetadataValue::from_static("invalid"));
+            r.metadata_mut().insert(
+                MetadataKey::unchecked_from_header_name(header.clone()),
+                MetadataValue::from_static("invalid"),
+            );
         }
 
         let http_response = r.into_http();

diff --git a/tonic/src/status.rs b/tonic/src/status.rs
@@ -2,7 +2,10 @@ use crate::metadata::MetadataMap;
 use crate::{body::BoxBody, metadata::GRPC_CONTENT_TYPE};
 use base64::Engine as _;
 use bytes::Bytes;
-use http::header::{HeaderMap, HeaderValue};
+use http::{
+    header::{HeaderMap, HeaderValue},
+    HeaderName,
+};
 use percent_encoding::{percent_decode, percent_encode, AsciiSet, CONTROLS};
 use std::{borrow::Cow, error::Error, fmt, sync::Arc};
 use tracing::{debug, trace, warn};
@@ -18,9 +21,9 @@ const ENCODING_SET: &AsciiSet = &CONTROLS
     .add(b'{')
     .add(b'}');
 
-const GRPC_STATUS_HEADER_CODE: &str = "grpc-status";
-const GRPC_STATUS_MESSAGE_HEADER: &str = "grpc-message";
-const GRPC_STATUS_DETAILS_HEADER: &str = "grpc-status-details-bin";
+const GRPC_STATUS_HEADER_CODE: HeaderName = HeaderName::from_static("grpc-status");
+const GRPC_STATUS_MESSAGE_HEADER: HeaderName = HeaderName::from_static("grpc-message");
+const GRPC_STATUS_DETAILS_HEADER: HeaderName = HeaderName::from_static("grpc-status-details-bin");
 
 /// A gRPC status describing the result of an RPC call.
 ///