Add event state key cache (#2576)

5 files changed, 58 insertions, 15 deletions

diff --git a/internal/caching/cache_eventstatekeys.go b/internal/caching/cache_eventstatekeys.go
new file mode 100644
index 00000000..05580ab0
--- /dev/null
+++ b/internal/caching/cache_eventstatekeys.go
@@ -0,0 +1,18 @@
+package caching
+
+import "github.com/matrix-org/dendrite/roomserver/types"
+
+// EventStateKeyCache contains the subset of functions needed for
+// a room event state key cache.
+type EventStateKeyCache interface {
+	GetEventStateKey(eventStateKeyNID types.EventStateKeyNID) (string, bool)
+	StoreEventStateKey(eventStateKeyNID types.EventStateKeyNID, eventStateKey string)
+}
+
+func (c Caches) GetEventStateKey(eventStateKeyNID types.EventStateKeyNID) (string, bool) {
+	return c.RoomServerStateKeys.Get(eventStateKeyNID)
+}
+
+func (c Caches) StoreEventStateKey(eventStateKeyNID types.EventStateKeyNID, eventStateKey string) {
+	c.RoomServerStateKeys.Set(eventStateKeyNID, eventStateKey)
+}
diff --git a/internal/caching/cache_roomservernids.go b/internal/caching/cache_roomservernids.go
index b409aeef..f27154f1 100644
--- a/internal/caching/cache_roomservernids.go
+++ b/internal/caching/cache_roomservernids.go
@@ -9,6 +9,7 @@ type RoomServerCaches interface {
 	RoomVersionCache
 	RoomInfoCache
 	RoomServerEventsCache
+	EventStateKeyCache
 }
 
 // RoomServerNIDsCache contains the subset of functions needed for
@@ -19,9 +20,9 @@ type RoomServerNIDsCache interface {
 }
 
 func (c Caches) GetRoomServerRoomID(roomNID types.RoomNID) (string, bool) {
-	return c.RoomServerRoomIDs.Get(int64(roomNID))
+	return c.RoomServerRoomIDs.Get(roomNID)
 }
 
 func (c Caches) StoreRoomServerRoomID(roomNID types.RoomNID, roomID string) {
-	c.RoomServerRoomIDs.Set(int64(roomNID), roomID)
+	c.RoomServerRoomIDs.Set(roomNID, roomID)
 }
diff --git a/internal/caching/caches.go b/internal/caching/caches.go
index e7914ce7..f13f743d 100644
--- a/internal/caching/caches.go
+++ b/internal/caching/caches.go
@@ -23,16 +23,17 @@ import (
 // different implementations as long as they satisfy the Cache
 // interface.
 type Caches struct {
-	RoomVersions       Cache[string, gomatrixserverlib.RoomVersion]           // room ID -> room version
-	ServerKeys         Cache[string, gomatrixserverlib.PublicKeyLookupResult] // server name -> server keys
-	RoomServerRoomNIDs Cache[string, types.RoomNID]                           // room ID -> room NID
-	RoomServerRoomIDs  Cache[int64, string]                                   // room NID -> room ID
-	RoomServerEvents   Cache[int64, *gomatrixserverlib.Event]                 // event NID -> event
-	RoomInfos          Cache[string, *types.RoomInfo]                         // room ID -> room info
-	FederationPDUs     Cache[int64, *gomatrixserverlib.HeaderedEvent]         // queue NID -> PDU
-	FederationEDUs     Cache[int64, *gomatrixserverlib.EDU]                   // queue NID -> EDU
-	SpaceSummaryRooms  Cache[string, gomatrixserverlib.MSC2946SpacesResponse] // room ID -> space response
-	LazyLoading        Cache[lazyLoadingCacheKey, string]                     // composite key -> event ID
+	RoomVersions        Cache[string, gomatrixserverlib.RoomVersion]           // room ID -> room version
+	ServerKeys          Cache[string, gomatrixserverlib.PublicKeyLookupResult] // server name -> server keys
+	RoomServerRoomNIDs  Cache[string, types.RoomNID]                           // room ID -> room NID
+	RoomServerRoomIDs   Cache[types.RoomNID, string]                           // room NID -> room ID
+	RoomServerEvents    Cache[int64, *gomatrixserverlib.Event]                 // event NID -> event
+	RoomServerStateKeys Cache[types.EventStateKeyNID, string]                  // event NID -> event state key
+	RoomInfos           Cache[string, *types.RoomInfo]                         // room ID -> room info
+	FederationPDUs      Cache[int64, *gomatrixserverlib.HeaderedEvent]         // queue NID -> PDU
+	FederationEDUs      Cache[int64, *gomatrixserverlib.EDU]                   // queue NID -> EDU
+	SpaceSummaryRooms   Cache[string, gomatrixserverlib.MSC2946SpacesResponse] // room ID -> space response
+	LazyLoading         Cache[lazyLoadingCacheKey, string]                     // composite key -> event ID
 }
 
 // Cache is the interface that an implementation must satisfy.
@@ -44,7 +45,7 @@ type Cache[K keyable, T any] interface {
 
 type keyable interface {
 	// from https://github.com/dgraph-io/ristretto/blob/8e850b710d6df0383c375ec6a7beae4ce48fc8d5/z/z.go#L34
-	uint64 | string | []byte | byte | int | int32 | uint32 | int64 | lazyLoadingCacheKey
+	~uint64 | ~string | []byte | byte | ~int | ~int32 | ~uint32 | ~int64 | lazyLoadingCacheKey
 }
 
 type costable interface {
diff --git a/internal/caching/impl_ristretto.go b/internal/caching/impl_ristretto.go
index 01c97ad5..fdbbb4d7 100644
--- a/internal/caching/impl_ristretto.go
+++ b/internal/caching/impl_ristretto.go
@@ -40,6 +40,7 @@ const (
 	federationEDUsCache
 	spaceSummaryRoomsCache
 	lazyLoadingCache
+	eventStateKeyCache
 )
 
 func NewRistrettoCache(maxCost config.DataUnit, maxAge time.Duration, enablePrometheus bool) *Caches {
@@ -88,7 +89,7 @@ func NewRistrettoCache(maxCost config.DataUnit, maxAge time.Duration, enableProm
 			Prefix: roomNIDsCache,
 			MaxAge: maxAge,
 		},
-		RoomServerRoomIDs: &RistrettoCachePartition[int64, string]{ // room NID -> room ID
+		RoomServerRoomIDs: &RistrettoCachePartition[types.RoomNID, string]{ // room NID -> room ID
 			cache:  cache,
 			Prefix: roomIDsCache,
 			MaxAge: maxAge,
@@ -100,6 +101,11 @@ func NewRistrettoCache(maxCost config.DataUnit, maxAge time.Duration, enableProm
 				MaxAge: maxAge,
 			},
 		},
+		RoomServerStateKeys: &RistrettoCachePartition[types.EventStateKeyNID, string]{ // event NID -> event state key
+			cache:  cache,
+			Prefix: eventStateKeyCache,
+			MaxAge: maxAge,
+		},
 		RoomInfos: &RistrettoCachePartition[string, *types.RoomInfo]{ // room ID -> room info
 			cache:   cache,
 			Prefix:  roomInfosCache,
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index d8d5f67c..2f9932ff 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -72,7 +72,24 @@ func (d *Database) eventTypeNIDs(
 func (d *Database) EventStateKeys(
 	ctx context.Context, eventStateKeyNIDs []types.EventStateKeyNID,
 ) (map[types.EventStateKeyNID]string, error) {
-	return d.EventStateKeysTable.BulkSelectEventStateKey(ctx, nil, eventStateKeyNIDs)
+	result := make(map[types.EventStateKeyNID]string, len(eventStateKeyNIDs))
+	fetch := make([]types.EventStateKeyNID, 0, len(eventStateKeyNIDs))
+	for _, nid := range eventStateKeyNIDs {
+		if key, ok := d.Cache.GetEventStateKey(nid); ok {
+			result[nid] = key
+		} else {
+			fetch = append(fetch, nid)
+		}
+	}
+	fromDB, err := d.EventStateKeysTable.BulkSelectEventStateKey(ctx, nil, fetch)
+	if err != nil {
+		return nil, err
+	}
+	for nid, key := range fromDB {
+		result[nid] = key
+		d.Cache.StoreEventStateKey(nid, key)
+	}
+	return result, nil
 }
 
 func (d *Database) EventStateKeyNIDs(