leadership: avoid potential data race (tikv#6636) (tikv#6703)

close tikv#6635, ref tikv#6636 Signed-off-by: ti-chi-bot <[email protected]> Signed-off-by: lhy1024 <[email protected]> Signed-off-by: Ryan Leung <[email protected]> Co-authored-by: lhy1024 <[email protected]> Co-authored-by: Ryan Leung <[email protected]> Co-authored-by: Ti Chi Robot <[email protected]>
ti-chi-bot · Jun 29, 2023 · 1198030 · 1198030
1 parent cb1af02
commit 1198030
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/server/election/leadership.go b/server/election/leadership.go
@@ -16,6 +16,7 @@ package election
 
 import (
 	"context"
+	"sync"
 	"sync/atomic"
 
 	"github.com/pingcap/failpoint"
@@ -54,8 +55,9 @@ type Leadership struct {
 	leaderKey   string
 	leaderValue string
 
-	keepAliveCtx        context.Context
-	keepAliceCancelFunc context.CancelFunc
+	keepAliveCtx            context.Context
+	keepAliveCancelFunc     context.CancelFunc
+	keepAliveCancelFuncLock sync.Mutex
 }
 
 // NewLeadership creates a new Leadership.
@@ -137,8 +139,10 @@ func (ls *Leadership) Keep(ctx context.Context) {
 	if ls == nil {
 		return
 	}
-	ls.keepAliveCtx, ls.keepAliceCancelFunc = context.WithCancel(ctx)
-	ls.getLease().KeepAlive(ls.keepAliveCtx)
+	ls.keepAliveCancelFuncLock.Lock()
+	ls.keepAliveCtx, ls.keepAliveCancelFunc = context.WithCancel(ctx)
+	ls.keepAliveCancelFuncLock.Unlock()
+	go ls.getLease().KeepAlive(ls.keepAliveCtx)
 }
 
 // Check returns whether the leadership is still available.
@@ -230,8 +234,10 @@ func (ls *Leadership) Reset() {
 	if ls == nil || ls.getLease() == nil {
 		return
 	}
-	if ls.keepAliceCancelFunc != nil {
-		ls.keepAliceCancelFunc()
+	ls.keepAliveCancelFuncLock.Lock()
+	if ls.keepAliveCancelFunc != nil {
+		ls.keepAliveCancelFunc()
 	}
+	ls.keepAliveCancelFuncLock.Unlock()
 	ls.getLease().Close()
 }
diff --git a/server/server.go b/server/server.go
@@ -1248,7 +1248,7 @@ func (s *Server) campaignLeader() {
 	})
 
 	// maintain the PD leadership, after this, TSO can be service.
-	go s.member.KeepLeader(ctx)
+	s.member.KeepLeader(ctx)
 	log.Info("campaign pd leader ok", zap.String("campaign-pd-leader-name", s.Name()))
 
 	alllocator, err := s.tsoAllocatorManager.GetAllocator(tso.GlobalDCLocation)