Skip to content

Commit

Permalink
fix: fix the bug that temporarily saved the header (#15)
Browse files Browse the repository at this point in the history
* fix: fix the bug that temporarily saved the header

* feat: use sync.Pool to improve
  • Loading branch information
cqqqq777 authored Aug 20, 2023
1 parent 1bb9078 commit c06b823
Showing 1 changed file with 15 additions and 2 deletions.
17 changes: 15 additions & 2 deletions reverse_proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import (
"net/textproto"
"reflect"
"strings"
"sync"
"unsafe"

"github.com/cloudwego/hertz/pkg/app"
Expand Down Expand Up @@ -213,18 +214,24 @@ func (r *ReverseProxy) defaultErrorHandler(c *app.RequestContext, _ error) {
c.Response.Header.SetStatusCode(consts.StatusBadGateway)
}

var respTmpHeaderPool = sync.Pool{
New: func() interface{} {
return make(map[string][]string)
},
}

func (r *ReverseProxy) ServeHTTP(c context.Context, ctx *app.RequestContext) {
req := &ctx.Request
resp := &ctx.Response

// save tmp resp header
respTmpHeader := map[string][]string{}
respTmpHeader := respTmpHeaderPool.Get().(map[string][]string)
if r.saveOriginResHeader {
resp.Header.SetNoDefaultContentType(true)
resp.Header.VisitAll(func(key, value []byte) {
keyStr := string(key)
valueStr := string(value)
if _, ok := respTmpHeader[keyStr]; ok {
if _, ok := respTmpHeader[keyStr]; !ok {
respTmpHeader[keyStr] = []string{valueStr}
} else {
respTmpHeader[keyStr] = append(respTmpHeader[keyStr], valueStr)
Expand Down Expand Up @@ -286,6 +293,12 @@ func (r *ReverseProxy) ServeHTTP(c context.Context, ctx *app.RequestContext) {
}
}

// Clear and put respTmpHeader back to respTmpHeaderPool
for k := range respTmpHeader {
delete(respTmpHeader, k)
}
respTmpHeaderPool.Put(respTmpHeader)

removeResponseConnHeaders(ctx)

for _, h := range hopHeaders {
Expand Down

0 comments on commit c06b823

Please sign in to comment.