pd: Add pd request retry logic (#1088)

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
 )
@@ -147,6 +150,19 @@ func pdRequest(
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
+	count := 0
+	for {
+		count++
+		if count > pdRequestRetryTime || resp.StatusCode < 500 {
+			break
+		}
+		resp.Body.Close()
+		time.Sleep(time.Second)
+		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)

pkg/pdutil/pd_test.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"testing"
 
@@ -168,3 +169,34 @@ 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) {
+	ctx := context.Background()
+	count := 0
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		count++
+		if count <= 5 {
+			w.WriteHeader(http.StatusGatewayTimeout)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	cli := http.DefaultClient
+	taddr := ts.URL
+	_, reqErr := pdRequest(ctx, taddr, "", cli, http.MethodGet, nil)
+	c.Assert(reqErr, IsNil)
+	ts.Close()
+	count = 0
+	ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		count++
+		if count <= 11 {
+			w.WriteHeader(http.StatusGatewayTimeout)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer ts.Close()
+	taddr = ts.URL
+	_, reqErr = pdRequest(ctx, taddr, "", cli, http.MethodGet, nil)
+	c.Assert(reqErr, NotNil)
+}