Add RoomInfo cache, remove RoomServerRoomNID cache, ensure caches are thread-safe

internal/caching/cache_roominfo.go

@@ -0,0 +1,32 @@
+package caching
+
+import (
+	"github.com/matrix-org/dendrite/roomserver/types"
+)
+
+const (
+	RoomInfoCacheName       = "roominfo"
+	RoomInfoCacheMaxEntries = 1024
+	RoomInfoCacheMutable    = true
+)
+
+// RoomInfosCache contains the subset of functions needed for
+// a room Info cache.
+type RoomInfoCache interface {
+	GetRoomInfo(roomID string) (roomInfo types.RoomInfo, ok bool)
+	StoreRoomInfo(roomID string, roomInfo types.RoomInfo)
+}
+
+func (c Caches) GetRoomInfo(roomID string) (types.RoomInfo, bool) {
+	val, found := c.RoomInfos.Get(roomID)
+	if found && val != nil {
+		if roomInfo, ok := val.(types.RoomInfo); ok {
+			return roomInfo, true
+		}
+	}
+	return types.RoomInfo{}, false
+}
+
+func (c Caches) StoreRoomInfo(roomID string, roomInfo types.RoomInfo) {
+	c.RoomInfos.Set(roomID, roomInfo)
+}

internal/caching/cache_roomservernids.go

@@ -15,10 +15,6 @@ const (
 	RoomServerEventTypeNIDsCacheMaxEntries = 64
 	RoomServerEventTypeNIDsCacheMutable    = false
 
-	RoomServerRoomNIDsCacheName       = "roomserver_room_nids"
-	RoomServerRoomNIDsCacheMaxEntries = 1024
-	RoomServerRoomNIDsCacheMutable    = false
-
 	RoomServerRoomIDsCacheName       = "roomserver_room_ids"
 	RoomServerRoomIDsCacheMaxEntries = 1024
 	RoomServerRoomIDsCacheMutable    = false
@@ -27,6 +23,7 @@ const (
 type RoomServerCaches interface {
 	RoomServerNIDsCache
 	RoomVersionCache
+	RoomInfoCache
 }
 
 // RoomServerNIDsCache contains the subset of functions needed for
@@ -38,9 +35,6 @@ type RoomServerNIDsCache interface {
 	GetRoomServerEventTypeNID(eventType string) (types.EventTypeNID, bool)
 	StoreRoomServerEventTypeNID(eventType string, nid types.EventTypeNID)
 
-	GetRoomServerRoomNID(roomID string) (types.RoomNID, bool)
-	StoreRoomServerRoomNID(roomID string, nid types.RoomNID)
-
 	GetRoomServerRoomID(roomNID types.RoomNID) (string, bool)
 	StoreRoomServerRoomID(roomNID types.RoomNID, roomID string)
 }
@@ -73,21 +67,6 @@ func (c Caches) StoreRoomServerEventTypeNID(eventType string, nid types.EventTyp
 	c.RoomServerEventTypeNIDs.Set(eventType, nid)
 }
 
-func (c Caches) GetRoomServerRoomNID(roomID string) (types.RoomNID, bool) {
-	val, found := c.RoomServerRoomNIDs.Get(roomID)
-	if found && val != nil {
-		if roomNID, ok := val.(types.RoomNID); ok {
-			return roomNID, true
-		}
-	}
-	return 0, false
-}
-
-func (c Caches) StoreRoomServerRoomNID(roomID string, roomNID types.RoomNID) {
-	c.RoomServerRoomNIDs.Set(roomID, roomNID)
-	c.RoomServerRoomIDs.Set(strconv.Itoa(int(roomNID)), roomID)
-}
-
 func (c Caches) GetRoomServerRoomID(roomNID types.RoomNID) (string, bool) {
 	val, found := c.RoomServerRoomIDs.Get(strconv.Itoa(int(roomNID)))
 	if found && val != nil {
@@ -99,5 +78,5 @@ func (c Caches) GetRoomServerRoomID(roomNID types.RoomNID) (string, bool) {
 }
 
 func (c Caches) StoreRoomServerRoomID(roomNID types.RoomNID, roomID string) {
-	c.StoreRoomServerRoomNID(roomID, roomNID)
+	c.RoomServerRoomIDs.Set(strconv.Itoa(int(roomNID)), roomID)
 }

internal/caching/caches.go

@@ -10,6 +10,7 @@ type Caches struct {
 	RoomServerEventTypeNIDs Cache // RoomServerNIDsCache
 	RoomServerRoomNIDs      Cache // RoomServerNIDsCache
 	RoomServerRoomIDs       Cache // RoomServerNIDsCache
+	RoomInfos               Cache // RoomInfoCache
 	FederationEvents        Cache // FederationEventsCache
 }
 

internal/caching/impl_inmemorylru.go

@@ -2,6 +2,7 @@ package caching
 
 import (
 	"fmt"
+	"sync"
 
 	lru "github.com/hashicorp/golang-lru"
 	"github.com/prometheus/client_golang/prometheus"
@@ -45,19 +46,19 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 	if err != nil {
 		return nil, err
 	}
-	roomServerRoomNIDs, err := NewInMemoryLRUCachePartition(
-		RoomServerRoomNIDsCacheName,
-		RoomServerRoomNIDsCacheMutable,
-		RoomServerRoomNIDsCacheMaxEntries,
+	roomServerRoomIDs, err := NewInMemoryLRUCachePartition(
+		RoomServerRoomIDsCacheName,
+		RoomServerRoomIDsCacheMutable,
+		RoomServerRoomIDsCacheMaxEntries,
 		enablePrometheus,
 	)
 	if err != nil {
 		return nil, err
 	}
-	roomServerRoomIDs, err := NewInMemoryLRUCachePartition(
-		RoomServerRoomIDsCacheName,
-		RoomServerRoomIDsCacheMutable,
-		RoomServerRoomIDsCacheMaxEntries,
+	roomInfos, err := NewInMemoryLRUCachePartition(
+		RoomInfoCacheName,
+		RoomInfoCacheMutable,
+		RoomInfoCacheMaxEntries,
 		enablePrometheus,
 	)
 	if err != nil {
@@ -77,8 +78,8 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 		ServerKeys:              serverKeys,
 		RoomServerStateKeyNIDs:  roomServerStateKeyNIDs,
 		RoomServerEventTypeNIDs: roomServerEventTypeNIDs,
-		RoomServerRoomNIDs:      roomServerRoomNIDs,
 		RoomServerRoomIDs:       roomServerRoomIDs,
+		RoomInfos:               roomInfos,
 		FederationEvents:        federationEvents,
 	}, nil
 }
@@ -87,6 +88,7 @@ type InMemoryLRUCachePartition struct {
 	name       string
 	mutable    bool
 	maxEntries int
+	mutex      sync.RWMutex
 	lru        *lru.Cache
 }
 
@@ -114,6 +116,8 @@ func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int, ena
 }
 
 func (c *InMemoryLRUCachePartition) Set(key string, value interface{}) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
 	if !c.mutable {
 		if peek, ok := c.lru.Peek(key); ok && peek != value {
 			panic(fmt.Sprintf("invalid use of immutable cache tries to mutate existing value of %q", key))
@@ -123,6 +127,8 @@ func (c *InMemoryLRUCachePartition) Set(key string, value interface{}) {
 }
 
 func (c *InMemoryLRUCachePartition) Unset(key string) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
 	if !c.mutable {
 		panic(fmt.Sprintf("invalid use of immutable cache tries to unset value of %q", key))
 	}
@@ -130,5 +136,7 @@ func (c *InMemoryLRUCachePartition) Unset(key string) {
 }
 
 func (c *InMemoryLRUCachePartition) Get(key string) (value interface{}, ok bool) {
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
 	return c.lru.Get(key)
 }

roomserver/storage/shared/latest_events_updater.go

@@ -105,6 +105,13 @@ func (u *LatestEventsUpdater) SetLatestEvents(
 		if err := u.d.RoomsTable.UpdateLatestEventNIDs(u.ctx, txn, roomNID, eventNIDs, lastEventNIDSent, currentStateSnapshotNID); err != nil {
 			return fmt.Errorf("u.d.RoomsTable.updateLatestEventNIDs: %w", err)
 		}
+		if roomID, ok := u.d.Cache.GetRoomServerRoomID(roomNID); ok {
+			if roomInfo, ok := u.d.Cache.GetRoomInfo(roomID); ok {
+				roomInfo.StateSnapshotNID = currentStateSnapshotNID
+				roomInfo.IsStub = false
+				u.d.Cache.StoreRoomInfo(roomID, roomInfo)
+			}
+		}
 		return nil
 	})
 }

roomserver/storage/shared/storage.go

@@ -124,7 +124,15 @@ func (d *Database) StateEntriesForTuples(
 }
 
 func (d *Database) RoomInfo(ctx context.Context, roomID string) (*types.RoomInfo, error) {
-	return d.RoomsTable.SelectRoomInfo(ctx, roomID)
+	if roomInfo, ok := d.Cache.GetRoomInfo(roomID); ok {
+		return &roomInfo, nil
+	}
+	roomInfo, err := d.RoomsTable.SelectRoomInfo(ctx, roomID)
+	if err == nil {
+		d.Cache.StoreRoomServerRoomID(roomInfo.RoomNID, roomID)
+		d.Cache.StoreRoomInfo(roomID, *roomInfo)
+	}
+	return roomInfo, err
 }
 
 func (d *Database) AddState(
@@ -556,8 +564,8 @@ func (d *Database) assignRoomNID(
 	ctx context.Context, txn *sql.Tx,
 	roomID string, roomVersion gomatrixserverlib.RoomVersion,
 ) (types.RoomNID, error) {
-	if roomNID, ok := d.Cache.GetRoomServerRoomNID(roomID); ok {
-		return roomNID, nil
+	if roomInfo, ok := d.Cache.GetRoomInfo(roomID); ok {
+		return roomInfo.RoomNID, nil
 	}
 	// Check if we already have a numeric ID in the database.
 	roomNID, err := d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
@@ -569,9 +577,6 @@ func (d *Database) assignRoomNID(
 			roomNID, err = d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
 		}
 	}
-	if err == nil {
-		d.Cache.StoreRoomServerRoomNID(roomID, roomNID)
-	}
 	return roomNID, err
 }