From 3385784411ac24a8be403f7938890ec67ef6e0d6 Mon Sep 17 00:00:00 2001
From: Karl Seguin <karl@openmymind.net>
Date: Sat, 26 Jan 2019 12:33:50 +0700
Subject: [PATCH] Add cache.ItemCount() intt64 API

---
 bucket.go            |  6 ++++++
 cache.go             |  8 ++++++++
 cache_test.go        |  8 +++++++-
 layeredbucket.go     | 10 ++++++++++
 layeredcache.go      |  8 ++++++++
 layeredcache_test.go |  5 +++++
 6 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/bucket.go b/bucket.go
index d675351..cef9349 100644
--- a/bucket.go
+++ b/bucket.go
@@ -10,6 +10,12 @@ type bucket struct {
 	lookup map[string]*Item
 }
 
+func (b *bucket) itemCount() int {
+	b.RLock()
+	defer b.RUnlock()
+	return len(b.lookup)
+}
+
 func (b *bucket) get(key string) *Item {
 	b.RLock()
 	defer b.RUnlock()
diff --git a/cache.go b/cache.go
index d8d74ed..28feb6e 100644
--- a/cache.go
+++ b/cache.go
@@ -37,6 +37,14 @@ func New(config *Configuration) *Cache {
 	return c
 }
 
+func (c *Cache) ItemCount() int {
+	count := 0
+	for _, b := range c.buckets {
+		count += b.itemCount()
+	}
+	return count
+}
+
 // Get an item from the cache. Returns nil if the item wasn't found.
 // This can return an expired item. Use item.Expired() to see if the item
 // is expired and item.TTL() to see how long until the item expires (which
diff --git a/cache_test.go b/cache_test.go
index bcbefaa..67bb52e 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -16,15 +16,19 @@ func Test_Cache(t *testing.T) {
 
 func (_ CacheTests) DeletesAValue() {
 	cache := New(Configure())
+	Expect(cache.ItemCount()).To.Equal(0)
+
 	cache.Set("spice", "flow", time.Minute)
 	cache.Set("worm", "sand", time.Minute)
+	Expect(cache.ItemCount()).To.Equal(2)
+
 	cache.Delete("spice")
 	Expect(cache.Get("spice")).To.Equal(nil)
 	Expect(cache.Get("worm").Value()).To.Equal("sand")
+	Expect(cache.ItemCount()).To.Equal(1)
 }
 
 func (_ CacheTests) OnDeleteCallbackCalled() {
-
 	onDeleteFnCalled := false
 	onDeleteFn := func(item *Item) {
 		if item.key == "spice" {
@@ -66,6 +70,7 @@ func (_ CacheTests) GCsTheOldestItems() {
 	gcCache(cache)
 	Expect(cache.Get("9")).To.Equal(nil)
 	Expect(cache.Get("10").Value()).To.Equal(10)
+	Expect(cache.ItemCount()).To.Equal(490)
 }
 
 func (_ CacheTests) PromotedItemsDontGetPruned() {
@@ -114,6 +119,7 @@ func (_ CacheTests) RemovesOldestItemWhenFull() {
 	Expect(cache.Get("1")).To.Equal(nil)
 	Expect(cache.Get("2").Value()).To.Equal(2)
 	Expect(onDeleteFnCalled).To.Equal(true)
+	Expect(cache.ItemCount()).To.Equal(5)
 }
 
 func (_ CacheTests) RemovesOldestItemWhenFullBySizer() {
diff --git a/layeredbucket.go b/layeredbucket.go
index 88f3def..57676ab 100644
--- a/layeredbucket.go
+++ b/layeredbucket.go
@@ -10,6 +10,16 @@ type layeredBucket struct {
 	buckets map[string]*bucket
 }
 
+func (b *layeredBucket) itemCount() int {
+	count := 0
+	b.RLock()
+	defer b.RUnlock()
+	for _, b := range b.buckets {
+		count += b.itemCount()
+	}
+	return count
+}
+
 func (b *layeredBucket) get(primary, secondary string) *Item {
 	bucket := b.getSecondaryBucket(primary)
 	if bucket == nil {
diff --git a/layeredcache.go b/layeredcache.go
index 1027d1b..f57f5bf 100644
--- a/layeredcache.go
+++ b/layeredcache.go
@@ -49,6 +49,14 @@ func Layered(config *Configuration) *LayeredCache {
 	return c
 }
 
+func (c *LayeredCache) ItemCount() int {
+	count := 0
+	for _, b := range c.buckets {
+		count += b.itemCount()
+	}
+	return count
+}
+
 // Get an item from the cache. Returns nil if the item wasn't found.
 // This can return an expired item. Use item.Expired() to see if the item
 // is expired and item.TTL() to see how long until the item expires (which
diff --git a/layeredcache_test.go b/layeredcache_test.go
index a9c45fd..e53afa4 100644
--- a/layeredcache_test.go
+++ b/layeredcache_test.go
@@ -17,6 +17,7 @@ func Test_LayeredCache(t *testing.T) {
 func (_ *LayeredCacheTests) GetsANonExistantValue() {
 	cache := newLayered()
 	Expect(cache.Get("spice", "flow")).To.Equal(nil)
+	Expect(cache.ItemCount()).To.Equal(0)
 }
 
 func (_ *LayeredCacheTests) SetANewValue() {
@@ -24,6 +25,7 @@ func (_ *LayeredCacheTests) SetANewValue() {
 	cache.Set("spice", "flow", "a value", time.Minute)
 	Expect(cache.Get("spice", "flow").Value()).To.Equal("a value")
 	Expect(cache.Get("spice", "stop")).To.Equal(nil)
+	Expect(cache.ItemCount()).To.Equal(1)
 }
 
 func (_ *LayeredCacheTests) SetsMultipleValueWithinTheSameLayer() {
@@ -38,6 +40,7 @@ func (_ *LayeredCacheTests) SetsMultipleValueWithinTheSameLayer() {
 	Expect(cache.Get("leto", "sister").Value()).To.Equal("ghanima")
 	Expect(cache.Get("leto", "brother")).To.Equal(nil)
 	Expect(cache.Get("baron", "friend")).To.Equal(nil)
+	Expect(cache.ItemCount()).To.Equal(3)
 }
 
 func (_ *LayeredCacheTests) ReplaceDoesNothingIfKeyDoesNotExist() {
@@ -51,6 +54,7 @@ func (_ *LayeredCacheTests) ReplaceUpdatesTheValue() {
 	cache.Set("spice", "flow", "value-a", time.Minute)
 	Expect(cache.Replace("spice", "flow", "value-b")).To.Equal(true)
 	Expect(cache.Get("spice", "flow").Value().(string)).To.Equal("value-b")
+	Expect(cache.ItemCount()).To.Equal(1)
 	//not sure how to test that the TTL hasn't changed sort of a sleep..
 }
 
@@ -64,6 +68,7 @@ func (_ *LayeredCacheTests) DeletesAValue() {
 	Expect(cache.Get("spice", "must").Value()).To.Equal("value-b")
 	Expect(cache.Get("spice", "worm")).To.Equal(nil)
 	Expect(cache.Get("leto", "sister").Value()).To.Equal("ghanima")
+	Expect(cache.ItemCount()).To.Equal(2)
 }
 
 func (_ *LayeredCacheTests) OnDeleteCallbackCalled() {