Fix potential LRU cache size issue.

Signed-off-by: kuba-- <kuba@sourced.tech>

4 files changed, 68 insertions, 22 deletions

diff --git a/plumbing/cache/buffer_lru.go b/plumbing/cache/buffer_lru.go
index f2c0f90..e86ccb2 100644
--- a/plumbing/cache/buffer_lru.go
+++ b/plumbing/cache/buffer_lru.go
@@ -45,19 +45,25 @@ func (c *BufferLRU) Put(key int64, slice []byte) {
 		c.ll = list.New()
 	}
 
+	bufSize := FileSize(len(slice))
 	if ee, ok := c.cache[key]; ok {
+		oldBuf := ee.Value.(buffer)
+		// in this case bufSize is a delta: new size - old size
+		bufSize -= FileSize(len(oldBuf.Slice))
+
 		c.ll.MoveToFront(ee)
 		ee.Value = buffer{key, slice}
-		return
-	}
+	} else {
+		if bufSize > c.MaxSize {
+			return
+		}
 
-	objSize := FileSize(len(slice))
-
-	if objSize > c.MaxSize {
-		return
+		ee := c.ll.PushFront(buffer{key, slice})
+		c.cache[key] = ee
 	}
 
-	for c.actualSize+objSize > c.MaxSize {
+	c.actualSize += bufSize
+	for c.actualSize > c.MaxSize {
 		last := c.ll.Back()
 		lastObj := last.Value.(buffer)
 		lastSize := FileSize(len(lastObj.Slice))
@@ -66,10 +72,6 @@ func (c *BufferLRU) Put(key int64, slice []byte) {
 		delete(c.cache, lastObj.Key)
 		c.actualSize -= lastSize
 	}
-
-	ee := c.ll.PushFront(buffer{key, slice})
-	c.cache[key] = ee
-	c.actualSize += objSize
 }
 
 // Get returns a buffer by its key. It marks the buffer as used. If the buffer
diff --git a/plumbing/cache/buffer_test.go b/plumbing/cache/buffer_test.go
index 262138a..3e3adc2 100644
--- a/plumbing/cache/buffer_test.go
+++ b/plumbing/cache/buffer_test.go
@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"bytes"
 	"sync"
 
 	. "gopkg.in/check.v1"
@@ -38,6 +39,28 @@ func (s *BufferSuite) TestPutSameBuffer(c *C) {
 	}
 }
 
+func (s *ObjectSuite) TestPutSameBufferWithDifferentSize(c *C) {
+	aBuffer := []byte("a")
+	bBuffer := []byte("bbb")
+	cBuffer := []byte("ccccc")
+	dBuffer := []byte("ddddddd")
+
+	cache := NewBufferLRU(7 * Byte)
+	cache.Put(1, aBuffer)
+	cache.Put(1, bBuffer)
+	cache.Put(1, cBuffer)
+	cache.Put(1, dBuffer)
+
+	c.Assert(cache.MaxSize, Equals, 7*Byte)
+	c.Assert(cache.actualSize, Equals, 7*Byte)
+	c.Assert(cache.ll.Len(), Equals, 1)
+
+	buf, ok := cache.Get(1)
+	c.Assert(bytes.Equal(buf, dBuffer), Equals, true)
+	c.Assert(FileSize(len(buf)), Equals, 7*Byte)
+	c.Assert(ok, Equals, true)
+}
+
 func (s *BufferSuite) TestPutBigBuffer(c *C) {
 	for _, o := range s.c {
 		o.Put(1, s.bBuffer)
diff --git a/plumbing/cache/object_lru.go b/plumbing/cache/object_lru.go
index 0494539..31c0202 100644
--- a/plumbing/cache/object_lru.go
+++ b/plumbing/cache/object_lru.go
@@ -42,20 +42,26 @@ func (c *ObjectLRU) Put(obj plumbing.EncodedObject) {
 		c.ll = list.New()
 	}
 
+	objSize := FileSize(obj.Size())
 	key := obj.Hash()
 	if ee, ok := c.cache[key]; ok {
+		oldObj := ee.Value.(plumbing.EncodedObject)
+		// in this case objSize is a delta: new size - old size
+		objSize -= FileSize(oldObj.Size())
+
 		c.ll.MoveToFront(ee)
 		ee.Value = obj
-		return
-	}
-
-	objSize := FileSize(obj.Size())
+	} else {
+		if objSize > c.MaxSize {
+			return
+		}
 
-	if objSize > c.MaxSize {
-		return
+		ee := c.ll.PushFront(obj)
+		c.cache[key] = ee
 	}
 
-	for c.actualSize+objSize > c.MaxSize {
+	c.actualSize += objSize
+	for c.actualSize > c.MaxSize {
 		last := c.ll.Back()
 		lastObj := last.Value.(plumbing.EncodedObject)
 		lastSize := FileSize(lastObj.Size())
@@ -64,10 +70,6 @@ func (c *ObjectLRU) Put(obj plumbing.EncodedObject) {
 		delete(c.cache, lastObj.Hash())
 		c.actualSize -= lastSize
 	}
-
-	ee := c.ll.PushFront(obj)
-	c.cache[key] = ee
-	c.actualSize += objSize
 }
 
 // Get returns an object by its hash. It marks the object as used. If the object
diff --git a/plumbing/cache/object_test.go b/plumbing/cache/object_test.go
index ac3f0a3..b3e5f79 100644
--- a/plumbing/cache/object_test.go
+++ b/plumbing/cache/object_test.go
@@ -45,6 +45,25 @@ func (s *ObjectSuite) TestPutSameObject(c *C) {
 	}
 }
 
+func (s *ObjectSuite) TestPutSameObjectWithDifferentSize(c *C) {
+	const hash = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+
+	cache := NewObjectLRU(7 * Byte)
+	cache.Put(newObject(hash, 1*Byte))
+	cache.Put(newObject(hash, 3*Byte))
+	cache.Put(newObject(hash, 5*Byte))
+	cache.Put(newObject(hash, 7*Byte))
+
+	c.Assert(cache.MaxSize, Equals, 7*Byte)
+	c.Assert(cache.actualSize, Equals, 7*Byte)
+	c.Assert(cache.ll.Len(), Equals, 1)
+
+	obj, ok := cache.Get(plumbing.NewHash(hash))
+	c.Assert(obj.Hash(), Equals, plumbing.NewHash(hash))
+	c.Assert(FileSize(obj.Size()), Equals, 7*Byte)
+	c.Assert(ok, Equals, true)
+}
+
 func (s *ObjectSuite) TestPutBigObject(c *C) {
 	for _, o := range s.c {
 		o.Put(s.bObject)