transport/http: make inbound header preallocation configurable

transport/http/config.go

@@ -170,6 +170,7 @@ func (ts *transportSpec) buildTransport(tc *TransportConfig, k *yarpcconfig.Kit)
 //	      x-bar:
 //	        - X-Bar
 //	        - X-BAR
+//	    headerPreallocation: scan
 //	    shutdownTimeout: 5s
 //	    readHeaderTimeout: 5s
 //	    readTimeout: 10s
@@ -205,6 +206,10 @@ type InboundConfig struct {
 	// Keys must be lowercase. Values are the desired original casings.
 	// Ignored if CanonicalizeHeaderKeys is true.
 	HeaderCaseMapping map[string][]string `config:"headerCaseMapping"`
+	// HeaderPreallocation controls how inbound transport header maps are
+	// preallocated. Supported values are "unfiltered" (default), "scan", and
+	// "disabled".
+	HeaderPreallocation *HeaderPreallocationStrategy `config:"headerPreallocation"`
 }
 
 // TLSConfig specifies the TLS configuration of the HTTP inbound.
@@ -275,6 +280,19 @@ func (ts *transportSpec) buildInbound(ic *InboundConfig, t transport.Transport,
 		}
 	}
 
+	if ic.HeaderPreallocation != nil {
+		if !ic.HeaderPreallocation.valid() {
+			return nil, fmt.Errorf(
+				"headerPreallocation must be one of unfiltered, scan, or disabled, got %q",
+				*ic.HeaderPreallocation,
+			)
+		}
+		inboundOptions = append(
+			inboundOptions,
+			InboundHeaderPreallocation(*ic.HeaderPreallocation),
+		)
+	}
+
 	inboundOptions = append(inboundOptions, DisableHTTP2(ic.DisableHTTP2))
 
 	return t.(*Transport).NewInbound(ic.Address, inboundOptions...), nil

transport/http/config_test.go

@@ -88,6 +88,7 @@ func TestTransportSpec(t *testing.T) {
 		IdleTimeout            time.Duration
 		CanonicalizeHeaderKeys bool
 		HeaderCaseMapping      map[string][]string
+		HeaderPreallocation    HeaderPreallocationStrategy
 	}
 
 	type inboundTest struct {
@@ -349,6 +350,77 @@ func TestTransportSpec(t *testing.T) {
 				CanonicalizeHeaderKeys: true,
 			},
 		},
+		{
+			desc: "header preallocation scan",
+			cfg: attrs{
+				"address":             ":8080",
+				"headerPreallocation": "scan",
+			},
+			wantInbound: &wantInbound{
+				Address:             ":8080",
+				ShutdownTimeout:     defaultShutdownTimeout,
+				HeaderPreallocation: HeaderPreallocationScan,
+			},
+		},
+		{
+			desc: "header preallocation disabled",
+			cfg: attrs{
+				"address":             ":8080",
+				"headerPreallocation": "disabled",
+			},
+			wantInbound: &wantInbound{
+				Address:             ":8080",
+				ShutdownTimeout:     defaultShutdownTimeout,
+				HeaderPreallocation: HeaderPreallocationDisabled,
+			},
+		},
+		{
+			desc: "header preallocation config overrides option",
+			cfg: attrs{
+				"address":             ":8080",
+				"headerPreallocation": "unfiltered",
+			},
+			opts: []Option{
+				InboundHeaderPreallocation(HeaderPreallocationScan),
+			},
+			wantInbound: &wantInbound{
+				Address:             ":8080",
+				ShutdownTimeout:     defaultShutdownTimeout,
+				HeaderPreallocation: HeaderPreallocationUnfiltered,
+			},
+		},
+		{
+			desc: "header preallocation option without config",
+			cfg:  attrs{"address": ":8080"},
+			opts: []Option{
+				InboundHeaderPreallocation(HeaderPreallocationScan),
+			},
+			wantInbound: &wantInbound{
+				Address:             ":8080",
+				ShutdownTimeout:     defaultShutdownTimeout,
+				HeaderPreallocation: HeaderPreallocationScan,
+			},
+		},
+		{
+			desc: "invalid header preallocation",
+			cfg: attrs{
+				"address":             ":8080",
+				"headerPreallocation": "unknown",
+			},
+			wantErrors: []string{
+				`headerPreallocation must be one of unfiltered, scan, or disabled, got "unknown"`,
+			},
+		},
+		{
+			desc: "empty header preallocation",
+			cfg: attrs{
+				"address":             ":8080",
+				"headerPreallocation": "",
+			},
+			wantErrors: []string{
+				`headerPreallocation must be one of unfiltered, scan, or disabled, got ""`,
+			},
+		},
 	}
 
 	outboundTests := []outboundTest{
@@ -696,6 +768,11 @@ func TestTransportSpec(t *testing.T) {
 				assert.Equal(t, want.ReadTimeout, ib.server.ReadTimeout, "ReadTimeout should match")
 				assert.Equal(t, want.IdleTimeout, ib.server.IdleTimeout, "IdleTimeout should match")
 				assert.Equal(t, want.CanonicalizeHeaderKeys, ib.overrideOriginalItemWithCanonicalizedKey, "canonicalizeHeaderKeys should match")
+				wantHeaderPreallocation := want.HeaderPreallocation
+				if wantHeaderPreallocation == "" {
+					wantHeaderPreallocation = HeaderPreallocationUnfiltered
+				}
+				assert.Equal(t, wantHeaderPreallocation, ib.headerPreallocationStrategy, "headerPreallocation should match")
 				if len(want.HeaderCaseMapping) > 0 {
 					assert.Equal(t, want.HeaderCaseMapping, ib.headerCaseMapping, "headerCaseMapping should match")
 				} else {

transport/http/handler.go

@@ -48,6 +48,86 @@ func popHeader(h http.Header, n string) string {
 	return v
 }
 
+func inboundHeaderCapacity(
+	strategy HeaderPreallocationStrategy,
+	headers http.Header,
+	grabHeaders []headerPreallocationGrabHeader,
+) int {
+	switch strategy {
+	case HeaderPreallocationUnfiltered:
+		return len(headers)
+	case HeaderPreallocationScan:
+		return scannedInboundHeaderCapacity(headers, grabHeaders)
+	case HeaderPreallocationDisabled:
+		return 0
+	default:
+		// Unreachable: inbounds validate the strategy before installing the
+		// handler.
+		return 0
+	}
+}
+
+type headerPreallocationGrabHeader struct {
+	key          string
+	canonicalKey string
+}
+
+func newHeaderPreallocationGrabHeaders(
+	grabHeaders map[string]struct{},
+) []headerPreallocationGrabHeader {
+	headers := make([]headerPreallocationGrabHeader, 0, len(grabHeaders))
+	for key := range grabHeaders {
+		headers = append(headers, headerPreallocationGrabHeader{
+			key:          key,
+			canonicalKey: http.CanonicalHeaderKey(key),
+		})
+	}
+	return headers
+}
+
+func scannedInboundHeaderCapacity(
+	headers http.Header,
+	grabHeaders []headerPreallocationGrabHeader,
+) int {
+	capacity := 0
+	// Count eligible wire forms independently to keep the scan linear.
+	// Raw and prefixed forms may therefore slightly overestimate capacity.
+	for key, values := range headers {
+		if len(values) == 0 {
+			continue
+		}
+
+		if hasPrefixFold(key, ApplicationHeaderPrefix) ||
+			isTracingHeader(key) ||
+			isProxyHeader(key) {
+			capacity++
+		}
+	}
+
+	for _, header := range grabHeaders {
+		if isTracingHeader(header.key) || isProxyHeader(header.key) {
+			continue
+		}
+		if headerValue(headers, header.canonicalKey, header.key) != "" {
+			capacity++
+		}
+	}
+
+	return capacity
+}
+
+func headerValue(headers http.Header, canonicalKey, lowerKey string) string {
+	if values := headers[canonicalKey]; len(values) > 0 && values[0] != "" {
+		return values[0]
+	}
+	if lowerKey != canonicalKey {
+		if values := headers[lowerKey]; len(values) > 0 {
+			return values[0]
+		}
+	}
+	return ""
+}
+
 // handler adapts a transport.Handler into a handler for net/http.
 type handler struct {
 	router                                   transport.Router
@@ -58,6 +138,8 @@ type handler struct {
 	transport                                *Transport
 	overrideOriginalItemWithCanonicalizedKey bool
 	headerCaseMapping                        map[string][]string
+	headerPreallocationStrategy              HeaderPreallocationStrategy
+	headerPreallocationGrabHeaders           []headerPreallocationGrabHeader
 	// duplicate header counter vector
 	duplicateHeaderCounterVec *metrics.CounterVector
 }
@@ -123,7 +205,11 @@ func (h handler) callHandler(responseWriter *responseWriter, req *http.Request,
 	callerProcedure := popHeader(req.Header, CallerProcedureHeader)
 	ttl := popHeader(req.Header, TTLMSHeader)
 
-	transportHeader := transport.NewHeadersWithCapacity(len(req.Header))
+	transportHeader := transport.NewHeadersWithCapacity(inboundHeaderCapacity(
+		h.headerPreallocationStrategy,
+		req.Header,
+		h.headerPreallocationGrabHeaders,
+	))
 	if h.overrideOriginalItemWithCanonicalizedKey {
 		transportHeader = transportHeader.EnableOverrideOriginalItemsWithCanonicalizedKeys()
 	}

transport/http/handler_test.go

@@ -45,6 +45,126 @@ import (
 	"go.uber.org/yarpc/yarpcerrors"
 )
 
+func TestInboundHeaderCapacity(t *testing.T) {
+	t.Run("unfiltered", func(t *testing.T) {
+		headers := makeInboundPreallocationTestHeaders(2, 3)
+		assert.Equal(t, len(headers), inboundHeaderCapacity(
+			HeaderPreallocationUnfiltered,
+			headers,
+			nil,
+		))
+	})
+
+	t.Run("disabled", func(t *testing.T) {
+		assert.Zero(t, inboundHeaderCapacity(
+			HeaderPreallocationDisabled,
+			makeInboundPreallocationTestHeaders(20, 0),
+			nil,
+		))
+	})
+
+	t.Run("scan uses eligible count for small header sets", func(t *testing.T) {
+		headers := makeInboundPreallocationTestHeaders(8, 20)
+		assert.Equal(t, 8, inboundHeaderCapacity(
+			HeaderPreallocationScan,
+			headers,
+			nil,
+		))
+	})
+
+	t.Run("scan supports lowercase HTTP2 headers", func(t *testing.T) {
+		headers := make(http.Header, 10)
+		for i := 0; i < 7; i++ {
+			headers[fmt.Sprintf("rpc-header-application-%d", i)] = []string{"value"}
+		}
+		headers["rpc-header-x-grabbed"] = []string{"prefixed"}
+		headers["x-grabbed"] = []string{"raw"}
+		headers[XForwardedForHeader] = []string{"203.0.113.42"}
+
+		grabHeaders := map[string]struct{}{
+			"x-grabbed":         {},
+			XForwardedForHeader: {},
+		}
+		// Count both raw and prefixed grabbed forms.
+		assert.Equal(t, 10, inboundHeaderCapacity(
+			HeaderPreallocationScan,
+			headers,
+			newHeaderPreallocationGrabHeaders(grabHeaders),
+		))
+	})
+
+	t.Run("scan checks lowercase value after empty canonical value", func(t *testing.T) {
+		headers := http.Header{
+			"X-Grabbed": {""},
+			"x-grabbed": {"raw"},
+		}
+		grabHeaders := map[string]struct{}{
+			"x-grabbed": {},
+		}
+
+		assert.Equal(t, 1, inboundHeaderCapacity(
+			HeaderPreallocationScan,
+			headers,
+			newHeaderPreallocationGrabHeaders(grabHeaders),
+		))
+	})
+
+	t.Run("scan allows small overestimate for duplicate forms", func(t *testing.T) {
+		headers := makeInboundPreallocationTestHeaders(9, 1)
+		headers.Set(UberTraceContextHeaderKey, "trace")
+		headers.Set(UberBaggageHeaderKeyPrefix+"test", "baggage")
+
+		proxyHeaders := map[string]string{
+			XForwardedForHeader:   "203.0.113.42",
+			XForwardedProtoHeader: "https",
+			XForwardedPortHeader:  "443",
+			XRequestIDHeader:      "request-id",
+			XUberSourceHeader:     "source",
+			ViaHeader:             "1.1 proxy",
+			UserAgentHeader:       "yarpc/1.0",
+		}
+		for key, value := range proxyHeaders {
+			headers.Set(key, value)
+		}
+
+		// Both pairs produce one canonical item; the grabbed pair also preserves
+		// two differently-cased original items.
+		headers.Set(ApplicationHeaderPrefix+XForwardedForHeader, "prefixed-proxy")
+		headers.Set("X-Grabbed", "raw-grabbed")
+		headers.Set(ApplicationHeaderPrefix+"X-Grabbed", "prefixed-grabbed")
+		headers.Set("X-Grabbed-Only", "raw-grabbed")
+		headers.Set("X-Empty", "")
+		headers["Rpc-Header-No-Values"] = nil
+
+		grabHeaders := map[string]struct{}{
+			XForwardedForHeader: {},
+			"x-grabbed":         {},
+			"x-grabbed-only":    {},
+			"x-empty":           {},
+		}
+
+		assert.Equal(t, 22, inboundHeaderCapacity(
+			HeaderPreallocationScan,
+			headers,
+			newHeaderPreallocationGrabHeaders(grabHeaders),
+		))
+	})
+}
+
+func makeInboundPreallocationTestHeaders(application, ignored int) http.Header {
+	headers := make(http.Header, application+ignored)
+	for i := 0; i < application; i++ {
+		headers.Set(
+			fmt.Sprintf("%sApplication-%d", ApplicationHeaderPrefix, i),
+			"application-value",
+		)
+	}
+	for i := 0; i < ignored; i++ {
+		headers.Set(fmt.Sprintf("X-Ignored-%d", i), "ignored-value")
+	}
+	return headers
+}
+
 func TestHandlerSuccess(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()