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(