From a6984c53deb5c9beb262dcf260a09e01d59339b7 Mon Sep 17 00:00:00 2001
From: Gyuho Lee <gyuhox@gmail.com>
Date: Mon, 2 Apr 2018 16:27:16 -0700
Subject: [PATCH] lease: add "expireExists"

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>
---
 lease/lease_queue_test.go | 11 ++++++++++
 lease/lessor.go           | 44 +++++++++++++++++++++++++++++----------
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/lease/lease_queue_test.go b/lease/lease_queue_test.go
index 575fbc169..72346a1b8 100644
--- a/lease/lease_queue_test.go
+++ b/lease/lease_queue_test.go
@@ -41,4 +41,15 @@ func TestLeaseQueue(t *testing.T) {
 	if le.leaseHeap[0].id != LeaseID(1) {
 		t.Fatalf("first item expected lease ID %d, got %d", LeaseID(1), le.leaseHeap[0].id)
 	}
+
+	l, ok, more := le.expireExists()
+	if l.ID != 1 {
+		t.Fatalf("first item expected lease ID %d, got %d", 1, l.ID)
+	}
+	if !ok {
+		t.Fatal("expect expiry lease exists")
+	}
+	if more {
+		t.Fatal("expect no more expiry lease")
+	}
 }
diff --git a/lease/lessor.go b/lease/lessor.go
index 85d2f90b6..21569394c 100644
--- a/lease/lessor.go
+++ b/lease/lessor.go
@@ -521,28 +521,50 @@ func (le *lessor) runLoop() {
 	}
 }
 
+// expireExists returns true if expiry items exist.
+// It pops "first" expired item from heap.
+// If it's not expired yet, add it back.
+// "next" is true, to indicate that it may exist in next attempt.
+func (le *lessor) expireExists() (l *Lease, ok bool, next bool) {
+	if le.leaseHeap.Len() == 0 {
+		return nil, false, false
+	}
+
+	item := heap.Pop(&le.leaseHeap).(*LeaseWithTime) // O(log N)
+	l = le.leaseMap[item.id]
+	if l == nil {
+		// lease has expired or been revoked
+		// no need to revoke (nothing is expiry)
+		return nil, false, true
+	}
+
+	if time.Now().UnixNano() < item.expiration {
+		// Candidate expirations are caught up, reinsert this item
+		// and no need to revoke (nothing is expiry)
+		heap.Push(&le.leaseHeap, item) // O(log N)
+		return l, false, false
+	}
+	// if the lease is actually expired, add to the removal list. If it is not expired, we can ignore it because another entry will have been inserted into the heap
+
+	return l, true, false
+}
+
 // findExpiredLeases loops leases in the leaseMap until reaching expired limit
 // and returns the expired leases that needed to be revoked.
 func (le *lessor) findExpiredLeases(limit int) []*Lease {
 	leases := make([]*Lease, 0, 16)
 
 	for {
-		if le.leaseHeap.Len() == 0 {
+		l, ok, next := le.expireExists()
+		if !ok && !next {
 			break
 		}
-
-		item := heap.Pop(&le.leaseHeap).(*LeaseWithTime)
-		l := le.leaseMap[item.id]
-		if l == nil {
-			// lease has expired or been revoked, continue
+		if !ok {
 			continue
 		}
-		if time.Now().UnixNano() < item.expiration {
-			// Candidate expirations are caught up, reinsert this item
-			heap.Push(&le.leaseHeap, item)
-			break
+		if next {
+			continue
 		}
-		// if the lease is actually expired, add to the removal list. If it is not expired, we can ignore it because another entry will have been inserted into the heap
 
 		if l.expired() {
 			leases = append(leases, l)