From e0466c77774d396dee9a97fcdd921e5aa8e297ab Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 15 Jun 2022 14:54:13 +0100
Subject: [PATCH] Use a shared cache again

---
 internal/caching/impl_ristretto.go | 63 ++++++++++++++++++++++--------
 1 file changed, 46 insertions(+), 17 deletions(-)

diff --git a/internal/caching/impl_ristretto.go b/internal/caching/impl_ristretto.go
index 11e871fb2..f4f71f751 100644
--- a/internal/caching/impl_ristretto.go
+++ b/internal/caching/impl_ristretto.go
@@ -7,18 +7,22 @@ import (
 	"unsafe"
 
 	"github.com/dgraph-io/ristretto"
+	"github.com/dgraph-io/ristretto/z"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
-func MustCreateCache(name string, maxCost CacheSize, enablePrometheus bool) *ristretto.Cache {
+func MustCreateCache(maxCost CacheSize, enablePrometheus bool) *ristretto.Cache {
 	cache, err := ristretto.NewCache(&ristretto.Config{
 		NumCounters: 1e6,
 		MaxCost:     int64(maxCost),
 		BufferItems: 64,
 		Metrics:     true,
+		KeyToHash: func(key interface{}) (uint64, uint64) {
+			return z.KeyToHash(key)
+		},
 	})
 	if err != nil {
 		panic(err)
@@ -26,62 +30,83 @@ func MustCreateCache(name string, maxCost CacheSize, enablePrometheus bool) *ris
 	promauto.NewGaugeFunc(prometheus.GaugeOpts{
 		Namespace: "dendrite",
 		Subsystem: "caching_ristretto",
-		Name:      name + "_ratio",
+		Name:      "ratio",
 	}, func() float64 {
 		return float64(cache.Metrics.Ratio())
 	})
 	promauto.NewGaugeFunc(prometheus.GaugeOpts{
 		Namespace: "dendrite",
 		Subsystem: "caching_ristretto",
-		Name:      name + "_cost",
+		Name:      "cost",
 	}, func() float64 {
 		return float64(cache.Metrics.CostAdded() - cache.Metrics.CostEvicted())
 	})
 	return cache
 }
 
-func NewRistrettoCache(maxCost CacheSize, enablePrometheus bool) (*Caches, error) {
+const (
+	roomVersionsCache byte = iota + 1
+	serverKeysCache
+	roomIDsCache
+	roomEventsCache
+	roomInfosCache
+	federationPDUsCache
+	federationEDUsCache
+	spaceSummaryRoomsCache
+	lazyLoadingCache
+)
 
+func NewRistrettoCache(maxCost CacheSize, enablePrometheus bool) (*Caches, error) {
+	cache := MustCreateCache(maxCost, enablePrometheus)
 	return &Caches{
 		RoomVersions: &RistrettoCachePartition[string, gomatrixserverlib.RoomVersion]{
-			cache:  MustCreateCache("room_versions", 1*MB, enablePrometheus),
+			cache:  cache,
+			Prefix: roomVersionsCache,
 			MaxAge: time.Hour,
 		},
 		ServerKeys: &RistrettoCachePartition[string, gomatrixserverlib.PublicKeyLookupResult]{
-			cache:   MustCreateCache("server_keys", 32*MB, enablePrometheus),
+			cache:   cache,
+			Prefix:  serverKeysCache,
 			Mutable: true,
 			MaxAge:  time.Hour,
 		},
 		RoomServerRoomIDs: &RistrettoCachePartition[int64, string]{
-			cache:  MustCreateCache("room_ids", 1*MB, enablePrometheus),
+			cache:  cache,
+			Prefix: roomIDsCache,
 			MaxAge: time.Hour,
 		},
 		RoomServerEvents: &RistrettoCachePartition[int64, *gomatrixserverlib.Event]{
-			cache:  MustCreateCache("room_events", 1*GB, enablePrometheus),
+			cache:  cache,
+			Prefix: roomEventsCache,
 			MaxAge: time.Hour,
 		},
 		RoomInfos: &RistrettoCachePartition[string, types.RoomInfo]{
-			cache:   MustCreateCache("room_infos", 16*MB, enablePrometheus),
+			cache:   cache,
+			Prefix:  roomInfosCache,
 			Mutable: true,
 			MaxAge:  time.Hour,
 		},
 		FederationPDUs: &RistrettoCachePartition[int64, *gomatrixserverlib.HeaderedEvent]{
-			cache:   MustCreateCache("federation_pdus", 128*MB, enablePrometheus),
+			cache:   cache,
+			Prefix:  federationPDUsCache,
 			Mutable: true,
 			MaxAge:  time.Hour / 2,
 		},
 		FederationEDUs: &RistrettoCachePartition[int64, *gomatrixserverlib.EDU]{
-			cache:   MustCreateCache("federation_edus", 128*MB, enablePrometheus),
+			cache:   cache,
+			Prefix:  federationEDUsCache,
 			Mutable: true,
 			MaxAge:  time.Hour / 2,
 		},
 		SpaceSummaryRooms: &RistrettoCachePartition[string, gomatrixserverlib.MSC2946SpacesResponse]{
-			cache:   MustCreateCache("space_summary_rooms", 128, enablePrometheus), // TODO: not costed
+			cache:   cache,
+			Prefix:  spaceSummaryRoomsCache,
 			Mutable: true,
 			MaxAge:  time.Hour,
 		},
 		LazyLoading: &RistrettoCachePartition[string, any]{ // TODO: type
-			cache:   MustCreateCache("lazy_loading", 256, enablePrometheus), // TODO: not costed
+			cache:   cache,
+			Prefix:  lazyLoadingCache,
 			Mutable: true,
 			MaxAge:  time.Hour,
 		},
@@ -90,13 +115,15 @@ func NewRistrettoCache(maxCost CacheSize, enablePrometheus bool) (*Caches, error
 
 type RistrettoCachePartition[K keyable, V any] struct {
 	cache   *ristretto.Cache
+	Prefix  byte
 	Mutable bool
 	MaxAge  time.Duration
 }
 
 func (c *RistrettoCachePartition[K, V]) Set(key K, value V) {
+	bkey := fmt.Sprintf("%c%v", c.Prefix, key)
 	if !c.Mutable {
-		if v, ok := c.cache.Get(key); ok && v != nil && !reflect.DeepEqual(v, value) {
+		if v, ok := c.cache.Get(bkey); ok && v != nil && !reflect.DeepEqual(v, value) {
 			panic(fmt.Sprintf("invalid use of immutable cache tries to change value of %v from %v to %v", key, v, value))
 		}
 	}
@@ -108,18 +135,20 @@ func (c *RistrettoCachePartition[K, V]) Set(key K, value V) {
 	} else {
 		cost = int64(unsafe.Sizeof(value))
 	}
-	c.cache.SetWithTTL(key, value, cost, c.MaxAge)
+	c.cache.SetWithTTL(bkey, value, cost, c.MaxAge)
 }
 
 func (c *RistrettoCachePartition[K, V]) Unset(key K) {
+	bkey := fmt.Sprintf("%c%v", c.Prefix, key)
 	if !c.Mutable {
 		panic(fmt.Sprintf("invalid use of immutable cache tries to unset value of %v", key))
 	}
-	c.cache.Del(key)
+	c.cache.Del(bkey)
 }
 
 func (c *RistrettoCachePartition[K, V]) Get(key K) (value V, ok bool) {
-	v, ok := c.cache.Get(key)
+	bkey := fmt.Sprintf("%c%v", c.Prefix, key)
+	v, ok := c.cache.Get(bkey)
 	if !ok || v == nil {
 		var empty V
 		return empty, false