Merge pull request #4 from geniussportsgroup/lz4-compress

Lz4 compress

go.mod

@@ -2,10 +2,14 @@ module github.com/geniussportsgroup/simple_cache
 
 go 1.18
 
-require github.com/stretchr/testify v1.7.1
+require (
+	github.com/pierrec/lz4 v2.6.1+incompatible
+	github.com/stretchr/testify v1.7.1
+)
 
 require (
 	github.com/davecgh/go-spew v1.1.0 // indirect
+	github.com/frankban/quicktest v1.14.4 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
 )

simple_cache.go

@@ -1,9 +1,12 @@
 package simple_cache
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/pierrec/lz4"
+	"io"
 	"math"
 	"sync"
 	"time"
@@ -36,8 +39,10 @@ type SimpleCache struct {
 	capacity         int
 	extendedCapacity int
 	numEntries       int
-
-	toMapKey func(key interface{}) (string, error)
+	toCompress       bool
+	toMapKey         func(key interface{}) (string, error)
+	valueToBytes     func(value interface{}) ([]byte, error)
+	bytesToValue     func([]byte) (interface{}, error)
 }
 
 func (cache *SimpleCache) MissCount() int {
@@ -99,6 +104,22 @@ func New(capacity int, capFactor float64, ttl time.Duration,
 	return ret
 }
 
+func NewWithCompression(capacity int, capFactor float64, ttl time.Duration,
+	toMapKey func(key interface{}) (string, error),
+	valueToBytes func(value interface{}) ([]byte, error),
+	bytesToValue func([]byte) (interface{}, error),
+) *SimpleCache {
+
+	cache := New(capacity, capFactor, ttl, toMapKey)
+	if cache != nil {
+		cache.toCompress = true
+		cache.valueToBytes = valueToBytes
+		cache.bytesToValue = bytesToValue
+	}
+
+	return cache
+}
+
 func (entry *SimpleCacheEntry) hasExpired(currTime time.Time) bool {
 	return entry.expirationTime.Before(currTime)
 }
@@ -205,18 +226,31 @@ func (cache *SimpleCache) InsertOrUpdate(key interface{}, value interface{}) (in
 		}
 	}
 
+	if cache.toCompress {
+		buf, err := cache.valueToBytes(value)
+		if err != nil {
+			return nil, err
+		}
+		entry.value, err = lz4Compress(buf)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		entry.value = value
+	}
+
 	cache.hitCount++
-	entry.value = value
 	entry.timestamp = currTime
 	entry.expirationTime = currTime.Add(cache.ttl)
 	return entry.value, nil
 }
 
 // Read Retrieves the associates value to key. Return error if the key stringification fails,
 // the key is not in the cache, or if the key has expired
-func (cache *SimpleCache) Read(key interface{}) (interface{}, error) {
+func (cache *SimpleCache) Read(key interface{}) (value interface{}, err error) {
 
-	stringKey, err := cache.toMapKey(key)
+	var stringKey string
+	stringKey, err = cache.toMapKey(key)
 	if err != nil {
 		return nil, err
 	}
@@ -241,7 +275,22 @@ func (cache *SimpleCache) Read(key interface{}) (interface{}, error) {
 	entry.expirationTime = currTime.Add(cache.ttl)
 	cache.becomeMru(entry)
 
-	return entry.value, nil
+	if cache.toCompress {
+		var buf []byte
+		buf, err = lz4Decompress(entry.value.([]byte))
+		if err != nil {
+			return nil, err
+		}
+
+		value, err = cache.bytesToValue(buf)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		value = entry.value
+	}
+
+	return value, nil
 }
 
 // GetMRU Return the most recently used entry in the cache. The method do not refresh the entry
@@ -333,7 +382,7 @@ func (cache *SimpleCache) clean() error {
 	return nil
 }
 
-// Clean Clean the cache. All the entries are deleted and counters reset.
+// Clean the cache. All the entries are deleted and counters reset.
 //
 // Uses internal lock
 //
@@ -344,3 +393,29 @@ func (cache *SimpleCache) Clean() error {
 
 	return cache.clean()
 }
+
+func lz4Compress(in []byte) ([]byte, error) {
+	r := bytes.NewReader(in)
+	w := &bytes.Buffer{}
+	zw := lz4.NewWriter(w)
+	_, err := io.Copy(zw, r)
+	if err != nil {
+		return nil, err
+	}
+	// Closing is *very* important
+	if err := zw.Close(); err != nil {
+		return nil, err
+	}
+	return w.Bytes(), nil
+}
+
+func lz4Decompress(in []byte) ([]byte, error) {
+	r := bytes.NewReader(in)
+	w := &bytes.Buffer{}
+	zr := lz4.NewReader(r)
+	_, err := io.Copy(w, zr)
+	if err != nil {
+		return nil, err
+	}
+	return w.Bytes(), nil
+}

simple_cache_test.go

@@ -1,6 +1,7 @@
 package simple_cache
 
 import (
+	"encoding/json"
 	"fmt"
 	"github.com/stretchr/testify/assert"
 	"strconv"
@@ -105,3 +106,49 @@ func TestSimpleCache(t *testing.T) {
 	assert.Equal(t, key, strconv.Itoa(Capacity))
 	assert.Equal(t, mruValue.(int), Capacity)
 }
+
+type ValueType struct {
+	Num  int
+	Text string
+}
+
+func TestCompress(t *testing.T) {
+
+	cache := NewWithCompression(Capacity, Factor, 5*time.Hour,
+		func(key interface{}) (string, error) {
+			return strconv.Itoa(key.(int)), nil
+		}, func(value interface{}) ([]byte, error) {
+			content := value.(*ValueType)
+			b, err := json.Marshal(content)
+			if err != nil {
+				return nil, err
+			}
+			return b, nil
+		},
+		func(buf []byte) (interface{}, error) {
+			value := &ValueType{}
+			err := json.Unmarshal(buf, value)
+			if err != nil {
+				return nil, err
+			}
+			return value, nil
+		})
+
+	for i := 0; i < Capacity; i++ {
+		str := fmt.Sprintf("This is the %d-th string", i)
+		_, err := cache.InsertOrUpdate(i, &ValueType{
+			Num:  i,
+			Text: str,
+		})
+		assert.NoError(t, err)
+	}
+
+	for i := 0; i < Capacity; i++ {
+		expStr := fmt.Sprintf("This is the %d-th string", i)
+		inter, err := cache.Read(i)
+		assert.NoError(t, err)
+		value := inter.(*ValueType)
+		assert.NotNil(t, value)
+		assert.Equal(t, expStr, value.Text)
+	}
+}