LFU using ristretto

lfu.go

package cache

import (
	"time"
	"errors"

	"github.com/tinylib/msgp/msgp"
	"github.com/vmihailenco/msgpack/v5"

	"github.com/dgraph-io/ristretto"
)

// ErrKeyNotFound is the error when the given key is not found
var ErrKeyNotFound = errors.New("key not found")

// ErrKeyNotSet is the error when given key is not set
var ErrKeyNotSet = errors.New("key not set")

// LFU cache implementation using the Ristretto library.
type LFU struct {
	name                   string
	size                   int
	cache                  *ristretto.Cache
	defaultExpiry          time.Duration
	invalidateClusterEvent string
	len                    int
}

// LFUOptions contains options for initializing LFU cache
type LFUOptions struct {
	Name                   string
	Size                   int
	DefaultExpiry          time.Duration
	InvalidateClusterEvent string
	Metrics                bool
}

// NewLFU creates an LFU of the given size.
func NewLFU(opts *LFUOptions) LFU {
	lfu, err := ristretto.NewCache(&ristretto.Config{
		NumCounters: int64(10 * opts.Size),
		MaxCost:     int64(opts.Size),
		Metrics:     opts.Metrics,
		BufferItems: 64,
	})

	if err != nil {
		panic(err)
	}

	return LFU{
		name:                   opts.Name,
		size:                   opts.Size,
		cache:                  lfu,
		defaultExpiry:          opts.DefaultExpiry,
		invalidateClusterEvent: opts.InvalidateClusterEvent,
	}
}

// Purge is used to completely clear the cache.
func (l *LFU) Purge() error {
	l.len = 0
	l.cache.Clear()
	return nil
}

// Set adds the given key and value to the store without an expiry. If the key already exists,
// it will overwrite the previous value.
func (l *LFU) Set(key string, value interface{}) error {
	return l.set(key, value, 0)
}

// SetWithDefaultExpiry adds the given key and value to the store with the default expiry. If
// the key already exists, it will overwrite the previoous value
func (l *LFU) SetWithDefaultExpiry(key string, value interface{}) error {
	return l.SetWithExpiry(key, value, l.defaultExpiry)
}

// SetWithExpiry adds the given key and value to the cache with the given expiry. If the key
// already exists, it will overwrite the previoous value
func (l *LFU) SetWithExpiry(key string, value interface{}, ttl time.Duration) error {
	return l.set(key, value, ttl)
}

// Get the content stored in the cache for the given key, and decode it into the value interface.
// return ErrKeyNotFound if the key is missing from the cache
func (l *LFU) Get(key string, value interface{}) error {
	return l.get(key, value)
}

// Remove deletes the value for a key.
func (l *LFU) Remove(key string) error {
	return l.remove(key)
}

// Keys returns a slice of the keys in the cache.
func (l *LFU) Keys() ([]string, error) {
	keys := []string{"foo", "bar"}
	return keys, nil
}

// Len returns the number of items in the cache.
func (l *LFU) Len() (int, error) {
	return l.len, nil
}

// GetInvalidateClusterEvent returns the cluster event configured when this cache was created.
func (l *LFU) GetInvalidateClusterEvent() string {
	return l.invalidateClusterEvent
}

// Name returns the name of the cache
func (l *LFU) Name() string {
	return l.name
}

func (l *LFU) set(key string, value interface{}, ttl time.Duration) error {
	var buf []byte
	var err error

	// We use a fast path for hot structs.
	if msgpVal, ok := value.(msgp.Marshaler); ok {
		buf, err = msgpVal.MarshalMsg(nil)
		if err != nil {
			return err
		}
	} else {
		// Slow path for other structs.
		buf, err = msgpack.Marshal(value)
		if err != nil {
			return err
		}
	}

	set := l.cache.SetWithTTL(key, buf, 1, ttl)

	if !set {
		return ErrKeyNotSet
	}

	if l.len < l.size {
		l.len++
	}

	return nil
}

func (l *LFU) get(key string, value interface{}) error {
	val, found := l.cache.Get(key)

	if !found {
		return ErrKeyNotFound
	}

	// We use a fast path for hot structs.
	if msgpVal, ok := value.(msgp.Unmarshaler); ok {
		_, err := msgpVal.UnmarshalMsg(val.([]byte))
		return err
	}

	// Slow path for other structs.
	return msgpack.Unmarshal(val.([]byte), value)
}

func (l *LFU) remove(key string) error {
	l.cache.Del(key)
	if l.len > 0 {
		l.len--
	}
	return nil
}

lfu_test.go

package cache

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const wait = 10 * time.Millisecond

// This test is passing
func TestLFU(t *testing.T) {
	l := NewLFU(&LFUOptions{
		Size:                   128,
		DefaultExpiry:          0,
		InvalidateClusterEvent: "",
	})

	len, _ := l.Len()
	assert.Equal(t, 0, len, "length should be 0 after cache is initialized.")

	err := l.Set("key", "value")
	require.Nil(t, err)

	time.Sleep(wait)

	var v string
	err = l.Get("key", &v)
	require.Nil(t, err, "should exist")
	require.Equalf(t, "value", v, "bad value: %v", v)

	len, err = l.Len()
	require.Nil(t, err)
	assert.Equal(t, 1, len, "length should be 1 after item is added to the cache.")

	err = l.Remove("key")
	err = l.Get("key", &v)
	require.EqualError(t, err, "key not found")

	_ = l.Purge()

	len, _ = l.Len()
	assert.Equal(t, 0, len, "length should be 0 after cache is purged.")
}

// This test is passing
func TestLFUExpire(t *testing.T) {
	l := NewLFU(&LFUOptions{
		Size:                   128,
		DefaultExpiry:          1 * time.Second,
		InvalidateClusterEvent: "",
	})

	l.SetWithDefaultExpiry("1", 1)
	l.SetWithExpiry("3", 3, 0*time.Second)

	time.Sleep(time.Second * 2)

	var r1 int
	err := l.Get("1", &r1)
	require.Equal(t, err, ErrKeyNotFound, "should not exist")

	var r2 int
	err2 := l.Get("3", &r2)
	require.Nil(t, err2, "should exist")
	require.Equal(t, 3, r2)
}

func TestLFUMarshalUnMarshal(t *testing.T) {
	l := NewLFU(&LFUOptions{
		Size:                   1,
		DefaultExpiry:          0,
		InvalidateClusterEvent: "",
	})

	value1 := map[string]interface{}{
		"key1": 1,
		"key2": "value2",
	}

	err := l.Set("test", value1)
	require.Nil(t, err)

	time.Sleep(wait)

	var value2 map[string]interface{}
	err = l.Get("test", &value2)
	require.Nil(t, err)

	v1, ok := value2["key1"].(int64)
	require.True(t, ok, "unable to cast value")
	assert.Equal(t, int64(1), v1)

	v2, ok := value2["key2"].(string)
	require.True(t, ok, "unable to cast value")
	assert.Equal(t, "value2", v2)

}

func BenchmarkLFU(b *testing.B) {

	value1 := "simplestring"

	b.Run("simple=new", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			l2 := NewLFU(&LFUOptions{
				Size:                   1,
				DefaultExpiry:          0,
				InvalidateClusterEvent: "",
			})
			err := l2.Set("test", value1)
			require.Nil(b, err)
			time.Sleep(wait)
			var val string
			err = l2.Get("test", &val)
			require.Nil(b, err)
		}
	})

	type obj struct {
		Field1 int
		Field2 string
		Field3 struct {
			Field4 int
			Field5 string
		}
		Field6 map[string]string
	}

	value2 := obj{
		1,
		"field2",
		struct {
			Field4 int
			Field5 string
		}{
			6,
			"field5 is a looooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong string",
		},
		map[string]string{
			"key0": "value0",
			"key1": "value value1",
			"key2": "value value value2",
			"key3": "value value value value3",
			"key4": "value value value value value4",
			"key5": "value value value value value value5",
			"key6": "value value value value value value value6",
			"key7": "value value value value value value value value7",
			"key8": "value value value value value value value value value8",
			"key9": "value value value value value value value value value value9",
		},
	}

	b.Run("complex=new", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			l2 := NewLFU(&LFUOptions{
				Size:                   1,
				DefaultExpiry:          0,
				InvalidateClusterEvent: "",
			})
			err := l2.Set("test", value2)
			require.Nil(b, err)
			time.Sleep(wait)
			var val obj
			err = l2.Get("test", &val)
			require.Nil(b, err)
		}
	})
}