pd-client: fix a data race in pd client (#660)

pd-client/client.go

@@ -451,15 +451,15 @@ type TSFuture interface {
 }
 
 func (req *tsoRequest) Wait() (int64, int64, error) {
-	defer func() { cmdDuration.WithLabelValues("tso").Observe(time.Since(req.start).Seconds()) }()
 	select {
 	case err := <-req.done:
+		defer tsoReqPool.Put(req)
 		if err != nil {
 			cmdFailedDuration.WithLabelValues("tso").Observe(time.Since(req.start).Seconds())
 			return 0, 0, errors.Trace(err)
 		}
 		physical, logical := req.physical, req.logical
-		tsoReqPool.Put(req)
+		cmdDuration.WithLabelValues("tso").Observe(time.Since(req.start).Seconds())
 		return physical, logical, err
 	case <-req.ctx.Done():
 		return 0, 0, errors.Trace(req.ctx.Err())

pd-client/client_test.go

@@ -17,6 +17,7 @@ import (
 	"net"
 	"os"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -175,6 +176,25 @@ func (s *testClientSuite) TestTSO(c *C) {
 	}
 }
 
+func (s *testClientSuite) TestTSORace(c *C) {
+	var wg sync.WaitGroup
+	begin := make(chan struct{})
+	count := 10
+	wg.Add(count)
+	for i := 0; i < count; i++ {
+		go func() {
+			<-begin
+			for i := 0; i < 100; i++ {
+				_, _, err := s.client.GetTS(context.Background())
+				c.Assert(err, IsNil)
+			}
+			wg.Done()
+		}()
+	}
+	close(begin)
+	wg.Wait()
+}
+
 func (s *testClientSuite) TestGetRegion(c *C) {
 	req := &pdpb.RegionHeartbeatRequest{
 		Header: newHeader(s.srv),