Fix request header being pinned after pd profiling (#1069)

pkg/apiserver/profiling/fetcher.go

@@ -4,7 +4,6 @@ package profiling
 
 import (
 	"fmt"
-	"net/http"
 	"time"
 
 	"go.uber.org/fx"
@@ -92,8 +91,9 @@ type pdFetcher struct {
 
 func (f *pdFetcher) fetch(op *fetchOptions) ([]byte, error) {
 	baseURL := fmt.Sprintf("%s://%s:%d", f.statusAPIHTTPScheme, op.ip, op.port)
-	f.client.WithBeforeRequest(func(req *http.Request) {
-		req.Header.Add("PD-Allow-follower-handle", "true")
-	})
-	return f.client.WithTimeout(maxProfilingTimeout).WithBaseURL(baseURL).SendGetRequest(op.path)
+	return f.client.
+		WithTimeout(maxProfilingTimeout).
+		WithBaseURL(baseURL).
+		AddRequestHeader("PD-Allow-follower-handle", "true").
+		SendGetRequest(op.path)
 }

pkg/httpc/client.go

@@ -25,7 +25,8 @@ const (
 
 type Client struct {
 	http.Client
-	BeforeRequest func(req *http.Request)
+
+	header http.Header
 }
 
 func NewHTTPClient(lc fx.Lifecycle, config *config.Config) *Client {
@@ -52,14 +53,27 @@ func NewHTTPClient(lc fx.Lifecycle, config *config.Config) *Client {
 	}
 }
 
+// Clone is a temporary solution to the unexpected shared pointer field and race problem
+// TODO: use latest `/util/client` for better api experience.
+func (c *Client) Clone() *Client {
+	return &Client{
+		Client: c.Client,
+		header: c.header.Clone(),
+	}
+}
+
 func (c Client) WithTimeout(timeout time.Duration) *Client {
 	c.Timeout = timeout
 	return &c
 }
 
-func (c Client) WithBeforeRequest(callback func(req *http.Request)) *Client {
-	c.BeforeRequest = callback
-	return &c
+func (c *Client) CloneAndAddRequestHeader(key, value string) *Client {
+	cc := c.Clone()
+	if cc.header == nil {
+		cc.header = http.Header{}
+	}
+	cc.header.Add(key, value)
+	return cc
 }
 
 // TODO: Replace using go-resty.
@@ -90,10 +104,7 @@ func (c *Client) Send(
 		log.Warn("SendRequest failed", zap.String("uri", uri), zap.Error(err))
 		return nil, e
 	}
-
-	if c.BeforeRequest != nil {
-		c.BeforeRequest(req)
-	}
+	req.Header = c.header
 
 	resp, err := c.Do(req)
 	if err != nil {

pkg/httpc/client_test.go

@@ -0,0 +1,67 @@
+// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0.
+
+package httpc
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/fx/fxtest"
+
+	"github.com/pingcap/tidb-dashboard/pkg/config"
+)
+
+func newTestClient(t *testing.T) *Client {
+	lc := fxtest.NewLifecycle(t)
+	config := &config.Config{}
+	return NewHTTPClient(lc, config)
+}
+
+func Test_Clone(t *testing.T) {
+	c := newTestClient(t)
+	cc := c.Clone()
+
+	require.NotSame(t, c, cc)
+
+	require.Nil(t, c.header)
+	require.Nil(t, cc.header)
+	require.NotSame(t, c.header, cc.header)
+}
+
+func Test_CloneAndAddRequestHeader(t *testing.T) {
+	c := newTestClient(t)
+	cc := c.CloneAndAddRequestHeader("1", "11")
+
+	require.Nil(t, c.header)
+	require.Equal(t, "11", cc.header.Get("1"))
+
+	cc2 := cc.CloneAndAddRequestHeader("2", "22")
+	require.Equal(t, "11", cc.header.Get("1"))
+	require.Equal(t, "", cc.header.Get("2"))
+	require.Equal(t, "11", cc2.header.Get("1"))
+	require.Equal(t, "22", cc2.header.Get("2"))
+}
+
+func Test_Send_withHeader(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte(r.Header.Get("1")))
+	}))
+	defer ts.Close()
+
+	c := newTestClient(t)
+	resp1, _ := c.Send(context.Background(), ts.URL, http.MethodGet, nil, nil, "")
+	d1, _ := resp1.Body()
+	require.Equal(t, "", string(d1))
+
+	cc := c.CloneAndAddRequestHeader("1", "11")
+	resp2, _ := cc.Send(context.Background(), ts.URL, http.MethodGet, nil, nil, "")
+	d2, _ := resp2.Body()
+	require.Equal(t, "11", string(d2))
+
+	resp3, _ := c.Send(context.Background(), ts.URL, http.MethodGet, nil, nil, "")
+	d3, _ := resp3.Body()
+	require.Equal(t, "", string(d3))
+}

pkg/pd/client.go

@@ -64,8 +64,8 @@ func (c Client) WithTimeout(timeout time.Duration) *Client {
 	return &c
 }
 
-func (c Client) WithBeforeRequest(callback func(req *http.Request)) *Client {
-	c.httpClient.BeforeRequest = callback
+func (c Client) AddRequestHeader(key, value string) *Client {
+	c.httpClient = c.httpClient.CloneAndAddRequestHeader(key, value)
 	return &c
 }
 

pkg/pd/client_test.go

@@ -0,0 +1,52 @@
+// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0.
+
+package pd
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/fx/fxtest"
+
+	"github.com/pingcap/tidb-dashboard/pkg/config"
+	"github.com/pingcap/tidb-dashboard/pkg/httpc"
+)
+
+func newTestClient(t *testing.T) *Client {
+	lc := fxtest.NewLifecycle(t)
+	config := &config.Config{}
+	c := NewPDClient(lc, httpc.NewHTTPClient(lc, config), config)
+	c.lifecycleCtx = context.Background()
+	return c
+}
+
+func Test_AddRequestHeader_returnDifferentHTTPClient(t *testing.T) {
+	c := newTestClient(t)
+	cc := c.AddRequestHeader("1", "11")
+
+	require.NotSame(t, c.httpClient, cc.httpClient)
+}
+
+func Test_Get_withHeader(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte(r.Header.Get("1")))
+	}))
+	defer ts.Close()
+
+	c := newTestClient(t).WithBaseURL(ts.URL)
+	resp1, _ := c.Get("")
+	d1, _ := resp1.Body()
+	require.Equal(t, "", string(d1))
+
+	cc := c.AddRequestHeader("1", "11")
+	resp2, _ := cc.Get("")
+	d2, _ := resp2.Body()
+	require.Equal(t, "11", string(d2))
+
+	resp3, _ := c.Get("")
+	d3, _ := resp3.Body()
+	require.Equal(t, "", string(d3))
+}