From 79a4d263b24d96787957d325c6f950cb4a7c410a 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             |  2 ++
 tests/integration/v3_lease_test.go | 15 ++++++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/server/lease/lessor.go b/server/lease/lessor.go
index f4902d275ce7..7ed94c08d6af 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)
@@ -787,6 +788,7 @@ func (le *lessor) initAndRecover() {
 			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 0e649fd803e2..240e6bcbaf8f 100644
--- a/tests/integration/v3_lease_test.go
+++ b/tests/integration/v3_lease_test.go
@@ -236,6 +236,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 		ttl                   time.Duration
 		checkpointingInterval time.Duration
 		leaderChanges         int
+		clusterSize           int
 		expectTTLIsGT         time.Duration
 		expectTTLIsLT         time.Duration
 	}{
@@ -243,6 +244,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,
 		},
 		{
@@ -251,6 +253,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,
 		},
 		{
@@ -260,6 +272,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 			checkpointingEnabled:  true,
 			checkpointingInterval: 10 * time.Second,
 			leaderChanges:         2,
+			clusterSize:           3,
 			expectTTLIsLT:         280 * time.Second,
 		},
 	}
@@ -267,7 +280,7 @@ func TestV3LeaseCheckpoint(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			integration.BeforeTest(t)
 			config := &integration.ClusterConfig{
-				Size:                    3,
+				Size:                    tc.clusterSize,
 				EnableLeaseCheckpoint:   tc.checkpointingEnabled,
 				LeaseCheckpointInterval: tc.checkpointingInterval,
 			}