pd: Add pd request retry logic

pkg/pdutil/pd.go

@@ -40,6 +40,9 @@ const (
 	scheduleConfigPrefix = "pd/api/v1/config/schedule"
 	pauseTimeout         = 5 * time.Minute
 
+	// pd request retry time when connection fail
+	pdRequestRetryTime = 10
+
 	// set max-pending-peer-count to a large value to avoid scatter region failed.
 	maxPendingPeerUnlimited uint64 = math.MaxInt32
 )
@@ -143,21 +146,30 @@ func pdRequest(
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	resp, err := cli.Do(req)
-	if err != nil {
-		return nil, errors.Trace(err)
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != http.StatusOK {
-		res, _ := ioutil.ReadAll(resp.Body)
-		return nil, errors.Annotatef(berrors.ErrPDInvalidResponse, "[%d] %s %s", resp.StatusCode, res, reqURL)
-	}
+	var (
+		resp   *http.Response
+		reqErr error
+	)
 
-	r, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return nil, errors.Trace(err)
+	for i := 0; i < pdRequestRetryTime; i++ {
+		resp, reqErr = cli.Do(req)
+		if reqErr != nil {
+			log.Warn("pd request fail, retry", zap.Int("retry time", i), zap.Error(reqErr))
+			time.Sleep(time.Duration(1) * time.Second)
+			continue
+		}
+		if resp.StatusCode != http.StatusOK {
+			res, _ := ioutil.ReadAll(resp.Body)
+			return nil, errors.Annotatef(berrors.ErrPDInvalidResponse, "[%d] %s %s", resp.StatusCode, res, reqURL)
+		}
+		r, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		resp.Body.Close()
+		return r, nil
 	}
-	return r, nil
+	return nil, errors.Trace(reqErr)
 }
 
 // PdController manage get/update config from pd.

pkg/pdutil/pd_test.go

@@ -12,13 +12,16 @@ import (
 	"net/http"
 	"net/url"
 	"testing"
+	"time"
 
 	"github.com/coreos/go-semver/semver"
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/util/codec"
 	"github.com/tikv/pd/server/core"
 	"github.com/tikv/pd/server/statistics"
+	"go.uber.org/zap"
 )
 
 func TestT(t *testing.T) {
@@ -168,3 +171,36 @@ func (s *testPDControllerSuite) TestPDVersion(c *C) {
 	c.Assert(r.Minor, Equals, expectV.Minor)
 	c.Assert(r.PreRelease, Equals, expectV.PreRelease)
 }
+
+func (s *testPDControllerSuite) TestPDRequestRetry(c *C) {
+	cnt := 0
+	// mock cli.Do(req) function
+	mockDo := func(req *http.Request) (retres *http.Response, reterr error) {
+		cnt++
+		if cnt <= 5 {
+			return nil, errors.New("request fail")
+		}
+		return &http.Response{}, nil
+	}
+	ctx := context.Background()
+	addr := "https://mock"
+	u, err := url.Parse(addr)
+	c.Assert(err, IsNil)
+	reqURL := fmt.Sprintf("%s/%s", u, clusterVersionPrefix)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
+	c.Assert(err, IsNil)
+	var (
+		resp     *http.Response
+		reqError error
+	)
+	for i := 0; i < pdRequestRetryTime; i++ {
+		resp, reqError = mockDo(req)
+		if reqError != nil {
+			log.Warn("pd request fail, retry", zap.Int("retry time", i), zap.Error(reqError))
+			time.Sleep(time.Duration(1) * time.Second)
+			continue
+		}
+	}
+	c.Assert(resp, NotNil)
+	c.Assert(reqError, IsNil)
+}