etcd: add WithRequireLeader in lease and txn

Signed-off-by: lhy1024 <[email protected]>
lhy1024 · Jan 22, 2025 · ab7c0a5 · ab7c0a5
1 parent ae6df14
commit ab7c0a5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pkg/election/lease.go b/pkg/election/lease.go
@@ -117,7 +117,7 @@ func (l *Lease) KeepAlive(ctx context.Context) {
 	if l == nil {
 		return
 	}
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancel(clientv3.WithRequireLeader(ctx))
 	defer cancel()
 	timeCh := l.keepAliveWorker(ctx, l.leaseTimeout/3)
 	defer log.Info("lease keep alive stopped", zap.String("purpose", l.Purpose))

diff --git a/pkg/storage/kv/etcd_kv.go b/pkg/storage/kv/etcd_kv.go
@@ -210,7 +210,7 @@ type etcdTxn struct {
 func (kv *etcdKVBase) RunInTxn(ctx context.Context, f func(txn Txn) error) error {
 	txn := &etcdTxn{
 		kv:  kv,
-		ctx: ctx,
+		ctx: clientv3.WithRequireLeader(ctx),
 	}
 	err := f(txn)
 	if err != nil {