leadership: avoid potential data race (tikv#6636)

close tikv#6635

Signed-off-by: lhy1024 <[email protected]>

pkg/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
-	keepAliveCancelFunc context.CancelFunc
+	keepAliveCtx            context.Context
+	keepAliveCancelFunc     context.CancelFunc
+	keepAliveCancelFuncLock sync.Mutex
 }
 
 // NewLeadership creates a new Leadership.
@@ -137,7 +139,9 @@ func (ls *Leadership) Keep(ctx context.Context) {
 	if ls == nil {
 		return
 	}
+	ls.keepAliveCancelFuncLock.Lock()
 	ls.keepAliveCtx, ls.keepAliveCancelFunc = context.WithCancel(ctx)
+	ls.keepAliveCancelFuncLock.Unlock()
 	go ls.getLease().KeepAlive(ls.keepAliveCtx)
 }
 
@@ -230,8 +234,10 @@ func (ls *Leadership) Reset() {
 	if ls == nil || ls.getLease() == nil {
 		return
 	}
+	ls.keepAliveCancelFuncLock.Lock()
 	if ls.keepAliveCancelFunc != nil {
 		ls.keepAliveCancelFunc()
 	}
+	ls.keepAliveCancelFuncLock.Unlock()
 	ls.getLease().Close()
 }