leadership: avoid potential data race (#6636) (#6702)

ti-chi-bot · lhy1024 · web-flow · commit cfdcdff89441 · 2023-06-29T11:05:18.000Z
close #6635, ref #6636 Signed-off-by: lhy1024 <admin@liudos.us> Co-authored-by: lhy1024 <admin@liudos.us>
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
-	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()
 }

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ package election`
`16`	`16`
`17`	`17`	`import (`
`18`	`18`	`"context"`
	`19`	`+ "sync"`
`19`	`20`	`"sync/atomic"`
`20`	`21`
`21`	`22`	`"github.com/pingcap/failpoint"`
`@@ -54,8 +55,9 @@ type Leadership struct {`
`54`	`55`	`leaderKey string`
`55`	`56`	`leaderValue string`
`56`	`57`
`57`		`- keepAliveCtx context.Context`
`58`		`- keepAliveCancelFunc context.CancelFunc`
	`58`	`+ keepAliveCtx context.Context`
	`59`	`+ keepAliveCancelFunc context.CancelFunc`
	`60`	`+ keepAliveCancelFuncLock sync.Mutex`
`59`	`61`	`}`
`60`	`62`
`61`	`63`	`// NewLeadership creates a new Leadership.`
`@@ -137,7 +139,9 @@ func (ls *Leadership) Keep(ctx context.Context) {`
`137`	`139`	`if ls == nil {`
`138`	`140`	`return`
`139`	`141`	`}`
	`142`	`+ ls.keepAliveCancelFuncLock.Lock()`
`140`	`143`	`ls.keepAliveCtx, ls.keepAliveCancelFunc = context.WithCancel(ctx)`
	`144`	`+ ls.keepAliveCancelFuncLock.Unlock()`
`141`	`145`	`go ls.getLease().KeepAlive(ls.keepAliveCtx)`
`142`	`146`	`}`
`143`	`147`
`@@ -230,8 +234,10 @@ func (ls *Leadership) Reset() {`
`230`	`234`	`if ls == nil \|\| ls.getLease() == nil {`
`231`	`235`	`return`
`232`	`236`	`}`
	`237`	`+ ls.keepAliveCancelFuncLock.Lock()`
`233`	`238`	`if ls.keepAliveCancelFunc != nil {`
`234`	`239`	`ls.keepAliveCancelFunc()`
`235`	`240`	`}`
	`241`	`+ ls.keepAliveCancelFuncLock.Unlock()`
`236`	`242`	`ls.getLease().Close()`
`237`	`243`	`}`