From d76f55de05d608f099e915bfde5d9beabf762fce Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 11 Sep 2024 13:17:21 +0800
Subject: [PATCH] election: fix the keep alive worker (#6925) (#6939)

close tikv/pd#6926

Signed-off-by: Ryan Leung <rleungx@gmail.com>

Co-authored-by: Ryan Leung <rleungx@gmail.com>
---
 server/election/lease.go      |  5 ++++-
 server/election/lease_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/server/election/lease.go b/server/election/lease.go
index 4e368b98361..f3c00f47089 100644
--- a/server/election/lease.go
+++ b/server/election/lease.go
@@ -152,8 +152,11 @@ func (l *lease) keepAliveWorker(ctx context.Context, interval time.Duration) <-c
 					expire := start.Add(time.Duration(res.TTL) * time.Second)
 					select {
 					case ch <- expire:
-					case <-ctx1.Done():
+					// Here we don't use `ctx1.Done()` because we want to make sure if the keep alive success, we can update the expire time.
+					case <-ctx.Done():
 					}
+				} else {
+					log.Error("keep alive response ttl is zero", zap.String("purpose", l.Purpose))
 				}
 			}()
 
diff --git a/server/election/lease_test.go b/server/election/lease_test.go
index 0c0aa3c1687..0161e7f3d27 100644
--- a/server/election/lease_test.go
+++ b/server/election/lease_test.go
@@ -16,9 +16,11 @@ package election
 
 import (
 	"context"
+	"testing"
 	"time"
 
 	. "github.com/pingcap/check"
+	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/etcdutil"
 	"go.etcd.io/etcd/clientv3"
 	"go.etcd.io/etcd/embed"
@@ -104,3 +106,34 @@ func (s *testLeaseSuite) TestLease(c *C) {
 	time.Sleep((defaultLeaseTimeout + 1) * time.Second)
 	c.Check(lease1.IsExpired(), IsTrue)
 }
+
+func TestLeaseKeepAlive(t *testing.T) {
+	re := require.New(t)
+	cfg := etcdutil.NewTestSingleConfig()
+	etcd, err := embed.StartEtcd(cfg)
+	defer func() {
+		etcd.Close()
+	}()
+	re.NoError(err)
+
+	ep := cfg.LCUrls[0].String()
+	client, err := clientv3.New(clientv3.Config{
+		Endpoints: []string{ep},
+	})
+	re.NoError(err)
+
+	<-etcd.Server.ReadyNotify()
+
+	// Create the lease.
+	lease := &lease{
+		Purpose: "test_lease",
+		client:  client,
+		lease:   clientv3.NewLease(client),
+	}
+
+	re.NoError(lease.Grant(defaultLeaseTimeout))
+	ch := lease.keepAliveWorker(context.Background(), 2*time.Second)
+	time.Sleep(2 * time.Second)
+	<-ch
+	re.NoError(lease.Close())
+}