Ristretto cache (#2563)

* Try Ristretto cache

* Tweak

* It's beautiful

* Update GMSL

* More strict keyable interface

* Fix that some more

* Make less panicky

* Don't enforce mutability checks for now

* Determine mutability using deep equality

* Tweaks

* Namespace keys

* Make federation caches mutable

* Update cost estimation, add metric

* Update GMSL

* Estimate cost for metrics better

* Reduce counters a bit

* Try caching events

* Some guards

* Try again

* Try this

* Use separate caches for hopefully better hash distribution

* Fix bug with admitting events into cache

* Try to fix bugs

* Check nil

* Try that again

* Preserve order jeezo this is messy

* thanks VS Code for doing exactly the wrong thing

* Try this again

* Be more specific

* aaaaargh

* One more time

* That might be better

* Stronger sorting

* Cache expiries, async publishing of EDUs

* Put it back

* Use a shared cache again

* Cost estimation fixes

* Update ristretto

* Reduce counters a bit

* Clean up a bit

* Update GMSL

* 1GB

* Configurable cache sizees

* Tweaks

* Add `config.DataUnit` for specifying friendly cache sizes

* Various tweaks

* Update GMSL

* Add back some lazy loading caching

* Include key in cost

* Include key in cost

* Tweak max age handling, config key name

* Only register prometheus metrics if requested

* Review comments @S7evinK

* Don't return errors when creating caches (it is better just to crash since otherwise we'll `nil`-pointer exception everywhere)

* Review comments

* Update sample configs

* Update GHA Workflow

* Update Complement images to Go 1.18

* Remove the cache test from the federation API as we no longer guarantee immediate cache admission

* Don't check the caches in the renewal test

* Possibly fix the upgrade tests

* Update to matrix-org/gomatrixserverlib#322

* Update documentation to refer to Go 1.18

1 files changed, 200 insertions, 0 deletions

diff --git a/internal/caching/impl_ristretto.go b/internal/caching/impl_ristretto.go
new file mode 100644
index 00000000..6d625b55
--- /dev/null
+++ b/internal/caching/impl_ristretto.go
@@ -0,0 +1,200 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package caching
+
+import (
+	"fmt"
+	"reflect"
+	"time"
+	"unsafe"
+
+	"github.com/dgraph-io/ristretto"
+	"github.com/dgraph-io/ristretto/z"
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/dendrite/setup/config"
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+const (
+	roomVersionsCache byte = iota + 1
+	serverKeysCache
+	roomNIDsCache
+	roomIDsCache
+	roomEventsCache
+	roomInfosCache
+	federationPDUsCache
+	federationEDUsCache
+	spaceSummaryRoomsCache
+	lazyLoadingCache
+)
+
+func NewRistrettoCache(maxCost config.DataUnit, maxAge time.Duration, enablePrometheus bool) *Caches {
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e5, // 10x number of expected cache items, affects bloom filter size, gives us room for 10,000 currently
+		BufferItems: 64,  // recommended by the ristretto godocs as a sane buffer size value
+		MaxCost:     int64(maxCost),
+		Metrics:     true,
+		KeyToHash: func(key interface{}) (uint64, uint64) {
+			return z.KeyToHash(key)
+		},
+	})
+	if err != nil {
+		panic(err)
+	}
+	if enablePrometheus {
+		promauto.NewGaugeFunc(prometheus.GaugeOpts{
+			Namespace: "dendrite",
+			Subsystem: "caching_ristretto",
+			Name:      "ratio",
+		}, func() float64 {
+			return float64(cache.Metrics.Ratio())
+		})
+		promauto.NewGaugeFunc(prometheus.GaugeOpts{
+			Namespace: "dendrite",
+			Subsystem: "caching_ristretto",
+			Name:      "cost",
+		}, func() float64 {
+			return float64(cache.Metrics.CostAdded() - cache.Metrics.CostEvicted())
+		})
+	}
+	return &Caches{
+		RoomVersions: &RistrettoCachePartition[string, gomatrixserverlib.RoomVersion]{ // room ID -> room version
+			cache:  cache,
+			Prefix: roomVersionsCache,
+			MaxAge: maxAge,
+		},
+		ServerKeys: &RistrettoCachePartition[string, gomatrixserverlib.PublicKeyLookupResult]{ // server name -> server keys
+			cache:   cache,
+			Prefix:  serverKeysCache,
+			Mutable: true,
+			MaxAge:  maxAge,
+		},
+		RoomServerRoomNIDs: &RistrettoCachePartition[string, types.RoomNID]{ // room ID -> room NID
+			cache:  cache,
+			Prefix: roomNIDsCache,
+			MaxAge: maxAge,
+		},
+		RoomServerRoomIDs: &RistrettoCachePartition[int64, string]{ // room NID -> room ID
+			cache:  cache,
+			Prefix: roomIDsCache,
+			MaxAge: maxAge,
+		},
+		RoomServerEvents: &RistrettoCostedCachePartition[int64, *gomatrixserverlib.Event]{ // event NID -> event
+			&RistrettoCachePartition[int64, *gomatrixserverlib.Event]{
+				cache:  cache,
+				Prefix: roomEventsCache,
+				MaxAge: maxAge,
+			},
+		},
+		RoomInfos: &RistrettoCachePartition[string, types.RoomInfo]{ // room ID -> room info
+			cache:   cache,
+			Prefix:  roomInfosCache,
+			Mutable: true,
+			MaxAge:  maxAge,
+		},
+		FederationPDUs: &RistrettoCostedCachePartition[int64, *gomatrixserverlib.HeaderedEvent]{ // queue NID -> PDU
+			&RistrettoCachePartition[int64, *gomatrixserverlib.HeaderedEvent]{
+				cache:   cache,
+				Prefix:  federationPDUsCache,
+				Mutable: true,
+				MaxAge:  lesserOf(time.Hour/2, maxAge),
+			},
+		},
+		FederationEDUs: &RistrettoCostedCachePartition[int64, *gomatrixserverlib.EDU]{ // queue NID -> EDU
+			&RistrettoCachePartition[int64, *gomatrixserverlib.EDU]{
+				cache:   cache,
+				Prefix:  federationEDUsCache,
+				Mutable: true,
+				MaxAge:  lesserOf(time.Hour/2, maxAge),
+			},
+		},
+		SpaceSummaryRooms: &RistrettoCachePartition[string, gomatrixserverlib.MSC2946SpacesResponse]{ // room ID -> space response
+			cache:   cache,
+			Prefix:  spaceSummaryRoomsCache,
+			Mutable: true,
+			MaxAge:  maxAge,
+		},
+		LazyLoading: &RistrettoCachePartition[lazyLoadingCacheKey, string]{ // composite key -> event ID
+			cache:   cache,
+			Prefix:  lazyLoadingCache,
+			Mutable: true,
+			MaxAge:  maxAge,
+		},
+	}
+}
+
+type RistrettoCostedCachePartition[k keyable, v costable] struct {
+	*RistrettoCachePartition[k, v]
+}
+
+func (c *RistrettoCostedCachePartition[K, V]) Set(key K, value V) {
+	cost := value.CacheCost()
+	c.setWithCost(key, value, int64(cost))
+}
+
+type RistrettoCachePartition[K keyable, V any] struct {
+	cache   *ristretto.Cache
+	Prefix  byte
+	Mutable bool
+	MaxAge  time.Duration
+}
+
+func (c *RistrettoCachePartition[K, V]) setWithCost(key K, value V, cost int64) {
+	bkey := fmt.Sprintf("%c%v", c.Prefix, key)
+	if !c.Mutable {
+		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))
+		}
+	}
+	c.cache.SetWithTTL(bkey, value, int64(len(bkey))+cost, c.MaxAge)
+}
+
+func (c *RistrettoCachePartition[K, V]) Set(key K, value V) {
+	var cost int64
+	if cv, ok := any(value).(string); ok {
+		cost = int64(len(cv))
+	} else {
+		cost = int64(unsafe.Sizeof(value))
+	}
+	c.setWithCost(key, value, cost)
+}
+
+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(bkey)
+}
+
+func (c *RistrettoCachePartition[K, V]) Get(key K) (value V, ok bool) {
+	bkey := fmt.Sprintf("%c%v", c.Prefix, key)
+	v, ok := c.cache.Get(bkey)
+	if !ok || v == nil {
+		var empty V
+		return empty, false
+	}
+	value, ok = v.(V)
+	return
+}
+
+func lesserOf(a, b time.Duration) time.Duration {
+	if a < b {
+		return a
+	}
+	return b
+}