From 56801469f37a1f2f4a34f9fd4f8c6d3b4a610095 Mon Sep 17 00:00:00 2001
From: Brooke Bryan <brooke@bajb.net>
Date: Thu, 25 Feb 2021 11:25:00 +0000
Subject: [PATCH] Support cleanup

---
 cache.go      | 17 +++++++++++++----
 cache_test.go | 29 +++++++++++++++++++++++++++++
 go.mod        |  2 +-
 item.go       |  4 +++-
 4 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/cache.go b/cache.go
index acc2ce2..745a0ad 100644
--- a/cache.go
+++ b/cache.go
@@ -57,23 +57,28 @@ func (m CacheMap) MSet(data map[string]interface{}, duration time.Duration) {
 	for key, value := range data {
 		shard := m.GetShard(key)
 		shard.Lock()
-		shard.items[key] = newItem(value, duration, time.Now().Add(m.options.maxLifetime))
+		shard.items[key] = newItem(value, duration, time.Now().Add(m.options.maxLifetime), nil)
 		shard.Unlock()
 	}
 }
 
-// Sets the given value under the specified key
-func (m CacheMap) Set(key string, value interface{}, duration *time.Duration) {
+func (m CacheMap) SetWithCleanup(key string, value interface{}, duration *time.Duration, cleanup func(*Item)) {
 	// Get map shard.
 	shard := m.GetShard(key)
 	shard.Lock()
 	if duration == nil {
 		duration = &m.options.defaultCacheDuration
 	}
-	shard.items[key] = newItem(value, *duration, time.Now().Add(m.options.maxLifetime))
+	itm := newItem(value, *duration, time.Now().Add(m.options.maxLifetime), cleanup)
+	shard.items[key] = itm
 	shard.Unlock()
 }
 
+// Sets the given value under the specified key
+func (m CacheMap) Set(key string, value interface{}, duration *time.Duration) {
+	m.SetWithCleanup(key, value, duration, nil)
+}
+
 // Retrieves an item from the map with the given key, and optionally increase its expiry time if found
 func (m CacheMap) TouchGet(key string, touch bool) (interface{}, bool) {
 	shard := m.GetShard(key)
@@ -133,6 +138,10 @@ func (ms CacheMapShared) Remove(key string) {
 
 // Removes an element from the map
 func (ms CacheMapShared) remove(key string) {
+	if itm, ok := ms.items[key]; ok && itm.onDelete != nil {
+		itm.onDelete(itm)
+	}
+
 	delete(ms.items, key)
 }
 
diff --git a/cache_test.go b/cache_test.go
index 4a73260..6ee5ba1 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -96,3 +96,32 @@ func TestItems(t *testing.T) {
 		t.Errorf("Expected cache to return 3 items after cache expiry")
 	}
 }
+
+func TestCleanup(t *testing.T) {
+	dur := time.Millisecond * 100
+
+	cleanup := 0
+
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Millisecond * 5))
+
+	cache.SetWithCleanup("item1", "one", nil, func(item *ttlmap.Item) { cleanup++ })
+
+	if cleanup != 0 {
+		t.Errorf("Cache item cleaned up too early")
+	}
+
+	cache.Remove("item1")
+
+	if cleanup != 1 {
+		t.Errorf("Cache item cleanup not called on Remove")
+	}
+
+	cache.SetWithCleanup("item2", "two", &dur, func(item *ttlmap.Item) { cleanup++ })
+	time.Sleep(dur)
+	// Wait a few more milliseconds for cleanup to run
+	time.Sleep(time.Millisecond * 20)
+
+	if cleanup != 2 {
+		t.Errorf("Cache item cleanup not called on expiry")
+	}
+}
diff --git a/go.mod b/go.mod
index 26f2ba4..83475d4 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,3 @@
 module github.com/packaged/ttlmap
 
-go 1.12
+go 1.15
diff --git a/item.go b/item.go
index 85c3ca3..0a948df 100644
--- a/item.go
+++ b/item.go
@@ -12,13 +12,15 @@ type Item struct {
 	deadline time.Time
 	ttl      time.Duration
 	expires  *time.Time
+	onDelete func(*Item)
 }
 
-func newItem(value interface{}, duration time.Duration, deadline time.Time) *Item {
+func newItem(value interface{}, duration time.Duration, deadline time.Time, onDelete func(*Item)) *Item {
 	i := &Item{
 		data:     value,
 		ttl:      duration,
 		deadline: deadline,
+		onDelete: onDelete,
 	}
 	expiry := time.Now().Add(duration)
 	i.expires = &expiry