From 0be684b73ea90cc65f648ecc90fa55f0d3d24a33 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Thu, 29 Jun 2023 19:09:48 +0800
Subject: [PATCH] leadership: avoid potential data race (#6636) (#6701)

close tikv/pd#6635, ref tikv/pd#6636

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>
Signed-off-by: lhy1024 <admin@liudos.us>
Signed-off-by: Ryan Leung <rleungx@gmail.com>

Co-authored-by: lhy1024 <admin@liudos.us>
Co-authored-by: Ryan Leung <rleungx@gmail.com>
Co-authored-by: Ti Chi Robot <ti-community-prow-bot@tidb.io>
---
 server/election/leadership.go | 18 ++++++++++++------
 server/server.go              |  2 +-
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/server/election/leadership.go b/server/election/leadership.go
index d8daf84b694..0489a6c7e4e 100644
--- 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
index bd13193532f..b1ac379edfd 100644
--- a/server/server.go
+++ b/server/server.go
@@ -1348,7 +1348,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)