added TTL

cache.go

@@ -24,6 +24,7 @@ import (
 	"errors"
 	"fmt"
 	"sync/atomic"
+	"time"
 
 	"github.com/dgraph-io/ristretto/z"
 )
@@ -120,6 +121,7 @@ type item struct {
 	conflict uint64
 	value    interface{}
 	cost     int64
+	ttl      int64
 }
 
 // NewCache returns a new Cache instance and any configuration errors, if any.
@@ -183,17 +185,24 @@ func (c *Cache) Get(key interface{}) (interface{}, bool) {
 // To dynamically evaluate the items cost using the Config.Coster function, set
 // the cost parameter to 0 and Coster will be ran when needed in order to find
 // the items true cost.
-func (c *Cache) Set(key, value interface{}, cost int64) bool {
+//
+// TTL is the amount of time (in seconds) that the item will remain in the
+// cache. To make an item stay indefinitely, set TTL to -1.
+func (c *Cache) Set(key, value interface{}, cost, ttl int64) bool {
 	if c == nil || key == nil {
 		return false
 	}
 	keyHash, conflictHash := c.keyToHash(key)
+	if ttl != -1 {
+		ttl = time.Now().Unix() + ttl
+	}
 	i := &item{
 		flag:     itemNew,
 		key:      keyHash,
 		conflict: conflictHash,
 		value:    value,
 		cost:     cost,
+		ttl:      ttl,
 	}
 	// attempt to immediately update hashmap value and set flag to update so the
 	// cost is eventually updated
@@ -262,7 +271,7 @@ func (c *Cache) processItems() {
 			}
 			switch i.flag {
 			case itemNew:
-				victims, added := c.policy.Add(i.key, i.cost)
+				victims, added := c.policy.Add(i.key, i.cost, i.ttl)
 				if added {
 					c.store.Set(i.key, i.conflict, i.value)
 					c.Metrics.add(keyAdd, i.key, 1)

cache_test.go

@@ -1,6 +1,7 @@
 package ristretto
 
 import (
+	"fmt"
 	"math/rand"
 	"strings"
 	"sync"
@@ -12,6 +13,58 @@ import (
 
 var wait time.Duration = time.Millisecond * 10
 
+func TestCacheTTL(t *testing.T) {
+	m := &sync.Mutex{}
+	evicted := make(map[uint64]struct{})
+	c, err := NewCache(&Config{
+		NumCounters: 100,
+		MaxCost:     10,
+		BufferItems: 64,
+		OnEvict: func(key, conflict uint64, value interface{}, cost int64) {
+			m.Lock()
+			defer m.Unlock()
+			evicted[key] = struct{}{}
+		},
+	})
+	if err != nil {
+		panic(err)
+	}
+	// item 1 will live for 1 second
+	c.Set(1, 1, 2, 1)
+	// item 2 will live for 2 seconds
+	c.Set(2, 2, 3, 2)
+	// items 11-14 will live indefinitely
+	c.Set(11, 1, 1, -1)
+	c.Set(12, 1, 1, -1)
+	c.Set(13, 1, 1, -1)
+	c.Set(14, 1, 1, -1)
+	// sleep for 3 seconds (2+1 for good measure)
+	time.Sleep(time.Second * 3)
+	// gets to simulate load (1 and 2 should be evicted despite this because
+	// they're expiring)
+	c.Get(1)
+	c.Get(1)
+	c.Get(1)
+	c.Get(2)
+	c.Get(2)
+	// try to set a new item to force 1 and 2 expiration
+	c.Set(3, 3, 5, -1)
+	// wait for new set to go through
+	time.Sleep(time.Millisecond)
+	m.Lock()
+	defer m.Unlock()
+	fmt.Println(evicted)
+	if len(evicted) != 2 {
+		t.Fatal("items 1 and 2 should have expired")
+	}
+	if _, ok := evicted[1]; !ok {
+		t.Fatal("items 1 and 2 should have expired")
+	}
+	if _, ok := evicted[2]; !ok {
+		t.Fatal("items 1 and 2 should have expired")
+	}
+}
+
 func TestCacheKeyToHash(t *testing.T) {
 	keyToHashCount := 0
 	c, err := NewCache(&Config{
@@ -26,7 +79,7 @@ func TestCacheKeyToHash(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-	if c.Set(1, 1, 1) {
+	if c.Set(1, 1, 1, -1) {
 		time.Sleep(wait)
 		if val, ok := c.Get(1); val == nil || !ok {
 			t.Fatal("get should be successful")
@@ -75,7 +128,7 @@ func TestCacheMaxCost(t *testing.T) {
 						} else {
 							val = strings.Repeat("a", 1000)
 						}
-						c.Set(key(), val, int64(2+len(val)))
+						c.Set(key(), val, int64(2+len(val)), -1)
 					}
 				}
 			}
@@ -270,7 +323,7 @@ func TestCacheSet(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-	if c.Set(1, 1, 1) {
+	if c.Set(1, 1, 1, -1) {
 		time.Sleep(wait)
 		if val, ok := c.Get(1); val == nil || val.(int) != 1 || !ok {
 			t.Fatal("set/get returned wrong value")
@@ -280,7 +333,7 @@ func TestCacheSet(t *testing.T) {
 			t.Fatal("set was dropped but value still added")
 		}
 	}
-	c.Set(1, 2, 2)
+	c.Set(1, 2, 2, -1)
 	val, ok := c.store.Get(z.KeyToHash(1))
 	if val == nil || val.(int) != 2 || !ok {
 		t.Fatal("set/update was unsuccessful")
@@ -296,7 +349,7 @@ func TestCacheSet(t *testing.T) {
 			cost:     1,
 		}
 	}
-	if c.Set(2, 2, 1) {
+	if c.Set(2, 2, 1, -1) {
 		t.Fatal("set should be dropped with full setBuf")
 	}
 	if c.Metrics.SetsDropped() != 1 {
@@ -305,7 +358,7 @@ func TestCacheSet(t *testing.T) {
 	close(c.setBuf)
 	close(c.stop)
 	c = nil
-	if c.Set(1, 1, 1) {
+	if c.Set(1, 1, 1, -1) {
 		t.Fatal("set shouldn't be successful with nil cache")
 	}
 }
@@ -319,7 +372,7 @@ func TestCacheDel(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-	c.Set(1, 1, 1)
+	c.Set(1, 1, 1, -1)
 	c.Del(1)
 	time.Sleep(wait)
 	if val, ok := c.Get(1); val != nil || ok {
@@ -345,7 +398,7 @@ func TestCacheClear(t *testing.T) {
 		panic(err)
 	}
 	for i := 0; i < 10; i++ {
-		c.Set(i, i, 1)
+		c.Set(i, i, 1, -1)
 	}
 	time.Sleep(wait)
 	if c.Metrics.KeysAdded() != 10 {
@@ -373,7 +426,7 @@ func TestCacheMetrics(t *testing.T) {
 		panic(err)
 	}
 	for i := 0; i < 10; i++ {
-		c.Set(i, i, 1)
+		c.Set(i, i, 1, -1)
 	}
 	time.Sleep(wait)
 	m := c.Metrics
@@ -460,7 +513,7 @@ func TestCacheMetricsClear(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-	c.Set(1, 1, 1)
+	c.Set(1, 1, 1, -1)
 	stop := make(chan struct{})
 	go func() {
 		for {