From e0d964188cd27c4a528cf2f224778f374adb162e Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Thu, 19 Apr 2018 18:35:10 -0700
Subject: [PATCH] agent/cache: make edge case with prev/next idx == 0 handled
 better

---
 agent/cache/cache.go |  8 +-------
 agent/cache/entry.go | 30 ++++++++++++++++++++++++++----
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/agent/cache/cache.go b/agent/cache/cache.go
index 1d9b732b83..759d2bc1d3 100644
--- a/agent/cache/cache.go
+++ b/agent/cache/cache.go
@@ -210,13 +210,7 @@ RETRY_GET:
 			// Touch the expiration and fix the heap.
 			c.entriesLock.Lock()
 			entry.Expiry.Reset()
-			idx := entry.Expiry.HeapIndex
-			heap.Fix(c.entriesExpiryHeap, entry.Expiry.HeapIndex)
-			if idx == 0 && entry.Expiry.HeapIndex == 0 {
-				// We didn't move and we were at the head of the heap.
-				// We need to let the loop know that the value changed.
-				c.entriesExpiryHeap.Notify()
-			}
+			c.entriesExpiryHeap.Fix(entry.Expiry)
 			c.entriesLock.Unlock()
 
 			return entry.Value, entry.Error
diff --git a/agent/cache/entry.go b/agent/cache/entry.go
index b86f80ea8c..8174d3f12d 100644
--- a/agent/cache/entry.go
+++ b/agent/cache/entry.go
@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"container/heap"
 	"time"
 )
 
@@ -51,9 +52,34 @@ type expiryHeap struct {
 	// NotifyCh is sent a value whenever the 0 index value of the heap
 	// changes. This can be used to detect when the earliest value
 	// changes.
+	//
+	// There is a single edge case where the heap will not automatically
+	// send a notification: if heap.Fix is called manually and the index
+	// changed is 0 and the change doesn't result in any moves (stays at index
+	// 0), then we won't detect the change. To work around this, please
+	// always call the expiryHeap.Fix method instead.
 	NotifyCh chan struct{}
 }
 
+// Identical to heap.Fix for this heap instance but will properly handle
+// the edge case where idx == 0 and no heap modification is necessary,
+// and still notify the NotifyCh.
+//
+// This is important for cache expiry since the expiry time may have been
+// extended and if we don't send a message to the NotifyCh then we'll never
+// reset the timer and the entry will be evicted early.
+func (h *expiryHeap) Fix(entry *cacheEntryExpiry) {
+	idx := entry.HeapIndex
+	heap.Fix(h, idx)
+
+	// This is the edge case we handle: if the prev and current index
+	// is zero, it means the head-of-line didn't change while the value
+	// changed. Notify to reset our expiry worker.
+	if idx == 0 && entry.HeapIndex == 0 {
+		h.NotifyCh <- struct{}{}
+	}
+}
+
 func (h *expiryHeap) Len() int { return len(h.Entries) }
 
 func (h *expiryHeap) Swap(i, j int) {
@@ -97,7 +123,3 @@ func (h *expiryHeap) Pop() interface{} {
 	h.Entries = old[0 : n-1]
 	return x
 }
-
-func (h *expiryHeap) Notify() {
-	h.NotifyCh <- struct{}{}
-}