Fix potential LRU cache size issue.

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

1 files changed, 13 insertions, 11 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