From 0dbf3f125f135d89ecfa28bdc54c2230b9d1ab50 Mon Sep 17 00:00:00 2001
From: Karl Seguin <k@openmymind.io>
Date: Fri, 14 Aug 2020 11:15:13 +0800
Subject: [PATCH] add TrackingSet to LayeredCache

---
 layeredbucket.go     |  4 ++--
 layeredcache.go      | 13 +++++++++----
 layeredcache_test.go | 11 +++++++----
 readme.md            |  2 +-
 4 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/layeredbucket.go b/layeredbucket.go
index da83b08..18f4a71 100644
--- a/layeredbucket.go
+++ b/layeredbucket.go
@@ -38,7 +38,7 @@ func (b *layeredBucket) getSecondaryBucket(primary string) *bucket {
 	return bucket
 }
 
-func (b *layeredBucket) set(primary, secondary string, value interface{}, duration time.Duration) (*Item, *Item) {
+func (b *layeredBucket) set(primary, secondary string, value interface{}, duration time.Duration, track bool) (*Item, *Item) {
 	b.Lock()
 	bkt, exists := b.buckets[primary]
 	if exists == false {
@@ -46,7 +46,7 @@ func (b *layeredBucket) set(primary, secondary string, value interface{}, durati
 		b.buckets[primary] = bkt
 	}
 	b.Unlock()
-	item, existing := bkt.set(secondary, value, duration, false)
+	item, existing := bkt.set(secondary, value, duration, track)
 	item.group = primary
 	return item, existing
 }
diff --git a/layeredcache.go b/layeredcache.go
index e95977f..2a846df 100644
--- a/layeredcache.go
+++ b/layeredcache.go
@@ -102,9 +102,14 @@ func (c *LayeredCache) TrackingGet(primary, secondary string) TrackedItem {
 	return item
 }
 
+// Set the value in the cache for the specified duration
+func (c *LayeredCache) TrackingSet(primary, secondary string, value interface{}, duration time.Duration) TrackedItem {
+	return c.set(primary, secondary, value, duration, true)
+}
+
 // Set the value in the cache for the specified duration
 func (c *LayeredCache) Set(primary, secondary string, value interface{}, duration time.Duration) {
-	c.set(primary, secondary, value, duration)
+	c.set(primary, secondary, value, duration, false)
 }
 
 // Replace the value if it exists, does not set if it doesn't.
@@ -131,7 +136,7 @@ func (c *LayeredCache) Fetch(primary, secondary string, duration time.Duration,
 	if err != nil {
 		return nil, err
 	}
-	return c.set(primary, secondary, value, duration), nil
+	return c.set(primary, secondary, value, duration, false), nil
 }
 
 // Remove the item from the cache, return true if the item was present, false otherwise.
@@ -193,8 +198,8 @@ func (c *LayeredCache) restart() {
 	go c.worker()
 }
 
-func (c *LayeredCache) set(primary, secondary string, value interface{}, duration time.Duration) *Item {
-	item, existing := c.bucket(primary).set(primary, secondary, value, duration)
+func (c *LayeredCache) set(primary, secondary string, value interface{}, duration time.Duration, track bool) *Item {
+	item, existing := c.bucket(primary).set(primary, secondary, value, duration, track)
 	if existing != nil {
 		c.deletables <- existing
 	}
diff --git a/layeredcache_test.go b/layeredcache_test.go
index 6ca33ce..1338be0 100644
--- a/layeredcache_test.go
+++ b/layeredcache_test.go
@@ -196,17 +196,20 @@ func (_ LayeredCacheTests) PromotedItemsDontGetPruned() {
 
 func (_ LayeredCacheTests) TrackerDoesNotCleanupHeldInstance() {
 	cache := Layered(Configure().ItemsToPrune(10).Track())
-	for i := 0; i < 10; i++ {
+	item0 := cache.TrackingSet("0", "a", 0, time.Minute)
+	for i := 1; i < 11; i++ {
 		cache.Set(strconv.Itoa(i), "a", i, time.Minute)
 	}
-	item := cache.TrackingGet("0", "a")
+	item1 := cache.TrackingGet("1", "a")
 	time.Sleep(time.Millisecond * 10)
 	gcLayeredCache(cache)
 	Expect(cache.Get("0", "a").Value()).To.Equal(0)
-	Expect(cache.Get("1", "a")).To.Equal(nil)
-	item.Release()
+	Expect(cache.Get("1", "a").Value()).To.Equal(1)
+	item0.Release()
+	item1.Release()
 	gcLayeredCache(cache)
 	Expect(cache.Get("0", "a")).To.Equal(nil)
+	Expect(cache.Get("1", "a")).To.Equal(nil)
 }
 
 func (_ LayeredCacheTests) RemovesOldestItemWhenFull() {
diff --git a/readme.md b/readme.md
index 54f4de0..8efbe45 100644
--- a/readme.md
+++ b/readme.md
@@ -143,7 +143,7 @@ user := item.Value()   //will be nil if "user:4" didn't exist in the cache
 item.Release()  //can be called even if item.Value() returned nil
 ```
 
-In practice, `Release` wouldn't be called until later, at some other place in your code.
+In practice, `Release` wouldn't be called until later, at some other place in your code. `TrackingSet` can be used to set a value to be tracked.
 
 There's a couple reason to use the tracking mode if other parts of your code also hold references to objects. First, if you're already going to hold a reference to these objects, there's really no reason not to have them in the cache - the memory is used up anyways.