From eddfb4232f09cbaeaf1627452611f912eb98e335 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Jasionowski?= <jasion@google.com>
Date: Fri, 19 Nov 2021 15:02:25 +0000
Subject: [PATCH] etcdserver,integration: Store remaining TTL on checkpoint

To extend lease checkpointing mechanism to cases when the whole etcd
cluster is restarted.
---
 server/lease/lessor.go             |  8 +++++---
 tests/integration/v3_lease_test.go | 15 ++++++++++++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/server/lease/lessor.go b/server/lease/lessor.go
index 668752974..3249ef5ed 100644
--- a/server/lease/lessor.go
+++ b/server/lease/lessor.go
@@ -351,6 +351,7 @@ func (le *lessor) Checkpoint(id LeaseID, remainingTTL int64) error {
 	if l, ok := le.leaseMap[id]; ok {
 		// when checkpointing, we only update the remainingTTL, Promote is responsible for applying this to lease expiry
 		l.remainingTTL = remainingTTL
+		l.persistTo(le.b)
 		if le.isPrimary() {
 			// schedule the next checkpoint as needed
 			le.scheduleCheckpointIfNeeded(l)
@@ -790,9 +791,10 @@ func (le *lessor) initAndRecover() {
 			ttl: lpb.TTL,
 			// itemSet will be filled in when recover key-value pairs
 			// set expiry to forever, refresh when promoted
-			itemSet: make(map[LeaseItem]struct{}),
-			expiry:  forever,
-			revokec: make(chan struct{}),
+			itemSet:      make(map[LeaseItem]struct{}),
+			expiry:       forever,
+			revokec:      make(chan struct{}),
+			remainingTTL: lpb.RemainingTTL,
 		}
 	}
 	le.leaseExpiredNotifier.Init()
diff --git a/tests/integration/v3_lease_test.go b/tests/integration/v3_lease_test.go
index 6adc2f2b3..9fd0964d5 100644
--- a/tests/integration/v3_lease_test.go
+++ b/tests/integration/v3_lease_test.go
@@ -235,6 +235,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 		ttl                   time.Duration
 		checkpointingInterval time.Duration
 		leaderChanges         int
+		clusterSize           int
 		expectTTLIsGT         time.Duration
 		expectTTLIsLT         time.Duration
 	}{
@@ -242,6 +243,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 			name:          "Checkpointing disabled, lease TTL is reset",
 			ttl:           300 * time.Second,
 			leaderChanges: 1,
+			clusterSize:   3,
 			expectTTLIsGT: 298 * time.Second,
 		},
 		{
@@ -250,6 +252,16 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 			checkpointingEnabled:  true,
 			checkpointingInterval: 10 * time.Second,
 			leaderChanges:         1,
+			clusterSize:           3,
+			expectTTLIsLT:         290 * time.Second,
+		},
+		{
+			name:                  "Checkpointing enabled 10s, lease TTL is preserved after cluster restart",
+			ttl:                   300 * time.Second,
+			checkpointingEnabled:  true,
+			checkpointingInterval: 10 * time.Second,
+			leaderChanges:         1,
+			clusterSize:           1,
 			expectTTLIsLT:         290 * time.Second,
 		},
 		{
@@ -259,6 +271,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 			checkpointingEnabled:  true,
 			checkpointingInterval: 10 * time.Second,
 			leaderChanges:         2,
+			clusterSize:           3,
 			expectTTLIsLT:         280 * time.Second,
 		},
 	}
@@ -266,7 +279,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			BeforeTest(t)
 			config := &ClusterConfig{
-				Size:                    3,
+				Size:                    tc.clusterSize,
 				EnableLeaseCheckpoint:   tc.checkpointingEnabled,
 				LeaseCheckpointInterval: tc.checkpointingInterval,
 			}