From 440b91db74ea864d0827f5df09285e8899e10099 Mon Sep 17 00:00:00 2001
From: Haizhi Geng <ghzpotato@gmail.com>
Date: Mon, 20 Jul 2020 14:11:35 +0800
Subject: [PATCH] cherry pick #2665 to release-4.0

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 server/tso/tso.go | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/server/tso/tso.go b/server/tso/tso.go
index aa34ab75c44..3c39e44adee 100644
--- a/server/tso/tso.go
+++ b/server/tso/tso.go
@@ -15,6 +15,7 @@ package tso
 
 import (
 	"path"
+	"sync"
 	"sync/atomic"
 	"time"
 	"unsafe"
@@ -45,6 +46,7 @@ type TimestampOracle struct {
 	ts            unsafe.Pointer
 	lastSavedTime atomic.Value
 
+	mu    sync.RWMutex
 	lease *member.LeaderLease
 
 	rootPath      string
@@ -86,6 +88,12 @@ func (t *TimestampOracle) loadTimestamp() (time.Time, error) {
 	return typeutil.ParseTimestamp(data)
 }
 
+func (t *TimestampOracle) checkLease() bool {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.lease != nil && !t.lease.IsExpired()
+}
+
 // save timestamp, if lastTs is 0, we think the timestamp doesn't exist, so create it,
 // otherwise, update it.
 func (t *TimestampOracle) saveTimestamp(ts time.Time) error {
@@ -140,7 +148,9 @@ func (t *TimestampOracle) SyncTimestamp(lease *member.LeaderLease) error {
 	current := &atomicObject{
 		physical: next,
 	}
+	t.mu.Lock()
 	t.lease = lease
+	t.mu.Unlock()
 	atomic.StorePointer(&t.ts, unsafe.Pointer(current))
 
 	return nil
@@ -148,7 +158,7 @@ func (t *TimestampOracle) SyncTimestamp(lease *member.LeaderLease) error {
 
 // ResetUserTimestamp update the physical part with specified tso.
 func (t *TimestampOracle) ResetUserTimestamp(tso uint64) error {
-	if t.lease == nil || t.lease.IsExpired() {
+	if !t.checkLease() {
 		tsoCounter.WithLabelValues("err_lease_reset_ts").Inc()
 		return errors.New("Setup timestamp failed, lease expired")
 	}
@@ -273,6 +283,11 @@ func (t *TimestampOracle) GetRespTS(count uint32) (pdpb.Timestamp, error) {
 	for i := 0; i < maxRetryCount; i++ {
 		current := (*atomicObject)(atomic.LoadPointer(&t.ts))
 		if current == nil || current.physical == typeutil.ZeroTime {
+			// If it's leader, maybe SyncTimestamp hasn't completed yet
+			if t.checkLease() {
+				time.Sleep(200 * time.Millisecond)
+				continue
+			}
 			return pdpb.Timestamp{}, errors.New("can not get timestamp, may be not leader")
 		}
 
@@ -287,7 +302,7 @@ func (t *TimestampOracle) GetRespTS(count uint32) (pdpb.Timestamp, error) {
 			continue
 		}
 		// In case lease expired after the first check.
-		if t.lease == nil || t.lease.IsExpired() {
+		if !t.checkLease() {
 			return pdpb.Timestamp{}, errors.New("alloc timestamp failed, lease expired")
 		}
 		return resp, nil