From 0251553fe5575909d56819e1294bab72e5617488 Mon Sep 17 00:00:00 2001
From: ac-rappi <54147417+ac-rappi@users.noreply.github.com>
Date: Thu, 27 Aug 2020 06:35:44 -0300
Subject: [PATCH] Fix in initial capacity allocation (#237)

---
 bigcache.go      |  2 +-
 bigcache_test.go | 27 +++++++++++++++++++++++++++
 config.go        |  6 +++---
 shard.go         |  7 ++++++-
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/bigcache.go b/bigcache.go
index b59eba08..045d1cf8 100644
--- a/bigcache.go
+++ b/bigcache.go
@@ -63,7 +63,7 @@ func newBigCache(config Config, clock clock) (*BigCache, error) {
 		hash:         config.Hasher,
 		config:       config,
 		shardMask:    uint64(config.Shards - 1),
-		maxShardSize: uint32(config.maximumShardSize()),
+		maxShardSize: uint32(config.maximumShardSizeInBytes()),
 		close:        make(chan struct{}),
 	}
 
diff --git a/bigcache_test.go b/bigcache_test.go
index a28cc9fe..4433ffa2 100644
--- a/bigcache_test.go
+++ b/bigcache_test.go
@@ -490,6 +490,33 @@ func TestCacheCapacity(t *testing.T) {
 	assertEqual(t, 40960, cache.Capacity())
 }
 
+func TestCacheInitialCapacity(t *testing.T) {
+	t.Parallel()
+
+	// given
+	cache, _ := NewBigCache(Config{
+		Shards:             1,
+		LifeWindow:         time.Second,
+		MaxEntriesInWindow: 2 * 1024,
+		HardMaxCacheSize:   1,
+		MaxEntrySize:       1024,
+	})
+
+	assertEqual(t, 0, cache.Len())
+	assertEqual(t, 1024*1024, cache.Capacity())
+
+	keys := 1024 * 1024
+
+	// when
+	for i := 0; i < keys; i++ {
+		cache.Set(fmt.Sprintf("key%d", i), []byte("value"))
+	}
+
+	// then
+	assertEqual(t, true, cache.Len() < keys)
+	assertEqual(t, 1024*1024, cache.Capacity())
+}
+
 func TestRemoveEntriesWhenShardIsFull(t *testing.T) {
 	t.Parallel()
 
diff --git a/config.go b/config.go
index e8cce086..3f500bf4 100644
--- a/config.go
+++ b/config.go
@@ -68,11 +68,11 @@ func DefaultConfig(eviction time.Duration) Config {
 
 // initialShardSize computes initial shard size
 func (c Config) initialShardSize() int {
-	return max(min(c.MaxEntriesInWindow/c.Shards, c.maximumShardSize()), minimumEntriesInShard)
+	return max(c.MaxEntriesInWindow/c.Shards, minimumEntriesInShard)
 }
 
-// maximumShardSize computes maximum shard size
-func (c Config) maximumShardSize() int {
+// maximumShardSizeInBytes computes maximum shard size in bytes
+func (c Config) maximumShardSizeInBytes() int {
 	maxShardSize := 0
 
 	if c.HardMaxCacheSize > 0 {
diff --git a/shard.go b/shard.go
index a11e7060..b432b331 100644
--- a/shard.go
+++ b/shard.go
@@ -389,10 +389,15 @@ func (s *cacheShard) collision() {
 }
 
 func initNewShard(config Config, callback onRemoveCallback, clock clock) *cacheShard {
+	bytesQueueInitialCapacity := config.initialShardSize() * config.MaxEntrySize
+	maximumShardSizeInBytes := config.maximumShardSizeInBytes()
+	if maximumShardSizeInBytes > 0 && bytesQueueInitialCapacity > maximumShardSizeInBytes {
+		bytesQueueInitialCapacity = maximumShardSizeInBytes
+	}
 	return &cacheShard{
 		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),
+		entries:      *queue.NewBytesQueue(bytesQueueInitialCapacity, maximumShardSizeInBytes, config.Verbose),
 		entryBuffer:  make([]byte, config.MaxEntrySize+headersSizeInBytes),
 		onRemove:     callback,