fix test

Signed-off-by: lhy1024 <admin@liudos.us>

pkg/election/leadership.go

@@ -188,6 +188,12 @@ func (ls *Leadership) Watch(serverCtx context.Context, revision int64) {
 		return
 	}
 
+	var watcherCancel context.CancelFunc
+	defer func() {
+		if watcherCancel != nil {
+			watcherCancel()
+		}
+	}()
 	unhealthyTimeout := watchLoopUnhealthyTimeout
 	failpoint.Inject("fastTick", func() {
 		unhealthyTimeout = 5 * time.Second
@@ -198,7 +204,14 @@ func (ls *Leadership) Watch(serverCtx context.Context, revision int64) {
 
 	for {
 		failpoint.Inject("delayWatcher", nil)
+		if watcherCancel != nil {
+			watcherCancel()
+		}
+		// In order to prevent a watch stream being stuck in a partitioned node,
+		// make sure to wrap context with "WithRequireLeader".
 		watcher := clientv3.NewWatcher(ls.client)
+		watcherCtx, cancel := context.WithCancel(clientv3.WithRequireLeader(serverCtx))
+		watcherCancel = cancel
 
 		// When etcd is not available, the watcher.Watch will block,
 		// so we check the etcd availability first.
@@ -222,9 +235,6 @@ func (ls *Leadership) Watch(serverCtx context.Context, revision int64) {
 			}
 		}
 
-		// In order to prevent a watch stream being stuck in a partitioned node,
-		// make sure to wrap context with "WithRequireLeader".
-		watcherCtx := clientv3.WithRequireLeader(serverCtx)
 		watchChan := watcher.Watch(watcherCtx, ls.leaderKey,
 			clientv3.WithRev(revision), clientv3.WithProgressNotify())
 	WatchChanLoop:

pkg/utils/etcdutil/etcdutil.go

@@ -683,15 +683,25 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 {
 }
 
 func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision int64, err error) {
+	var watcherCancel context.CancelFunc
+	defer func() {
+		if watcherCancel != nil {
+			watcherCancel()
+		}
+	}()
 	ticker := time.NewTicker(RequestProgressInterval)
 	defer ticker.Stop()
 	lastReceivedResponseTime := time.Now()
 
 	for {
+		if watcherCancel != nil {
+			watcherCancel()
+		}
 		// In order to prevent a watch stream being stuck in a partitioned node,
 		// make sure to wrap context with "WithRequireLeader".
 		watcher := clientv3.NewWatcher(lw.client)
-		watcherCtx := clientv3.WithRequireLeader(ctx)
+		watcherCtx, cancel := context.WithCancel(clientv3.WithRequireLeader(ctx))
+		watcherCancel = cancel
 		opts := append(lw.opts, clientv3.WithRev(revision))
 		watchChan := watcher.Watch(watcherCtx, lw.key, opts...)
 	WatchChanLoop: