From e7c33abbf683b8bceb4d80dd9e57e54cfe718e5e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rene=CC=81=20Kroon?= <rkroon@bol.com>
Date: Thu, 9 Jul 2020 21:57:54 +0200
Subject: [PATCH] Fix Issue #28 call expiration on Close

---
 .travis.yml   |  4 +--
 cache.go      | 76 +++++++++++++++++++++++++++++++++------------------
 cache_test.go | 59 +++++++++++++++++++++++++++++++++++++++
 go.mod        |  2 +-
 4 files changed, 112 insertions(+), 29 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 70b965c..7862c32 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,8 +1,8 @@
 language: go
 
 go:
-  - 1.13
-  - 1.12
+  - "1.14"
+  - "1.13"
 git:
   depth: 1
 
diff --git a/cache.go b/cache.go
index 0212fc3..f772d0c 100644
--- a/cache.go
+++ b/cache.go
@@ -80,6 +80,11 @@ func (cache *Cache) startExpirationProcessing() {
 		select {
 		case shutdownFeedback := <-cache.shutdownSignal:
 			timer.Stop()
+			cache.mutex.Lock()
+			if cache.priorityQueue.Len() > 0 {
+				cache.evictjob()
+			}
+			cache.mutex.Unlock()
 			shutdownFeedback <- struct{}{}
 			return
 		case <-timer.C:
@@ -90,32 +95,7 @@ func (cache *Cache) startExpirationProcessing() {
 				continue
 			}
 
-			// index will only be advanced if the current entry will not be evicted
-			i := 0
-			for item := cache.priorityQueue.items[i]; item.expired(); item = cache.priorityQueue.items[i] {
-
-				if cache.checkExpireCallback != nil {
-					if !cache.checkExpireCallback(item.key, item.data) {
-						item.touch()
-						cache.priorityQueue.update(item)
-						i++
-						if i == cache.priorityQueue.Len() {
-							break
-						}
-						continue
-					}
-				}
-
-				cache.priorityQueue.remove(item)
-				delete(cache.items, item.key)
-				if cache.expireCallback != nil {
-					go cache.expireCallback(item.key, item.data)
-				}
-				if cache.priorityQueue.Len() == 0 {
-					goto done
-				}
-			}
-		done:
+			cache.cleanjob()
 			cache.mutex.Unlock()
 
 		case <-cache.expirationNotification:
@@ -125,6 +105,50 @@ func (cache *Cache) startExpirationProcessing() {
 	}
 }
 
+func (cache *Cache) evictjob() {
+	// index will only be advanced if the current entry will not be evicted
+	i := 0
+	for item := cache.priorityQueue.items[i]; ; item = cache.priorityQueue.items[i] {
+
+		cache.priorityQueue.remove(item)
+		delete(cache.items, item.key)
+		if cache.expireCallback != nil {
+			go cache.expireCallback(item.key, item.data)
+		}
+		if cache.priorityQueue.Len() == 0 {
+			return
+		}
+	}
+}
+
+func (cache *Cache) cleanjob() {
+	// index will only be advanced if the current entry will not be evicted
+	i := 0
+	for item := cache.priorityQueue.items[i]; item.expired(); item = cache.priorityQueue.items[i] {
+
+		if cache.checkExpireCallback != nil {
+			if !cache.checkExpireCallback(item.key, item.data) {
+				item.touch()
+				cache.priorityQueue.update(item)
+				i++
+				if i == cache.priorityQueue.Len() {
+					break
+				}
+				continue
+			}
+		}
+
+		cache.priorityQueue.remove(item)
+		delete(cache.items, item.key)
+		if cache.expireCallback != nil {
+			go cache.expireCallback(item.key, item.data)
+		}
+		if cache.priorityQueue.Len() == 0 {
+			return
+		}
+	}
+}
+
 // Close calls Purge, and then stops the goroutine that does ttl checking, for a clean shutdown.
 // The cache is no longer cleaning up after the first call to Close, repeated calls are safe though.
 func (cache *Cache) Close() {
diff --git a/cache_test.go b/cache_test.go
index 2dc4852..3bc3ec2 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -17,6 +17,65 @@ func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(m)
 }
 
+// Issue #28: call expirationCallback automatically on cache.Close()
+func TestCache_ExpirationOnClose(t *testing.T) {
+
+	cache := NewCache()
+
+	success := make(chan struct{})
+	defer close(success)
+
+	cache.SetTTL(time.Hour * 100)
+	cache.SetExpirationCallback(func(key string, value interface{}) {
+		t.Logf("%s\t%v", key, value)
+		success <- struct{}{}
+	})
+	cache.Set("1", 1)
+	cache.Set("2", 1)
+	cache.Set("3", 1)
+
+	found := 0
+	cache.Close()
+	wait := time.NewTimer(time.Millisecond * 100)
+	for found != 3 {
+		select {
+		case <-success:
+			found++
+		case <-wait.C:
+			t.Fail()
+		}
+	}
+
+}
+
+// # Issue 29: After Close() the behaviour of Get, Set, Remove is not defined.
+/*
+func TestCache_ModifyAfterClose(t *testing.T) {
+	cache := NewCache()
+
+	cache.SetTTL(time.Hour * 100)
+	cache.SetExpirationCallback(func(key string, value interface{}) {
+		t.Logf("%s\t%v", key, value)
+	})
+	cache.Set("1", 1)
+	cache.Set("2", 1)
+	cache.Set("3", 1)
+
+	cache.Close()
+
+	cache.Get("broken3")
+	cache.Set("broken", 1)
+	cache.Remove("broken2")
+
+	wait := time.NewTimer(time.Millisecond * 100)
+
+	select {
+	case <-wait.C:
+		t.Fail()
+	}
+
+}*/
+
 // Issue #23: Goroutine leak on closing. When adding a close method i would like to see
 // that it can be called in a repeated way without problems.
 func TestCache_MultipleCloseCalls(t *testing.T) {
diff --git a/go.mod b/go.mod
index 51ad6c6..6806b28 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module github.com/ReneKroon/ttlcache
 
-go 1.12
+go 1.14
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect