election: fix the keep alive worker (#6925)

server/election/lease.go

@@ -152,8 +152,11 @@
 					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

@@ -101,3 +101,34 @@ func TestLease(t *testing.T) {
 	time.Sleep((defaultLeaseTimeout + 1) * time.Second)
 	re.True(lease1.IsExpired())
 }
+
+func TestLeaseKeepAlive(t *testing.T) {
+	re := require.New(t)
+	cfg := etcdutil.NewTestSingleConfig(t)
+	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())
+}