Try Ristretto cache

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/Masterminds/semver/v3 v3.1.1
 	github.com/Microsoft/go-winio v0.5.1 // indirect
 	github.com/codeclysm/extract v2.2.0+incompatible
+	github.com/dgraph-io/ristretto v0.1.0 // indirect
 	github.com/docker/distribution v2.7.1+incompatible // indirect
 	github.com/docker/docker v20.10.16+incompatible
 	github.com/docker/go-connections v0.4.0

go.sum

@@ -126,6 +126,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/badger v1.6.0/go.mod h1:zwt7syl517jmP8s94KqSxTlM6IMsdhYy6psNgSztDR4=
+github.com/dgraph-io/ristretto v0.1.0 h1:Jv3CGQHp9OjuMBSne1485aDpUkTKEcUqF+jm/LuerPI=
+github.com/dgraph-io/ristretto v0.1.0/go.mod h1:fux0lOrBhrVCJd3lcTHsIJhq1T2rokOu6v9Vcb3Q9ug=
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/docker/distribution v2.7.1+incompatible h1:a5mlkVzth6W5A4fOsS3D2EO5BUmsJpcB+cRlLU7cSug=
 github.com/docker/distribution v2.7.1+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w=
@@ -137,6 +139,7 @@ github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw
 github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
 github.com/dustin/go-humanize v0.0.0-20180421182945-02af3965c54e/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
+github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/eknkc/amber v0.0.0-20171010120322-cdade1c07385/go.mod h1:0vRUJqYpeSZifjYj7uP3BG/gKcuzL9xWVV/Y+cK33KM=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
@@ -204,6 +207,7 @@ github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
 github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k=
+github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/groupcache v0.0.0-20191227052852-215e87163ea7/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=

internal/caching/impl_ristretto.go

@@ -0,0 +1,153 @@
+package caching
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/dgraph-io/ristretto"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+func NewRistrettoCache(enablePrometheus bool) (*Caches, error) {
+	roomVersions, err := NewRistrettoCachePartition(
+		RoomVersionCacheName,
+		RoomVersionCacheMutable,
+		RoomVersionCacheMaxEntries,
+		RoomVersionCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	serverKeys, err := NewRistrettoCachePartition(
+		ServerKeyCacheName,
+		ServerKeyCacheMutable,
+		ServerKeyCacheMaxEntries,
+		ServerKeyCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	roomServerRoomIDs, err := NewRistrettoCachePartition(
+		RoomServerRoomIDsCacheName,
+		RoomServerRoomIDsCacheMutable,
+		RoomServerRoomIDsCacheMaxEntries,
+		RoomServerRoomIDsCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	roomInfos, err := NewRistrettoCachePartition(
+		RoomInfoCacheName,
+		RoomInfoCacheMutable,
+		RoomInfoCacheMaxEntries,
+		RoomInfoCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	federationEvents, err := NewRistrettoCachePartition(
+		FederationEventCacheName,
+		FederationEventCacheMutable,
+		FederationEventCacheMaxEntries,
+		FederationEventCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	spaceRooms, err := NewRistrettoCachePartition(
+		SpaceSummaryRoomsCacheName,
+		SpaceSummaryRoomsCacheMutable,
+		SpaceSummaryRoomsCacheMaxEntries,
+		SpaceSummaryRoomsCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	lazyLoadCache, err := NewRistrettoCachePartition(
+		LazyLoadCacheName,
+		LazyLoadCacheMutable,
+		LazyLoadCacheMaxEntries,
+		LazyLoadCacheMaxAge,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Caches{
+		RoomVersions:      roomVersions,
+		ServerKeys:        serverKeys,
+		RoomServerRoomIDs: roomServerRoomIDs,
+		RoomInfos:         roomInfos,
+		FederationEvents:  federationEvents,
+		SpaceSummaryRooms: spaceRooms,
+		LazyLoading:       lazyLoadCache,
+	}, nil
+}
+
+type RistrettoCachePartition struct {
+	name       string
+	mutable    bool
+	maxEntries int
+	maxAge     time.Duration
+	ristretto  *ristretto.Cache
+}
+
+func NewRistrettoCachePartition(name string, mutable bool, maxEntries int, maxAge time.Duration, enablePrometheus bool) (*RistrettoCachePartition, error) {
+	var err error
+	cache := RistrettoCachePartition{
+		name:       name,
+		mutable:    mutable,
+		maxEntries: maxEntries,
+		maxAge:     maxAge,
+	}
+	cache.ristretto, err = ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e7,
+		MaxCost:     1 << 27,
+		BufferItems: 64,
+		Metrics:     true,
+	})
+	if err != nil {
+		return nil, err
+	}
+	if enablePrometheus {
+		promauto.NewGaugeFunc(prometheus.GaugeOpts{
+			Namespace: "dendrite",
+			Subsystem: "caching_ristretto",
+			Name:      name,
+		}, func() float64 {
+			return float64(cache.ristretto.Metrics.Ratio())
+		})
+	}
+	return &cache, nil
+}
+
+func (c *RistrettoCachePartition) Set(key string, value interface{}) {
+	if !c.mutable {
+		if peek, ok := c.ristretto.Get(key); ok {
+			if entry, ok := peek.(*inMemoryLRUCacheEntry); ok && entry.value != value {
+				panic(fmt.Sprintf("invalid use of immutable cache tries to mutate existing value of %q", key))
+			}
+		}
+	}
+	c.ristretto.SetWithTTL(key, value, 1, c.maxAge)
+}
+
+func (c *RistrettoCachePartition) Unset(key string) {
+	if !c.mutable {
+		panic(fmt.Sprintf("invalid use of immutable cache tries to unset value of %q", key))
+	}
+	c.ristretto.Del(key)
+}
+
+func (c *RistrettoCachePartition) Get(key string) (value interface{}, ok bool) {
+	return c.ristretto.Get(key)
+}

setup/base/base.go

@@ -161,7 +161,7 @@ func NewBaseDendrite(cfg *config.Dendrite, componentName string, options ...Base
 		}
 	}
 
-	cache, err := caching.NewInMemoryLRUCache(enableMetrics)
+	cache, err := caching.NewRistrettoCache(enableMetrics)
 	if err != nil {
 		logrus.WithError(err).Warnf("Failed to create cache")
 	}