*: Updating hashicorp LRU cache to v2 (thanos-io#7306)

* *: Updating hashicorp LRU cache to v2 Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com> * Adding some new comments regarding removing complexity of TTL Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com> * Using new version everywhere Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com> * rephrase the comment Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com> --------- Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com>
jnyi · Jun 3, 2024 · ee6963c · ee6963c
1 parent 0a3291f
commit ee6963c
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 28 deletions.
diff --git a/go.mod b/go.mod
@@ -36,7 +36,6 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware/providers/kit/v2 v2.0.0-20201002093600-73cf2ae9d891
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.0-rc.2.0.20201207153454-9f6bf00c00a7
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
-	github.com/hashicorp/golang-lru v0.6.0
 	github.com/jpillora/backoff v1.0.0
 	github.com/json-iterator/go v1.1.12
 	github.com/klauspost/compress v1.17.7
@@ -117,6 +116,7 @@ require (
 
 require (
 	github.com/cortexproject/promqlsmith v0.0.0-20240326071418-c2a9ca1e89f5
+	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/mitchellh/go-ps v1.0.0
 	github.com/onsi/gomega v1.29.0
 	github.com/prometheus-community/prom-label-proxy v0.8.1-0.20240127162815-c1195f9aabc0

diff --git a/go.sum b/go.sum
@@ -1153,6 +1153,8 @@ github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.6.0 h1:uL2shRDx7RTrOrTCUZEGP/wJUFiUI8QT6E7z5o8jga4=
 github.com/hashicorp/golang-lru v0.6.0/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hashicorp/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
 github.com/hashicorp/mdns v1.0.0/go.mod h1:tL+uN++7HEJ6SQLQ2/p+z2pH24WQKWjBPkE0mNTz8vQ=
 github.com/hashicorp/memberlist v0.1.3/go.mod h1:ajVTdAv/9Im8oMAAj5G31PhhMCZJV2pPBoIllUwCN7I=

diff --git a/pkg/cache/inmemory.go b/pkg/cache/inmemory.go
@@ -10,7 +10,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	lru "github.com/hashicorp/golang-lru/simplelru"
+	lru "github.com/hashicorp/golang-lru/v2/simplelru"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -46,7 +46,7 @@ type InMemoryCache struct {
 
 	mtx         sync.Mutex
 	curSize     uint64
-	lru         *lru.LRU
+	lru         *lru.LRU[string, cacheDataWithTTLWrapper]
 	evicted     prometheus.Counter
 	requests    prometheus.Counter
 	hits        prometheus.Counter
@@ -62,11 +62,12 @@ type InMemoryCache struct {
 
 type cacheDataWithTTLWrapper struct {
 	data []byte
-	// The objects that are over the TTL are not destroyed eagerly.
-	// When there is a hit for an item that is over the TTL, the object is removed from the cache
-	// and null is returned.
-	// There is ongoing effort to integrate TTL within the Hashicorp golang cache itself.
-	// This https://github.com/hashicorp/golang-lru/pull/41 can be used here once complete.
+	// Items exceeding their Time-To-Live (TTL) are not immediately removed from the cache.
+	// Instead, when an access attempt is made for an item past its TTL, the item is evicted from the cache, and a null value is returned.
+	// Efforts are underway to incorporate TTL directly into the Hashicorp golang cache.
+	// Although this pull request (https://github.com/hashicorp/golang-lru/pull/41) has been completed, it's challenging to apply here due to the following reasons:
+	// The Hashicorp LRU API requires setting the TTL during the constructor phase, whereas in Thanos, we set the TTL for each Set()/Store() operation.
+	// Refer to this link for more details: https://github.com/thanos-io/thanos/blob/23d205286436291fa0c55c25c392ee08f42d5fbf/pkg/store/cache/caching_bucket.go#L167-L175
 	expiryTime time.Time
 }
 
@@ -176,7 +177,7 @@ func NewInMemoryCacheWithConfig(name string, logger log.Logger, reg prometheus.R
 
 	// Initialize LRU cache with a high size limit since we will manage evictions ourselves
 	// based on stored size using `RemoveOldest` method.
-	l, err := lru.NewLRU(maxInt, c.onEvict)
+	l, err := lru.NewLRU[string, cacheDataWithTTLWrapper](maxInt, c.onEvict)
 	if err != nil {
 		return nil, err
 	}
@@ -191,9 +192,9 @@ func NewInMemoryCacheWithConfig(name string, logger log.Logger, reg prometheus.R
 	return c, nil
 }
 
-func (c *InMemoryCache) onEvict(key, val interface{}) {
-	keySize := uint64(len(key.(string)))
-	entrySize := uint64(len(val.(cacheDataWithTTLWrapper).data))
+func (c *InMemoryCache) onEvict(key string, val cacheDataWithTTLWrapper) {
+	keySize := uint64(len(key))
+	entrySize := uint64(len(val.data))
 
 	c.evicted.Inc()
 	c.current.Dec()
@@ -214,13 +215,13 @@ func (c *InMemoryCache) get(key string) ([]byte, bool) {
 	}
 	// If the present time is greater than the TTL for the object from cache, the object will be
 	// removed from the cache and a nil will be returned
-	if time.Now().After(v.(cacheDataWithTTLWrapper).expiryTime) {
+	if time.Now().After(v.expiryTime) {
 		c.hitsExpired.Inc()
 		c.lru.Remove(key)
 		return nil, false
 	}
 	c.hits.Inc()
-	return v.(cacheDataWithTTLWrapper).data, true
+	return v.data, true
 }
 
 func (c *InMemoryCache) set(key string, val []byte, ttl time.Duration) {

diff --git a/pkg/querysharding/analyzer.go b/pkg/querysharding/analyzer.go
@@ -19,7 +19,7 @@ package querysharding
 import (
 	"fmt"
 
-	lru "github.com/hashicorp/golang-lru"
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/promql/parser"
 )
@@ -38,14 +38,14 @@ type QueryAnalyzer struct{}
 
 type CachedQueryAnalyzer struct {
 	analyzer *QueryAnalyzer
-	cache    *lru.Cache
+	cache    *lru.Cache[string, cachedValue]
 }
 
 // NewQueryAnalyzer creates a new QueryAnalyzer.
 func NewQueryAnalyzer() *CachedQueryAnalyzer {
 	// Ignore the error check since it throws error
 	// only if size is <= 0.
-	cache, _ := lru.New(256)
+	cache, _ := lru.New[string, cachedValue](256)
 	return &CachedQueryAnalyzer{
 		analyzer: &QueryAnalyzer{},
 		cache:    cache,
@@ -61,7 +61,7 @@ func (a *CachedQueryAnalyzer) Analyze(query string) (QueryAnalysis, error) {
 	if a.cache.Contains(query) {
 		value, ok := a.cache.Get(query)
 		if ok {
-			return value.(cachedValue).QueryAnalysis, value.(cachedValue).err
+			return value.QueryAnalysis, value.err
 		}
 	}
 

diff --git a/pkg/store/cache/inmemory.go b/pkg/store/cache/inmemory.go
@@ -10,7 +10,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	lru "github.com/hashicorp/golang-lru/simplelru"
+	lru "github.com/hashicorp/golang-lru/v2/simplelru"
 	"github.com/oklog/ulid"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
@@ -36,7 +36,7 @@ type InMemoryIndexCache struct {
 	mtx sync.Mutex
 
 	logger           log.Logger
-	lru              *lru.LRU
+	lru              *lru.LRU[CacheKey, []byte]
 	maxSizeBytes     uint64
 	maxItemSizeBytes uint64
 
@@ -170,7 +170,7 @@ func NewInMemoryIndexCacheWithConfig(logger log.Logger, commonMetrics *CommonMet
 
 	// Initialize LRU cache with a high size limit since we will manage evictions ourselves
 	// based on stored size using `RemoveOldest` method.
-	l, err := lru.NewLRU(maxInt, c.onEvict)
+	l, err := lru.NewLRU[CacheKey, []byte](maxInt, c.onEvict)
 	if err != nil {
 		return nil, err
 	}
@@ -185,14 +185,14 @@ func NewInMemoryIndexCacheWithConfig(logger log.Logger, commonMetrics *CommonMet
 	return c, nil
 }
 
-func (c *InMemoryIndexCache) onEvict(key, val interface{}) {
-	k := key.(CacheKey).KeyType()
-	entrySize := sliceHeaderSize + uint64(len(val.([]byte)))
+func (c *InMemoryIndexCache) onEvict(key CacheKey, val []byte) {
+	k := key.KeyType()
+	entrySize := sliceHeaderSize + uint64(len(val))
 
 	c.evicted.WithLabelValues(k).Inc()
 	c.current.WithLabelValues(k).Dec()
 	c.currentSize.WithLabelValues(k).Sub(float64(entrySize))
-	c.totalCurrentSize.WithLabelValues(k).Sub(float64(entrySize + key.(CacheKey).Size()))
+	c.totalCurrentSize.WithLabelValues(k).Sub(float64(entrySize + key.Size()))
 
 	c.curSize -= entrySize
 }
@@ -205,7 +205,7 @@ func (c *InMemoryIndexCache) get(key CacheKey) ([]byte, bool) {
 	if !ok {
 		return nil, false
 	}
-	return v.([]byte), true
+	return v, true
 }
 
 func (c *InMemoryIndexCache) set(typ string, key CacheKey, val []byte) {

diff --git a/pkg/store/cache/inmemory_test.go b/pkg/store/cache/inmemory_test.go
@@ -12,7 +12,7 @@ import (
 	"testing"
 
 	"github.com/go-kit/log"
-	"github.com/hashicorp/golang-lru/simplelru"
+	"github.com/hashicorp/golang-lru/v2/simplelru"
 	"github.com/oklog/ulid"
 	"github.com/prometheus/client_golang/prometheus"
 	promtest "github.com/prometheus/client_golang/prometheus/testutil"
@@ -72,7 +72,7 @@ func TestInMemoryIndexCache_AvoidsDeadlock(t *testing.T) {
 	})
 	testutil.Ok(t, err)
 
-	l, err := simplelru.NewLRU(math.MaxInt64, func(key, val interface{}) {
+	l, err := simplelru.NewLRU[CacheKey, []byte](math.MaxInt64, func(key CacheKey, val []byte) {
 		// Hack LRU to simulate broken accounting: evictions do not reduce current size.
 		size := cache.curSize
 		cache.onEvict(key, val)