optimize: adjust the behavior of PeekAll based on VisitAll (#1403)

valyala · Oct 22, 2022 · 128e9b3 · 128e9b3
1 parent 2c8ce3b
commit 128e9b3
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 7 deletions.
diff --git a/header.go b/header.go
@@ -1811,11 +1811,17 @@ func (h *RequestHeader) peekAll(key []byte) [][]byte {
 	h.mulHeader = h.mulHeader[:0]
 	switch string(key) {
 	case HeaderHost:
-		h.mulHeader = append(h.mulHeader, h.Host())
+		if host := h.Host(); len(host) > 0 {
+			h.mulHeader = append(h.mulHeader, host)
+		}
 	case HeaderContentType:
-		h.mulHeader = append(h.mulHeader, h.ContentType())
+		if contentType := h.ContentType(); len(contentType) > 0 {
+			h.mulHeader = append(h.mulHeader, contentType)
+		}
 	case HeaderUserAgent:
-		h.mulHeader = append(h.mulHeader, h.UserAgent())
+		if ua := h.UserAgent(); len(ua) > 0 {
+			h.mulHeader = append(h.mulHeader, ua)
+		}
 	case HeaderConnection:
 		if h.ConnectionClose() {
 			h.mulHeader = append(h.mulHeader, strClose)
@@ -1827,7 +1833,6 @@ func (h *RequestHeader) peekAll(key []byte) [][]byte {
 	case HeaderCookie:
 		if h.cookiesCollected {
 			h.mulHeader = append(h.mulHeader, appendRequestCookieBytes(nil, h.cookies))
-			return [][]byte{appendRequestCookieBytes(nil, h.cookies)}
 		} else {
 			h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
 		}
@@ -1853,11 +1858,17 @@ func (h *ResponseHeader) peekAll(key []byte) [][]byte {
 	h.mulHeader = h.mulHeader[:0]
 	switch string(key) {
 	case HeaderContentType:
-		h.mulHeader = append(h.mulHeader, h.ContentType())
+		if contentType := h.ContentType(); len(contentType) > 0 {
+			h.mulHeader = append(h.mulHeader, contentType)
+		}
 	case HeaderContentEncoding:
-		h.mulHeader = append(h.mulHeader, h.ContentEncoding())
+		if contentEncoding := h.ContentEncoding(); len(contentEncoding) > 0 {
+			h.mulHeader = append(h.mulHeader, contentEncoding)
+		}
 	case HeaderServer:
-		h.mulHeader = append(h.mulHeader, h.Server())
+		if server := h.Server(); len(server) > 0 {
+			h.mulHeader = append(h.mulHeader, server)
+		}
 	case HeaderConnection:
 		if h.ConnectionClose() {
 			h.mulHeader = append(h.mulHeader, strClose)

diff --git a/header_test.go b/header_test.go
@@ -2883,6 +2883,13 @@ func TestRequestHeader_PeekAll(t *testing.T) {
 	expectRequestHeaderAll(t, h, "Cookie", [][]byte{s2b("foobar=baz")})
 	expectRequestHeaderAll(t, h, HeaderTrailer, [][]byte{s2b("Foo, Bar")})
 	expectRequestHeaderAll(t, h, "aaa", [][]byte{s2b("aaa"), s2b("bbb")})
+
+	h.Del("Content-Type")
+	h.Del(HeaderHost)
+	h.Del("aaa")
+	expectRequestHeaderAll(t, h, "Content-Type", [][]byte{})
+	expectRequestHeaderAll(t, h, HeaderHost, [][]byte{})
+	expectRequestHeaderAll(t, h, "aaa", [][]byte{})
 }
 func expectRequestHeaderAll(t *testing.T, h *RequestHeader, key string, expectedValue [][]byte) {
 	if len(h.PeekAll(key)) != len(expectedValue) {
@@ -2913,6 +2920,11 @@ func TestResponseHeader_PeekAll(t *testing.T) {
 	expectResponseHeaderAll(t, h, HeaderServer, [][]byte{s2b("aaaa")})
 	expectResponseHeaderAll(t, h, HeaderSetCookie, [][]byte{s2b("cccc")})
 	expectResponseHeaderAll(t, h, "aaa", [][]byte{s2b("aaa"), s2b("bbb")})
+
+	h.Del(HeaderContentType)
+	h.Del(HeaderContentEncoding)
+	expectResponseHeaderAll(t, h, HeaderContentType, [][]byte{defaultContentType})
+	expectResponseHeaderAll(t, h, HeaderContentEncoding, [][]byte{})
 }
 
 func expectResponseHeaderAll(t *testing.T, h *ResponseHeader, key string, expectedValue [][]byte) {