Repository: allegro/bigcache (go)

Issue:
Entry stats
Being able to see the number of times a cached resource was requested is a useful metric, especially at the point of removal.

As a production example, `OnRemove` could be used to trigger cache pre-warming based on the frequency of cache hits on a resource within a given window.

I'm aware of the overhead of storing such data by default and suggest this only as an optional parameter.

Thoughts?

Pull request description:
Add entry stats
This pull request resolves #159.

- Adding a way to access entry-level stats, that can be enabled/disabled using the config attribute `StatsEnabled` which is set to false by default.
- Adding a way to get metadata related to a particular key using `KeyMetaData(key string)`
-  OnRemove now can give you the possibility to access the entry metadata using `OnRemoveWithMetaData`. 
- Writing unit tests for the added functions.

Patch:
diff --git a/bigcache.go b/bigcache.go
index 950991e..afe10ad 100644
--- a/bigcache.go
+++ b/bigcache.go
@@ -69,7 +69,9 @@ func newBigCache(config Config, clock clock) (*BigCache, error) {
 	}
 
 	var onRemove func(wrappedEntry []byte, reason RemoveReason)
-	if config.OnRemove != nil {
+	if config.OnRemoveWithMetadata != nil {
+		onRemove = cache.providedOnRemoveWithMetadata
+	} else if config.OnRemove != nil {
 		onRemove = cache.providedOnRemove
 	} else if config.OnRemoveWithReason != nil {
 		onRemove = cache.providedOnRemoveWithReason
@@ -179,6 +181,13 @@ func (c *BigCache) Stats() Stats {
 	return s
 }
 
+// KeyMetadata returns number of times a cached resource was requested.
+func (c *BigCache) KeyMetadata(key string) Metadata {
+	hashedKey := c.hash.Sum64(key)
+	shard := c.getShard(hashedKey)
+	return shard.getKeyMetadata(hashedKey)
+}
+
 // Iterator returns iterator function to iterate over EntryInfo's from whole cache.
 func (c *BigCache) Iterator() *EntryInfoIterator {
 	return newIterator(c)
@@ -215,3 +224,9 @@ func (c *BigCache) providedOnRemoveWithReason(wrappedEntry []byte, reason Remove
 
 func (c *BigCache) notProvidedOnRemove(wrappedEntry []byte, reason RemoveReason) {
 }
+
+func (c *BigCache) providedOnRemoveWithMetadata(wrappedEntry []byte, reason RemoveReason) {
+	hashedKey := c.hash.Sum64(readKeyFromEntry(wrappedEntry))
+	shard := c.getShard(hashedKey)
+	c.config.OnRemoveWithMetadata(readKeyFromEntry(wrappedEntry), readEntry(wrappedEntry), shard.getKeyMetadata(hashedKey))
+}
diff --git a/config.go b/config.go
index 918908c..c94d7be 100644
--- a/config.go
+++ b/config.go
@@ -16,6 +16,8 @@ type Config struct {
 	MaxEntriesInWindow int
 	// Max size of entry in bytes. Used only to calculate initial size for cache shards.
 	MaxEntrySize int
+	// StatsEnabled if true calculate the number of times a cached resource was requested.
+	StatsEnabled bool
 	// Verbose mode prints information about new memory allocation
 	Verbose bool
 	// Hasher used to map between string keys and unsigned 64bit integers, by default fnv64 hashing is used.
@@ -28,7 +30,12 @@ type Config struct {
 	// OnRemove is a callback fired when the oldest entry is removed because of its expiration time or no space left
 	// for the new entry, or because delete was called.
 	// Default value is nil which means no callback and it prevents from unwrapping the oldest entry.
+	// ignored if OnRemoveWithMetadata is specified.
 	OnRemove func(key string, entry []byte)
+	// OnRemoveWithMetadata is a callback fired when the oldest entry is removed because of its expiration time or no space left
+	// for the new entry, or because delete was called. A structure representing details about that specific entry.
+	// Default value is nil which means no callback and it prevents from unwrapping the oldest entry.
+	OnRemoveWithMetadata func(key string, entry []byte, keyMetadata Metadata)
 	// OnRemoveWithReason is a callback fired when the oldest entry is removed because of its expiration time or no space left
 	// for the new entry, or because delete was called. A constant representing the reason will be passed through.
 	// Default value is nil which means no callback and it prevents from unwrapping the oldest entry.
@@ -51,6 +58,7 @@ func DefaultConfig(eviction time.Duration) Config {
 		CleanWindow:        1 * time.Second,
 		MaxEntriesInWindow: 1000 * 10 * 60,
 		MaxEntrySize:       500,
+		StatsEnabled:       false,
 		Verbose:            true,
 		Hasher:             newDefaultHasher(),
 		HardMaxCacheSize:   0,
diff --git a/shard.go b/shard.go
index 36a6d9d..e19ed4b 100644
--- a/shard.go
+++ b/shard.go
@@ -10,6 +10,11 @@ import (
 
 type onRemoveCallback func(wrappedEntry []byte, reason RemoveReason)
 
+// Metadata contains information of a spesific entry
+type Metadata struct {
+	RequestCount uint32
+}
+
 type cacheShard struct {
 	hashmap     map[uint64]uint32
 	entries     queue.BytesQueue
@@ -17,12 +22,14 @@ type cacheShard struct {
 	entryBuffer []byte
 	onRemove    onRemoveCallback
 
-	isVerbose  bool
-	logger     Logger
-	clock      clock
-	lifeWindow uint64
+	isVerbose    bool
+	statsEnabled bool
+	logger       Logger
+	clock        clock
+	lifeWindow   uint64
 
-	stats Stats
+	hashmapStats map[uint64]uint32
+	stats        Stats
 }
 
 func (s *cacheShard) getWithInfo(key string, hashedKey uint64) (entry []byte, resp Response, err error) {
@@ -48,7 +55,7 @@ func (s *cacheShard) getWithInfo(key string, hashedKey uint64) (entry []byte, re
 		}
 		entry := readEntry(wrappedEntry)
 		s.lock.RUnlock()
-		s.hit()
+		s.hit(hashedKey)
 		return entry, resp, nil
 	}
 	// it is nil & error
@@ -69,7 +76,8 @@ func (s *cacheShard) get(key string, hashedKey uint64) ([]byte, error) {
 		}
 		entry := readEntry(wrappedEntry)
 		s.lock.RUnlock()
-		s.hit()
+		s.hit(hashedKey)
+
 		return entry, nil
 	}
 	// it is nil & error
@@ -165,6 +173,9 @@ func (s *cacheShard) del(hashedKey uint64) error {
 
 		delete(s.hashmap, hashedKey)
 		s.onRemove(wrappedEntry, Deleted)
+		if s.statsEnabled {
+			delete(s.hashmapStats, hashedKey)
+		}
 		resetKeyFromEntry(wrappedEntry)
 	}
 	s.lock.Unlock()
@@ -227,6 +238,9 @@ func (s *cacheShard) removeOldestEntry(reason RemoveReason) error {
 		hash := readHashFromEntry(oldest)
 		delete(s.hashmap, hash)
 		s.onRemove(oldest, reason)
+		if s.statsEnabled {
+			delete(s.hashmapStats, hash)
+		}
 		return nil
 	}
 	return err
@@ -265,8 +279,19 @@ func (s *cacheShard) getStats() Stats {
 	return stats
 }
 
-func (s *cacheShard) hit() {
+func (s *cacheShard) getKeyMetadata(key uint64) Metadata {
+	return Metadata{
+		RequestCount: s.hashmapStats[key],
+	}
+}
+
+func (s *cacheShard) hit(key uint64) {
 	atomic.AddInt64(&s.stats.Hits, 1)
+	if s.statsEnabled {
+		s.lock.RLock()
+		s.hashmapStats[key]++
+		s.lock.RUnlock()
+	}
 }
 
 func (s *cacheShard) miss() {
@@ -287,14 +312,16 @@ func (s *cacheShard) collision() {
 
 func initNewShard(config Config, callback onRemoveCallback, clock clock) *cacheShard {
 	return &cacheShard{
-		hashmap:     make(map[uint64]uint32, config.initialShardSize()),
-		entries:     *queue.NewBytesQueue(config.initialShardSize()*config.MaxEntrySize, config.maximumShardSize(), config.Verbose),
-		entryBuffer: make([]byte, config.MaxEntrySize+headersSizeInBytes),
-		onRemove:    callback,
-
-		isVerbose:  config.Verbose,
-		logger:     newLogger(config.Logger),
-		clock:      clock,
-		lifeWindow: uint64(config.LifeWindow.Seconds()),
+		hashmap:      make(map[uint64]uint32, config.initialShardSize()),
+		hashmapStats: make(map[uint64]uint32, config.initialShardSize()),
+		entries:      *queue.NewBytesQueue(config.initialShardSize()*config.MaxEntrySize, config.maximumShardSize(), config.Verbose),
+		entryBuffer:  make([]byte, config.MaxEntrySize+headersSizeInBytes),
+		onRemove:     callback,
+
+		isVerbose:    config.Verbose,
+		logger:       newLogger(config.Logger),
+		clock:        clock,
+		lifeWindow:   uint64(config.LifeWindow.Seconds()),
+		statsEnabled: config.StatsEnabled,
 	}
 }


Test patch:
diff --git a/bigcache_test.go b/bigcache_test.go
index 4d5e2cb..d82ea0a 100644
--- a/bigcache_test.go
+++ b/bigcache_test.go
@@ -275,6 +275,39 @@ func TestOnRemoveFilter(t *testing.T) {
 	assert.True(t, onRemoveInvoked)
 }
 
+func TestOnRemoveGetEntryStats(t *testing.T) {
+	t.Parallel()
+
+	// given
+	clock := mockedClock{value: 0}
+	count := uint32(0)
+	onRemove := func(key string, entry []byte, keyMetadata Metadata) {
+		count = keyMetadata.RequestCount
+	}
+	c := Config{
+		Shards:               1,
+		LifeWindow:           time.Second,
+		MaxEntriesInWindow:   1,
+		MaxEntrySize:         256,
+		OnRemoveWithMetadata: onRemove,
+		StatsEnabled:         true,
+	}.OnRemoveFilterSet(Deleted, NoSpace)
+
+	cache, _ := newBigCache(c, &clock)
+
+	// when
+	cache.Set("key", []byte("value"))
+
+	for i := 0; i < 100; i++ {
+		cache.Get("key")
+	}
+
+	cache.Delete("key")
+
+	// then
+	assert.Equal(t, uint32(100), count)
+}
+
 func TestCacheLen(t *testing.T) {
 	t.Parallel()
 
@@ -359,6 +392,29 @@ func TestCacheStats(t *testing.T) {
 	assert.Equal(t, stats.DelHits, int64(10))
 	assert.Equal(t, stats.DelMisses, int64(10))
 }
+func TestCacheEntryStats(t *testing.T) {
+	t.Parallel()
+
+	// given
+	cache, _ := NewBigCache(Config{
+		Shards:             8,
+		LifeWindow:         time.Second,
+		MaxEntriesInWindow: 1,
+		MaxEntrySize:       256,
+		StatsEnabled:       true,
+	})
+
+	cache.Set("key0", []byte("value"))
+
+	for i := 0; i < 10; i++ {
+		_, err := cache.Get("key0")
+		assert.Nil(t, err)
+	}
+
+	// then
+	keyMetadata := cache.KeyMetadata("key0")
+	assert.Equal(t, uint32(10), keyMetadata.RequestCount)
+}
 
 func TestCacheDel(t *testing.T) {
 	t.Parallel()