Improve memory performance (#195)

- Use an int64 instead of a time.Time struct to represent the time.
- By default, include the cost of the storeItem in the cost calculation.

Related to DGRAPH-1378

Author: martinmr

cache.go

@@ -26,6 +26,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/dgraph-io/ristretto/z"
 )
@@ -37,6 +38,8 @@ var (
 
 type itemCallback func(*Item)
 
+const itemSize = int64(unsafe.Sizeof(storeItem{}))
+
 // Cache is a thread-safe implementation of a hashmap with a TinyLFU admission
 // policy and a Sampled LFU eviction policy. You can use the same Cache instance
 // from as many goroutines as you want.
@@ -65,6 +68,9 @@ type Cache struct {
 	stop chan struct{}
 	// cost calculates cost from a value.
 	cost func(value interface{}) int64
+	// ignoreInternalCost dictates whether to ignore the cost of internally storing
+	// the item in the cost calculation.
+	ignoreInternalCost bool
 	// cleanupTicker is used to periodically check for entries whose TTL has passed.
 	cleanupTicker *time.Ticker
 	// Metrics contains a running log of important statistics like hits, misses,
@@ -119,6 +125,11 @@ type Config struct {
 	// is ran after Set is called for a new item or an item update with a cost
 	// param of 0.
 	Cost func(value interface{}) int64
+	// IgnoreInternalCost set to true indicates to the cache that the cost of
+	// internally storing the value should be ignored. This is useful when the
+	// cost passed to set is not using bytes as units. Keep in mind that setting
+	// this to true will increase the memory usage.
+	IgnoreInternalCost bool
 }
 
 type itemFlag byte
@@ -136,7 +147,7 @@ type Item struct {
 	Conflict   uint64
 	Value      interface{}
 	Cost       int64
-	Expiration time.Time
+	Expiration int64
 	wg         *sync.WaitGroup
 }
 
@@ -152,14 +163,15 @@ func NewCache(config *Config) (*Cache, error) {
 	}
 	policy := newPolicy(config.NumCounters, config.MaxCost)
 	cache := &Cache{
-		store:         newStore(),
-		policy:        policy,
-		getBuf:        newRingBuffer(policy, config.BufferItems),
-		setBuf:        make(chan *Item, setBufSize),
-		keyToHash:     config.KeyToHash,
-		stop:          make(chan struct{}),
-		cost:          config.Cost,
-		cleanupTicker: time.NewTicker(time.Duration(bucketDurationSecs) * time.Second / 2),
+		store:              newStore(),
+		policy:             policy,
+		getBuf:             newRingBuffer(policy, config.BufferItems),
+		setBuf:             make(chan *Item, setBufSize),
+		keyToHash:          config.KeyToHash,
+		stop:               make(chan struct{}),
+		cost:               config.Cost,
+		ignoreInternalCost: config.IgnoreInternalCost,
+		cleanupTicker:      time.NewTicker(time.Duration(bucketDurationSecs) * time.Second / 2),
 	}
 	cache.onExit = func(val interface{}) {
 		if config.OnExit != nil && val != nil {
@@ -241,7 +253,7 @@ func (c *Cache) SetWithTTL(key, value interface{}, cost int64, ttl time.Duration
 		return false
 	}
 
-	var expiration time.Time
+	var expiration int64
 	switch {
 	case ttl == 0:
 		// No expiration.
@@ -250,7 +262,7 @@ func (c *Cache) SetWithTTL(key, value interface{}, cost int64, ttl time.Duration
 		// Treat this a a no-op.
 		return false
 	default:
-		expiration = time.Now().Add(ttl)
+		expiration = time.Now().Add(ttl).Unix()
 	}
 
 	keyHash, conflictHash := c.keyToHash(key)
@@ -393,6 +405,11 @@ func (c *Cache) processItems() {
 			if i.Cost == 0 && c.cost != nil && i.flag != itemDelete {
 				i.Cost = c.cost(i.Value)
 			}
+			if !c.ignoreInternalCost {
+				// Add the cost of internally storing the object.
+				i.Cost += itemSize
+			}
+
 			switch i.flag {
 			case itemNew:
 				victims, added := c.policy.Add(i.Key, i.Cost)

cache_test.go

@@ -19,9 +19,10 @@ var wait = time.Millisecond * 10
 func TestCacheKeyToHash(t *testing.T) {
 	keyToHashCount := 0
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
+		NumCounters:        10,
+		MaxCost:            1000,
+		BufferItems:        64,
+		IgnoreInternalCost: true,
 		KeyToHash: func(key interface{}) (uint64, uint64) {
 			keyToHashCount++
 			return z.KeyToHash(key)
@@ -150,9 +151,10 @@ func TestCacheProcessItems(t *testing.T) {
 	m := &sync.Mutex{}
 	evicted := make(map[uint64]struct{})
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
+		NumCounters:        100,
+		MaxCost:            10,
+		BufferItems:        64,
+		IgnoreInternalCost: true,
 		Cost: func(value interface{}) int64 {
 			return int64(value.(int))
 		},
@@ -249,10 +251,11 @@ func TestCacheProcessItems(t *testing.T) {
 
 func TestCacheGet(t *testing.T) {
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
-		Metrics:     true,
+		NumCounters:        100,
+		MaxCost:            10,
+		BufferItems:        64,
+		IgnoreInternalCost: true,
+		Metrics:            true,
 	})
 	require.NoError(t, err)
 
@@ -299,10 +302,11 @@ func retrySet(t *testing.T, c *Cache, key, value int, cost int64, ttl time.Durat
 
 func TestCacheSet(t *testing.T) {
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
-		Metrics:     true,
+		NumCounters:        100,
+		MaxCost:            10,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
+		Metrics:            true,
 	})
 	require.NoError(t, err)
 
@@ -333,13 +337,31 @@ func TestCacheSet(t *testing.T) {
 	require.False(t, c.Set(1, 1, 1))
 }
 
-func TestRecacheWithTTL(t *testing.T) {
+func TestCacheInternalCost(t *testing.T) {
 	c, err := NewCache(&Config{
 		NumCounters: 100,
 		MaxCost:     10,
 		BufferItems: 64,
 		Metrics:     true,
 	})
+	require.NoError(t, err)
+
+	// Get should return false because the cache's cost is too small to store the item
+	// when accounting for the internal cost.
+	c.SetWithTTL(1, 1, 1, 0)
+	time.Sleep(wait)
+	_, ok := c.Get(1)
+	require.False(t, ok)
+}
+
+func TestRecacheWithTTL(t *testing.T) {
+	c, err := NewCache(&Config{
+		NumCounters:        100,
+		MaxCost:            10,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
+		Metrics:            true,
+	})
 
 	require.NoError(t, err)
 
@@ -378,10 +400,11 @@ func TestCacheSetWithTTL(t *testing.T) {
 	m := &sync.Mutex{}
 	evicted := make(map[uint64]struct{})
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
-		Metrics:     true,
+		NumCounters:        100,
+		MaxCost:            10,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
+		Metrics:            true,
 		OnEvict: func(item *Item) {
 			m.Lock()
 			defer m.Unlock()
@@ -451,9 +474,10 @@ func TestCacheDel(t *testing.T) {
 
 func TestCacheDelWithTTL(t *testing.T) {
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
+		NumCounters:        100,
+		MaxCost:            10,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
 	})
 	require.NoError(t, err)
 	retrySet(t, c, 3, 1, 1, 10*time.Second)
@@ -468,10 +492,11 @@ func TestCacheDelWithTTL(t *testing.T) {
 
 func TestCacheClear(t *testing.T) {
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
-		Metrics:     true,
+		NumCounters:        100,
+		MaxCost:            10,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
+		Metrics:            true,
 	})
 	require.NoError(t, err)
 
@@ -493,10 +518,11 @@ func TestCacheClear(t *testing.T) {
 
 func TestCacheMetrics(t *testing.T) {
 	c, err := NewCache(&Config{
-		NumCounters: 100,
-		MaxCost:     10,
-		BufferItems: 64,
-		Metrics:     true,
+		NumCounters:        100,
+		MaxCost:            10,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
+		Metrics:            true,
 	})
 	require.NoError(t, err)
 

store.go

@@ -26,7 +26,7 @@ type storeItem struct {
 	key        uint64
 	conflict   uint64
 	value      interface{}
-	expiration time.Time
+	expiration int64
 }
 
 // store is the interface fulfilled by all hash map implementations in this
@@ -39,7 +39,7 @@ type store interface {
 	// Get returns the value associated with the key parameter.
 	Get(uint64, uint64) (interface{}, bool)
 	// Expiration returns the expiration time for this key.
-	Expiration(uint64) time.Time
+	Expiration(uint64) int64
 	// Set adds the key-value pair to the Map or updates the value if it's
 	// already present. The key-value pair is passed as a pointer to an
 	// item object.
@@ -82,7 +82,7 @@ func (sm *shardedMap) Get(key, conflict uint64) (interface{}, bool) {
 	return sm.shards[key%numShards].get(key, conflict)
 }
 
-func (sm *shardedMap) Expiration(key uint64) time.Time {
+func (sm *shardedMap) Expiration(key uint64) int64 {
 	return sm.shards[key%numShards].Expiration(key)
 }
 
@@ -138,13 +138,13 @@ func (m *lockedMap) get(key, conflict uint64) (interface{}, bool) {
 	}
 
 	// Handle expired items.
-	if !item.expiration.IsZero() && time.Now().After(item.expiration) {
+	if item.expiration != 0 && time.Now().Unix() > item.expiration {
 		return nil, false
 	}
 	return item.value, true
 }
 
-func (m *lockedMap) Expiration(key uint64) time.Time {
+func (m *lockedMap) Expiration(key uint64) int64 {
 	m.RLock()
 	defer m.RUnlock()
 	return m.data[key].expiration
@@ -193,7 +193,7 @@ func (m *lockedMap) Del(key, conflict uint64) (uint64, interface{}) {
 		return 0, nil
 	}
 
-	if !item.expiration.IsZero() {
+	if item.expiration != 0 {
 		m.em.del(key, item.expiration)
 	}
 

stress_test.go

@@ -15,10 +15,11 @@ import (
 
 func TestStressSetGet(t *testing.T) {
 	c, err := NewCache(&Config{
-		NumCounters: 1000,
-		MaxCost:     100,
-		BufferItems: 64,
-		Metrics:     true,
+		NumCounters:        1000,
+		MaxCost:            100,
+		IgnoreInternalCost: true,
+		BufferItems:        64,
+		Metrics:            true,
 	})
 	require.NoError(t, err)
 

ttl.go

@@ -26,11 +26,11 @@ var (
 	bucketDurationSecs = int64(5)
 )
 
-func storageBucket(t time.Time) int64 {
-	return (t.Unix() / bucketDurationSecs) + 1
+func storageBucket(t int64) int64 {
+	return (t / bucketDurationSecs) + 1
 }
 
-func cleanupBucket(t time.Time) int64 {
+func cleanupBucket(t int64) int64 {
 	// The bucket to cleanup is always behind the storage bucket by one so that
 	// no elements in that bucket (which might not have expired yet) are deleted.
 	return storageBucket(t) - 1
@@ -51,13 +51,13 @@ func newExpirationMap() *expirationMap {
 	}
 }
 
-func (m *expirationMap) add(key, conflict uint64, expiration time.Time) {
+func (m *expirationMap) add(key, conflict uint64, expiration int64) {
 	if m == nil {
 		return
 	}
 
 	// Items that don't expire don't need to be in the expiration map.
-	if expiration.IsZero() {
+	if expiration == 0 {
 		return
 	}
 
@@ -73,7 +73,7 @@ func (m *expirationMap) add(key, conflict uint64, expiration time.Time) {
 	b[key] = conflict
 }
 
-func (m *expirationMap) update(key, conflict uint64, oldExpTime, newExpTime time.Time) {
+func (m *expirationMap) update(key, conflict uint64, oldExpTime, newExpTime int64) {
 	if m == nil {
 		return
 	}
@@ -96,7 +96,7 @@ func (m *expirationMap) update(key, conflict uint64, oldExpTime, newExpTime time
 	newBucket[key] = conflict
 }
 
-func (m *expirationMap) del(key uint64, expiration time.Time) {
+func (m *expirationMap) del(key uint64, expiration int64) {
 	if m == nil {
 		return
 	}
@@ -120,15 +120,15 @@ func (m *expirationMap) cleanup(store store, policy policy, onEvict itemCallback
 	}
 
 	m.Lock()
-	now := time.Now()
+	now := time.Now().Unix()
 	bucketNum := cleanupBucket(now)
 	keys := m.buckets[bucketNum]
 	delete(m.buckets, bucketNum)
 	m.Unlock()
 
 	for key, conflict := range keys {
 		// Sanity check. Verify that the store agrees that this key is expired.
-		if store.Expiration(key).After(now) {
+		if store.Expiration(key) > now {
 			continue
 		}