From 787379640fb1a24c85624a5f798fdd0a3edc4b05 Mon Sep 17 00:00:00 2001
From: Ben Parli <ben.parli@doordash.com>
Date: Sun, 22 Nov 2020 16:45:22 -0800
Subject: [PATCH] fix size calculation

---
 simplelfuda/lfuda.go      | 23 ++++++++++++++---------
 simplelfuda/lfuda_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/simplelfuda/lfuda.go b/simplelfuda/lfuda.go
index 1c9015c..6e1240b 100644
--- a/simplelfuda/lfuda.go
+++ b/simplelfuda/lfuda.go
@@ -2,6 +2,7 @@ package simplelfuda
 
 import (
 	"container/list"
+	"encoding/binary"
 	"fmt"
 )
 
@@ -111,15 +112,7 @@ func (l *LFUDA) Set(key interface{}, value interface{}) bool {
 	} else {
 		// check if we need to evict
 		// convert to bytes so we can get the size of the value
-
-		var numBytes float64
-		// if the value is binary
-		if valBytes, ok := value.([]byte); ok {
-			numBytes = float64(len(valBytes))
-		} else {
-			// otherwise use the default format
-			numBytes = float64(len([]byte(fmt.Sprintf("%v", value.(interface{})))))
-		}
+		numBytes := calcBytes(value)
 
 		// check this value will even fit in the cache.  if not just return
 		if l.size < numBytes {
@@ -305,3 +298,15 @@ func gdsfPolicy(element *item, cacheAge float64) float64 {
 func lfuPolicy(element *item, cacheAge float64) float64 {
 	return element.hits
 }
+
+func calcBytes(value interface{}) float64 {
+	// if the value is binary
+	if valBytes, ok := value.([]byte); ok {
+		return float64(len(valBytes))
+	} else if valBytes := binary.Size(value); valBytes != -1 {
+		return float64(valBytes)
+	} else {
+		// otherwise use the default format
+		return float64(len([]byte(fmt.Sprintf("%v", value.(interface{})))))
+	}
+}
diff --git a/simplelfuda/lfuda_test.go b/simplelfuda/lfuda_test.go
index aac68e7..8dac91a 100644
--- a/simplelfuda/lfuda_test.go
+++ b/simplelfuda/lfuda_test.go
@@ -313,3 +313,38 @@ func TestEvictLFU(t *testing.T) {
 		t.Errorf("cache should still contain key a")
 	}
 }
+
+func TestCalcBytes(t *testing.T) {
+	a := make([]int16, 0, 1)       // 0
+	b := [...]int8{2, 3, 5, 7, 11} // 5
+	c := []int64{2, 3, 5, 7, 11}   // 40
+	d := []int32{}                 // 0
+
+	if res := calcBytes(append(a, 2)); res != float64(2) {
+		t.Errorf("Size is not correct.  Got %f but should be %d", res, 3)
+	}
+
+	if res := calcBytes(b); res != float64(5) {
+		t.Errorf("Size is not correct.  Got %f but should be %d", res, 5)
+	}
+	if res := calcBytes(c); res != float64(40) {
+		t.Errorf("Size is not correct.  Got %f but should be %d", res, 40)
+	}
+	if res := calcBytes(d); res != float64(0) {
+		t.Errorf("Size is not correct.  Got %f but should be %d", res, 0)
+	}
+
+	tests := map[interface{}]int{
+		true:             1,
+		int16(5):         2,
+		1278624387349734: 16,
+		"We must not confuse dissent with disloyalty. We must remember always that accusation is not proof and that conviction depends upon evidence and due process of law. We will not walk in fear, one of another. We will not be driven by fear into an age of unreason, if we dig deep in our history and our doctrine, and remember that we are not descended from fearful men — not from men who feared to write, to speak, to associate and to defend causes that were, for the moment, unpopular.": 484,
+		12: 2,
+	}
+
+	for test, res := range tests {
+		if calcBytes(test) != float64(res) {
+			t.Errorf("cache should have size 10 bytes at this point: %d", res)
+		}
+	}
+}