client: fix the pd client could be blocked in some cases

client/client.go

@@ -140,6 +140,7 @@ func WithRetry(retry uint64) RegionsOption {
 
 type tsoRequest struct {
 	start      time.Time
+	client     *client
 	ctx        context.Context
 	done       chan error
 	physical   int64
@@ -567,6 +568,7 @@ func (c *client) GetLocalTSAsync(ctx context.Context, dcLocation string) TSFutur
 	}
 	req := tsoReqPool.Get().(*tsoRequest)
 	req.ctx = ctx
+	req.client = c
 	req.start = time.Now()
 	req.dcLocation = dcLocation
 	c.waitForDispatcher()
@@ -630,6 +632,8 @@ func (req *tsoRequest) Wait() (physical int64, logical int64, err error) {
 		return
 	case <-req.ctx.Done():
 		return 0, 0, errors.WithStack(req.ctx.Err())
+	case <-req.client.ctx.Done():
+		return 0, 0, errors.WithStack(req.client.ctx.Err())
 	}
 }
 

client/client_test.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	. "github.com/pingcap/check"
+	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/tikv/pd/pkg/testutil"
 	"go.uber.org/goleak"
@@ -109,3 +110,33 @@ func (s *testClientDialOptionSuite) TestGRPCDialOption(c *C) {
 	c.Assert(err, NotNil)
 	c.Assert(time.Since(start), Greater, 500*time.Millisecond)
 }
+
+var _ = Suite(&testTsoRequestSuite{})
+
+type testTsoRequestSuite struct{}
+
+func (s *testTsoRequestSuite) TestTsoRequestWait(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cli := &client{
+		baseClient: &baseClient{
+			urls:            []string{"localhost:8080"},
+			checkLeaderCh:   make(chan struct{}, 1),
+			ctx:             ctx,
+			cancel:          cancel,
+			security:        SecurityOption{},
+			gRPCDialOptions: []grpc.DialOption{grpc.WithBlock()},
+		},
+		checkTSDeadlineCh: make(chan struct{}),
+	}
+
+	tsoRequest := &tsoRequest{
+		done:     make(chan error, 1),
+		physical: 0,
+		logical:  0,
+		ctx:      context.TODO(),
+		client:   cli,
+	}
+	cancel()
+	_, _, err := tsoRequest.Wait()
+	c.Assert(errors.Cause(err), Equals, context.Canceled)
+}