From e9a0a103e5bd068d32c14bcfaab3b449d98d9259 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Sun, 13 Mar 2016 20:19:47 -0700
Subject: [PATCH] *: refresh the lease TTL correctly when a leader is elected.

The new leader needs to refresh with an extened TTL to gracefully handle
the potential concurrent leader issue. Clients might still send keep alive
to old leader until the old leader itself gives up leadership at most after
an election timeout.
---
 etcdserver/config.go |  4 ++++
 etcdserver/raft.go   |  2 +-
 lease/lessor.go      | 17 +++++++++--------
 lease/lessor_test.go |  4 ++--
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/etcdserver/config.go b/etcdserver/config.go
index 481b908df..0a5384182 100644
--- a/etcdserver/config.go
+++ b/etcdserver/config.go
@@ -134,6 +134,10 @@ func (c *ServerConfig) ReqTimeout() time.Duration {
 	return 5*time.Second + 2*time.Duration(c.ElectionTicks)*time.Duration(c.TickMs)*time.Millisecond
 }
 
+func (c *ServerConfig) electionTimeout() time.Duration {
+	return time.Duration(c.ElectionTicks) * time.Duration(c.TickMs) * time.Millisecond
+}
+
 func (c *ServerConfig) peerDialTimeout() time.Duration {
 	// 1s for queue wait and system delay
 	// + one RTT, which is smaller than 1/5 election timeout
diff --git a/etcdserver/raft.go b/etcdserver/raft.go
index 6d657b3fa..8d8f93686 100644
--- a/etcdserver/raft.go
+++ b/etcdserver/raft.go
@@ -165,7 +165,7 @@ func (r *raftNode) start(s *EtcdServer) {
 						// it promotes or demotes instead of modifying server directly.
 						syncC = r.s.SyncTicker
 						if r.s.lessor != nil {
-							r.s.lessor.Promote()
+							r.s.lessor.Promote(r.s.cfg.electionTimeout())
 						}
 						// TODO: remove the nil checking
 						// current test utility does not provide the stats
diff --git a/lease/lessor.go b/lease/lessor.go
index cc368e03e..268eb8b49 100644
--- a/lease/lessor.go
+++ b/lease/lessor.go
@@ -78,7 +78,8 @@ type Lessor interface {
 
 	// Promote promotes the lessor to be the primary lessor. Primary lessor manages
 	// the expiration and renew of leases.
-	Promote()
+	// Newly promoted lessor renew the TTL of all lease to extend + previous TTL.
+	Promote(extend time.Duration)
 
 	// Demote demotes the lessor from being the primary lessor.
 	Demote()
@@ -188,7 +189,7 @@ func (le *lessor) Grant(id LeaseID, ttl int64) (*Lease, error) {
 	}
 
 	if le.primary {
-		l.refresh()
+		l.refresh(0)
 	} else {
 		l.forever()
 	}
@@ -240,7 +241,7 @@ func (le *lessor) Renew(id LeaseID) (int64, error) {
 		return -1, ErrLeaseNotFound
 	}
 
-	l.refresh()
+	l.refresh(0)
 	return l.TTL, nil
 }
 
@@ -253,7 +254,7 @@ func (le *lessor) Lookup(id LeaseID) *Lease {
 	return nil
 }
 
-func (le *lessor) Promote() {
+func (le *lessor) Promote(extend time.Duration) {
 	le.mu.Lock()
 	defer le.mu.Unlock()
 
@@ -261,7 +262,7 @@ func (le *lessor) Promote() {
 
 	// refresh the expiries of all leases.
 	for _, l := range le.leaseMap {
-		l.refresh()
+		l.refresh(extend)
 	}
 }
 
@@ -452,11 +453,11 @@ func (l Lease) removeFrom(b backend.Backend) {
 
 // refresh refreshes the expiry of the lease. It extends the expiry at least
 // minLeaseTTL second.
-func (l *Lease) refresh() {
+func (l *Lease) refresh(extend time.Duration) {
 	if l.TTL < minLeaseTTL {
 		l.TTL = minLeaseTTL
 	}
-	l.expiry = time.Now().Add(time.Second * time.Duration(l.TTL))
+	l.expiry = time.Now().Add(extend + time.Second*time.Duration(l.TTL))
 }
 
 // forever sets the expiry of lease to be forever.
@@ -491,7 +492,7 @@ func (fl *FakeLessor) Attach(id LeaseID, items []LeaseItem) error { return nil }
 
 func (fl *FakeLessor) Detach(id LeaseID, items []LeaseItem) error { return nil }
 
-func (fl *FakeLessor) Promote() {}
+func (fl *FakeLessor) Promote(extend time.Duration) {}
 
 func (fl *FakeLessor) Demote() {}
 
diff --git a/lease/lessor_test.go b/lease/lessor_test.go
index ef41c6452..c018719d3 100644
--- a/lease/lessor_test.go
+++ b/lease/lessor_test.go
@@ -34,7 +34,7 @@ func TestLessorGrant(t *testing.T) {
 	defer be.Close()
 
 	le := newLessor(be)
-	le.Promote()
+	le.Promote(0)
 
 	l, err := le.Grant(1, 1)
 	if err != nil {
@@ -128,7 +128,7 @@ func TestLessorRenew(t *testing.T) {
 	defer os.RemoveAll(dir)
 
 	le := newLessor(be)
-	le.Promote()
+	le.Promote(0)
 
 	l, err := le.Grant(1, 5)
 	if err != nil {