Minimize allocations on header access (connectrpc#445)

The standard library's `http.Header.{Get,Set,Add,Del}` automatically canonicalize header keys. This is nice, but expensive for us - we're using these APIs in many hot paths, and our constants are already in canonical form. This PR adds functions that bypass canonicalization and uses them wherever possible.

---------

Co-authored-by: Akshay Shah <[email protected]>

Author: mattrobenolt

error_writer.go

@@ -83,7 +83,7 @@ func NewErrorWriter(opts ...HandlerOption) *ErrorWriter {
 // IsSupported checks whether a request is using one of the ErrorWriter's
 // supported RPC protocols.
 func (w *ErrorWriter) IsSupported(request *http.Request) bool {
-	ctype := canonicalizeContentType(request.Header.Get(headerContentType))
+	ctype := canonicalizeContentType(getHeaderCanonical(request.Header, headerContentType))
 	_, ok := w.allContentTypes[ctype]
 	return ok
 }
@@ -94,22 +94,22 @@ func (w *ErrorWriter) IsSupported(request *http.Request) bool {
 //
 // Write does not read or close the request body.
 func (w *ErrorWriter) Write(response http.ResponseWriter, request *http.Request, err error) error {
-	ctype := canonicalizeContentType(request.Header.Get(headerContentType))
+	ctype := canonicalizeContentType(getHeaderCanonical(request.Header, headerContentType))
 	if _, ok := w.unaryConnectContentTypes[ctype]; ok {
 		// Unary errors are always JSON.
-		response.Header().Set(headerContentType, connectUnaryContentTypeJSON)
+		setHeaderCanonical(response.Header(), headerContentType, connectUnaryContentTypeJSON)
 		return w.writeConnectUnary(response, err)
 	}
 	if _, ok := w.streamingConnectContentTypes[ctype]; ok {
-		response.Header().Set(headerContentType, ctype)
+		setHeaderCanonical(response.Header(), headerContentType, ctype)
 		return w.writeConnectStreaming(response, err)
 	}
 	if _, ok := w.grpcContentTypes[ctype]; ok {
-		response.Header().Set(headerContentType, ctype)
+		setHeaderCanonical(response.Header(), headerContentType, ctype)
 		return w.writeGRPC(response, err)
 	}
 	if _, ok := w.grpcWebContentTypes[ctype]; ok {
-		response.Header().Set(headerContentType, ctype)
+		setHeaderCanonical(response.Header(), headerContentType, ctype)
 		return w.writeGRPCWeb(response, err)
 	}
 	return fmt.Errorf("unsupported Content-Type %q", ctype)
@@ -153,7 +153,7 @@ func (w *ErrorWriter) writeGRPC(response http.ResponseWriter, err error) error {
 	for k := range trailers {
 		keys = append(keys, k)
 	}
-	response.Header().Set("Trailer", strings.Join(keys, ","))
+	setHeaderCanonical(response.Header(), headerTrailer, strings.Join(keys, ","))
 	response.WriteHeader(http.StatusOK)
 	mergeHeaders(response.Header(), trailers)
 	return nil

handler.go

@@ -190,7 +190,7 @@ func (h *Handler) ServeHTTP(responseWriter http.ResponseWriter, request *http.Re
 	}
 
 	// Find our implementation of the RPC protocol in use.
-	contentType := canonicalizeContentType(request.Header.Get("Content-Type"))
+	contentType := canonicalizeContentType(getHeaderCanonical(request.Header, headerContentType))
 	var protocolHandler protocolHandler
 	for _, handler := range h.protocolHandlers {
 		if _, ok := handler.ContentTypes()[contentType]; ok {
@@ -205,7 +205,7 @@ func (h *Handler) ServeHTTP(responseWriter http.ResponseWriter, request *http.Re
 	}
 
 	// Establish a stream and serve the RPC.
-	request.Header.Set("Content-Type", contentType)                // prefer canonicalized value
+	setHeaderCanonical(request.Header, headerContentType, contentType)
 	ctx, cancel, timeoutErr := protocolHandler.SetTimeout(request) //nolint: contextcheck
 	if timeoutErr != nil {
 		ctx = request.Context()

header.go

@@ -50,3 +50,38 @@ func mergeHeaders(into, from http.Header) {
 		into[k] = append(into[k], vals...)
 	}
 }
+
+// getCanonicalHeader is a shortcut for Header.Get() which
+// bypasses the CanonicalMIMEHeaderKey operation when we
+// know the key is already in canonical form.
+func getHeaderCanonical(h http.Header, key string) string {
+	if h == nil {
+		return ""
+	}
+	v := h[key]
+	if len(v) == 0 {
+		return ""
+	}
+	return v[0]
+}
+
+// setHeaderCanonical is a shortcut for Header.Set() which
+// bypasses the CanonicalMIMEHeaderKey operation when we
+// know the key is already in canonical form.
+func setHeaderCanonical(h http.Header, key, value string) {
+	h[key] = []string{value}
+}
+
+// delHeaderCanonical is a shortcut for Header.Del() which
+// bypasses the CanonicalMIMEHeaderKey operation when we
+// know the key is already in canonical form.
+func delHeaderCanonical(h http.Header, key string) {
+	delete(h, key)
+}
+
+// addHeaderCanonical is a shortcut for Header.Add() which
+// bypasses the CanonicalMIMEHeaderKey operation when we
+// know the key is already in canonical form.
+func addHeaderCanonical(h http.Header, key, value string) {
+	h[key] = append(h[key], value)
+}

protocol.go

@@ -37,6 +37,7 @@ const (
 const (
 	headerContentType = "Content-Type"
 	headerUserAgent   = "User-Agent"
+	headerTrailer     = "Trailer"
 
 	discardLimit = 1024 * 1024 * 4 // 4MiB
 )

protocol_connect.go

@@ -91,7 +91,7 @@ func (h *connectHandler) ContentTypes() map[string]struct{} {
 }
 
 func (*connectHandler) SetTimeout(request *http.Request) (context.Context, context.CancelFunc, error) {
-	timeout := request.Header.Get(connectHeaderTimeout)
+	timeout := getHeaderCanonical(request.Header, connectHeaderTimeout)
 	if timeout == "" {
 		return request.Context(), nil, nil
 	}
@@ -117,11 +117,11 @@ func (h *connectHandler) NewConn(
 	// send the error to the client later on.
 	var contentEncoding, acceptEncoding string
 	if h.Spec.StreamType == StreamTypeUnary {
-		contentEncoding = request.Header.Get(connectUnaryHeaderCompression)
-		acceptEncoding = request.Header.Get(connectUnaryHeaderAcceptCompression)
+		contentEncoding = getHeaderCanonical(request.Header, connectUnaryHeaderCompression)
+		acceptEncoding = getHeaderCanonical(request.Header, connectUnaryHeaderAcceptCompression)
 	} else {
-		contentEncoding = request.Header.Get(connectStreamingHeaderCompression)
-		acceptEncoding = request.Header.Get(connectStreamingHeaderAcceptCompression)
+		contentEncoding = getHeaderCanonical(request.Header, connectStreamingHeaderCompression)
+		acceptEncoding = getHeaderCanonical(request.Header, connectStreamingHeaderAcceptCompression)
 	}
 	requestCompression, responseCompression, failed := negotiateCompression(
 		h.CompressionPools,
@@ -132,7 +132,7 @@ func (h *connectHandler) NewConn(
 		failed = checkServerStreamsCanFlush(h.Spec, responseWriter)
 	}
 	if failed == nil {
-		version := request.Header.Get(connectHeaderProtocolVersion)
+		version := getHeaderCanonical(request.Header, connectHeaderProtocolVersion)
 		if version == "" && h.RequireConnectProtocolHeader {
 			failed = errorf(CodeInvalidArgument, "missing required header: set %s to %q", connectHeaderProtocolVersion, connectProtocolVersion)
 		} else if version != "" && version != connectProtocolVersion {
@@ -148,7 +148,7 @@ func (h *connectHandler) NewConn(
 	// Since we know that these header keys are already in canonical form, we can
 	// skip the normalization in Header.Set.
 	header := responseWriter.Header()
-	header[headerContentType] = []string{request.Header.Get(headerContentType)}
+	header[headerContentType] = []string{getHeaderCanonical(request.Header, headerContentType)}
 	acceptCompressionHeader := connectUnaryHeaderAcceptCompression
 	if h.Spec.StreamType != StreamTypeUnary {
 		acceptCompressionHeader = connectStreamingHeaderAcceptCompression
@@ -164,7 +164,7 @@ func (h *connectHandler) NewConn(
 
 	codecName := connectCodecFromContentType(
 		h.Spec.StreamType,
-		request.Header.Get(headerContentType),
+		getHeaderCanonical(request.Header, headerContentType),
 	)
 	codec := h.Codecs.Get(codecName) // handler.go guarantees this is not nil
 
@@ -247,7 +247,7 @@ func (c *connectClient) Peer() Peer {
 func (c *connectClient) WriteRequestHeader(streamType StreamType, header http.Header) {
 	// We know these header keys are in canonical form, so we can bypass all the
 	// checks in Header.Set.
-	if header.Get(headerUserAgent) == "" {
+	if getHeaderCanonical(header, headerUserAgent) == "" {
 		header[headerUserAgent] = []string{defaultConnectUserAgent}
 	}
 	header[connectHeaderProtocolVersion] = []string{connectProtocolVersion}
@@ -418,7 +418,7 @@ func (cc *connectUnaryClientConn) validateResponse(response *http.Response) *Err
 		}
 		cc.responseTrailer[strings.TrimPrefix(k, connectUnaryTrailerPrefix)] = v
 	}
-	compression := response.Header.Get(connectUnaryHeaderCompression)
+	compression := getHeaderCanonical(response.Header, connectUnaryHeaderCompression)
 	if compression != "" &&
 		compression != compressionIdentity &&
 		!cc.compressionPools.Contains(compression) {
@@ -534,7 +534,7 @@ func (cc *connectStreamingClientConn) validateResponse(response *http.Response)
 	if response.StatusCode != http.StatusOK {
 		return errorf(connectHTTPToCode(response.StatusCode), "HTTP status %v", response.Status)
 	}
-	compression := response.Header.Get(connectStreamingHeaderCompression)
+	compression := getHeaderCanonical(response.Header, connectStreamingHeaderCompression)
 	if compression != "" &&
 		compression != compressionIdentity &&
 		!cc.compressionPools.Contains(compression) {
@@ -605,7 +605,7 @@ func (hc *connectUnaryHandlerConn) Close(err error) error {
 		return hc.request.Body.Close()
 	}
 	// In unary Connect, errors always use application/json.
-	hc.responseWriter.Header().Set(headerContentType, connectUnaryContentTypeJSON)
+	setHeaderCanonical(hc.responseWriter.Header(), headerContentType, connectUnaryContentTypeJSON)
 	hc.responseWriter.WriteHeader(connectCodeToHTTP(CodeOf(err)))
 	data, marshalErr := json.Marshal(newConnectWireError(err))
 	if marshalErr != nil {
@@ -795,7 +795,7 @@ func (m *connectUnaryMarshaler) Marshal(message any) *Error {
 	if m.sendMaxBytes > 0 && compressed.Len() > m.sendMaxBytes {
 		return NewError(CodeResourceExhausted, fmt.Errorf("compressed message size %d exceeds sendMaxBytes %d", compressed.Len(), m.sendMaxBytes))
 	}
-	m.header.Set(connectUnaryHeaderCompression, m.compressionName)
+	setHeaderCanonical(m.header, connectUnaryHeaderCompression, m.compressionName)
 	return m.write(compressed.Bytes())
 }
 

protocol_grpc.go

@@ -131,7 +131,7 @@ func (g *grpcHandler) ContentTypes() map[string]struct{} {
 }
 
 func (*grpcHandler) SetTimeout(request *http.Request) (context.Context, context.CancelFunc, error) {
-	timeout, err := grpcParseTimeout(request.Header.Get(grpcHeaderTimeout))
+	timeout, err := grpcParseTimeout(getHeaderCanonical(request.Header, grpcHeaderTimeout))
 	if err != nil && !errors.Is(err, errNoTimeout) {
 		// Errors here indicate that the client sent an invalid timeout header, so
 		// the error text is safe to send back.
@@ -152,8 +152,8 @@ func (g *grpcHandler) NewConn(
 	// send the error to the client later on.
 	requestCompression, responseCompression, failed := negotiateCompression(
 		g.CompressionPools,
-		request.Header.Get(grpcHeaderCompression),
-		request.Header.Get(grpcHeaderAcceptCompression),
+		getHeaderCanonical(request.Header, grpcHeaderCompression),
+		getHeaderCanonical(request.Header, grpcHeaderAcceptCompression),
 	)
 	if failed == nil {
 		failed = checkServerStreamsCanFlush(g.Spec, responseWriter)
@@ -167,13 +167,13 @@ func (g *grpcHandler) NewConn(
 	// Since we know that these header keys are already in canonical form, we can
 	// skip the normalization in Header.Set.
 	header := responseWriter.Header()
-	header[headerContentType] = []string{request.Header.Get(headerContentType)}
+	header[headerContentType] = []string{getHeaderCanonical(request.Header, headerContentType)}
 	header[grpcHeaderAcceptCompression] = []string{g.CompressionPools.CommaSeparatedNames()}
 	if responseCompression != compressionIdentity {
 		header[grpcHeaderCompression] = []string{responseCompression}
 	}
 
-	codecName := grpcCodecFromContentType(g.web, request.Header.Get(headerContentType))
+	codecName := grpcCodecFromContentType(g.web, getHeaderCanonical(request.Header, headerContentType))
 	codec := g.Codecs.Get(codecName) // handler.go guarantees this is not nil
 	protocolName := ProtocolGRPC
 	if g.web {
@@ -237,7 +237,7 @@ func (g *grpcClient) Peer() Peer {
 func (g *grpcClient) WriteRequestHeader(_ StreamType, header http.Header) {
 	// We know these header keys are in canonical form, so we can bypass all the
 	// checks in Header.Set.
-	if header.Get(headerUserAgent) == "" {
+	if getHeaderCanonical(header, headerUserAgent) == "" {
 		header[headerUserAgent] = []string{defaultGrpcUserAgent}
 	}
 	header[headerContentType] = []string{grpcContentTypeFromCodecName(g.web, g.Codec.Name())}
@@ -365,7 +365,7 @@ func (cc *grpcClientConn) Receive(msg any) error {
 	if err == nil {
 		return nil
 	}
-	if cc.responseHeader.Get(grpcHeaderStatus) != "" {
+	if getHeaderCanonical(cc.responseHeader, grpcHeaderStatus) != "" {
 		// We got what gRPC calls a trailers-only response, which puts the trailing
 		// metadata (including errors) into HTTP headers. validateResponse has
 		// already extracted the error.
@@ -423,7 +423,7 @@ func (cc *grpcClientConn) validateResponse(response *http.Response) *Error {
 	); err != nil {
 		return err
 	}
-	compression := response.Header.Get(grpcHeaderCompression)
+	compression := getHeaderCanonical(response.Header, grpcHeaderCompression)
 	cc.unmarshaler.envelopeReader.compressionPool = cc.compressionPools.Get(compression)
 	return nil
 }
@@ -542,11 +542,13 @@ func (hc *grpcHandlerConn) Close(err error) (retErr error) {
 		// implement http.Flusher, we must pre-declare our HTTP trailers. We can
 		// remove this when Go 1.21 ships and we drop support for Go 1.19.
 		for key := range mergedTrailers {
-			hc.responseWriter.Header().Add("Trailer", key)
+			addHeaderCanonical(hc.responseWriter.Header(), headerTrailer, key)
 		}
 		hc.responseWriter.WriteHeader(http.StatusOK)
 		for key, values := range mergedTrailers {
 			for _, value := range values {
+				// These are potentially user-supplied, so we can't assume they're in
+				// canonical form. Don't use addHeaderCanonical.
 				hc.responseWriter.Header().Add(key, value)
 			}
 		}
@@ -561,6 +563,8 @@ func (hc *grpcHandlerConn) Close(err error) (retErr error) {
 	// logic breaks Envoy's gRPC-Web translation.
 	for key, values := range mergedTrailers {
 		for _, value := range values {
+			// These are potentially user-supplied, so we can't assume they're in
+			// canonical form. Don't use addHeaderCanonical.
 			hc.responseWriter.Header().Add(http.TrailerPrefix+key, value)
 		}
 	}
@@ -636,7 +640,7 @@ func grpcValidateResponse(
 	if response.StatusCode != http.StatusOK {
 		return errorf(grpcHTTPToCode(response.StatusCode), "HTTP status %v", response.Status)
 	}
-	if compression := response.Header.Get(grpcHeaderCompression); compression != "" &&
+	if compression := getHeaderCanonical(response.Header, grpcHeaderCompression); compression != "" &&
 		compression != compressionIdentity &&
 		!availableCompressors.Contains(compression) {
 		// Per https://github.com/grpc/grpc/blob/master/doc/compression.md, we
@@ -658,11 +662,11 @@ func grpcValidateResponse(
 	); err != nil && !errors.Is(err, errTrailersWithoutGRPCStatus) {
 		// Per the specification, only the HTTP status code and Content-Type should
 		// be treated as headers. The rest should be treated as trailing metadata.
-		if contentType := response.Header.Get(headerContentType); contentType != "" {
-			header.Set(headerContentType, contentType)
+		if contentType := getHeaderCanonical(response.Header, headerContentType); contentType != "" {
+			setHeaderCanonical(header, headerContentType, contentType)
 		}
 		mergeHeaders(trailer, response.Header)
-		trailer.Del(headerContentType)
+		delHeaderCanonical(trailer, headerContentType)
 		// Also set the error metadata
 		err.meta = header.Clone()
 		mergeHeaders(err.meta, trailer)
@@ -699,7 +703,7 @@ func grpcHTTPToCode(httpCode int) Code {
 // use a different codec. Consequently, this function needs a Protobuf codec to
 // unmarshal error information in the headers.
 func grpcErrorFromTrailer(bufferPool *bufferPool, protobuf Codec, trailer http.Header) *Error {
-	codeHeader := trailer.Get(grpcHeaderStatus)
+	codeHeader := getHeaderCanonical(trailer, grpcHeaderStatus)
 	if codeHeader == "" {
 		return NewError(CodeInternal, errTrailersWithoutGRPCStatus)
 	}
@@ -711,10 +715,10 @@ func grpcErrorFromTrailer(bufferPool *bufferPool, protobuf Codec, trailer http.H
 	if err != nil {
 		return errorf(CodeInternal, "gRPC protocol error: invalid error code %q", codeHeader)
 	}
-	message := grpcPercentDecode(bufferPool, trailer.Get(grpcHeaderMessage))
+	message := grpcPercentDecode(bufferPool, getHeaderCanonical(trailer, grpcHeaderMessage))
 	retErr := NewWireError(Code(code), errors.New(message))
 
-	detailsBinaryEncoded := trailer.Get(grpcHeaderDetails)
+	detailsBinaryEncoded := getHeaderCanonical(trailer, grpcHeaderDetails)
 	if len(detailsBinaryEncoded) > 0 {
 		detailsBinary, err := DecodeBinaryHeader(detailsBinaryEncoded)
 		if err != nil {
@@ -794,19 +798,21 @@ func grpcContentTypeFromCodecName(web bool, name string) string {
 
 func grpcErrorToTrailer(bufferPool *bufferPool, trailer http.Header, protobuf Codec, err error) {
 	if err == nil {
-		trailer.Set(grpcHeaderStatus, "0") // zero is the gRPC OK status
-		trailer.Set(grpcHeaderMessage, "")
+		setHeaderCanonical(trailer, grpcHeaderStatus, "0") // zero is the gRPC OK status
+		setHeaderCanonical(trailer, grpcHeaderMessage, "")
 		return
 	}
 	status := grpcStatusFromError(err)
 	code := strconv.Itoa(int(status.Code))
 	bin, binErr := protobuf.Marshal(status)
 	if binErr != nil {
-		trailer.Set(
+		setHeaderCanonical(
+			trailer,
 			grpcHeaderStatus,
 			strconv.FormatInt(int64(CodeInternal), 10 /* base */),
 		)
-		trailer.Set(
+		setHeaderCanonical(
+			trailer,
 			grpcHeaderMessage,
 			grpcPercentEncode(
 				bufferPool,
@@ -818,9 +824,9 @@ func grpcErrorToTrailer(bufferPool *bufferPool, trailer http.Header, protobuf Co
 	if connectErr, ok := asError(err); ok {
 		mergeHeaders(trailer, connectErr.meta)
 	}
-	trailer.Set(grpcHeaderStatus, code)
-	trailer.Set(grpcHeaderMessage, grpcPercentEncode(bufferPool, status.Message))
-	trailer.Set(grpcHeaderDetails, EncodeBinaryHeader(bin))
+	setHeaderCanonical(trailer, grpcHeaderStatus, code)
+	setHeaderCanonical(trailer, grpcHeaderMessage, grpcPercentEncode(bufferPool, status.Message))
+	setHeaderCanonical(trailer, grpcHeaderDetails, EncodeBinaryHeader(bin))
 }
 
 func grpcStatusFromError(err error) *statusv1.Status {