From ce8db87a8c3f2560750232e85ad5152a0a1c69df Mon Sep 17 00:00:00 2001
From: tiancaiamao <tiancaiamao@gmail.com>
Date: Sat, 10 Jun 2017 15:59:58 +0800
Subject: [PATCH] pd-client: fix a data race in pd client (#660)

---
 pd-client/client.go      |  4 ++--
 pd-client/client_test.go | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/pd-client/client.go b/pd-client/client.go
index 49661d186ab..6baca4f4e72 100644
--- a/pd-client/client.go
+++ b/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())
diff --git a/pd-client/client_test.go b/pd-client/client_test.go
index 1a63c3597c9..879e24d1d4c 100644
--- a/pd-client/client_test.go
+++ b/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),