election: fix the keep alive worker (#6925) (#6939)

close #6926

Signed-off-by: Ryan Leung <[email protected]>

Co-authored-by: Ryan Leung <[email protected]>

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))
 				}
 			}()
 

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())
+}