diff --git a/pkg/pdutil/pd.go b/pkg/pdutil/pd.go index 3ebf2921d..017d11e86 100644 --- a/pkg/pdutil/pd.go +++ b/pkg/pdutil/pd.go @@ -38,6 +38,11 @@ const ( maxMsgSize = int(128 * utils.MB) // pd.ScanRegion may return a large response 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 ) type pauseConfigExpectation uint8 @@ -126,6 +131,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) diff --git a/pkg/pdutil/pd_test.go b/pkg/pdutil/pd_test.go index c2a84932a..707dfeefc 100644 --- a/pkg/pdutil/pd_test.go +++ b/pkg/pdutil/pd_test.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "net/http" + "net/http/httptest" "net/url" "testing" @@ -167,3 +168,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) +}