Add RoomInfo cache, remove RoomServerRoomNIDsCache (#1646)

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

* Don't panic if the roomInfo isn't known yet

* LRU package is already threadsafe

* Use RoomInfo cache to find room version if possible in Events()

* Adding comments about RoomInfoCache safety

2 files changed, 31 insertions, 9 deletions

diff --git a/roomserver/storage/shared/latest_events_updater.go b/roomserver/storage/shared/latest_events_updater.go
index 8825dc46..36865081 100644
--- a/roomserver/storage/shared/latest_events_updater.go
+++ b/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
 	})
 }
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index d61fa61d..b4d9d562 100644
--- a/roomserver/storage/shared/storage.go
+++ b/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 && roomInfo != nil {
+		d.Cache.StoreRoomServerRoomID(roomInfo.RoomNID, roomID)
+		d.Cache.StoreRoomInfo(roomID, *roomInfo)
+	}
+	return roomInfo, err
 }
 
 func (d *Database) AddState(
@@ -322,14 +330,24 @@ func (d *Database) Events(
 	for _, n := range roomNIDs {
 		uniqueRoomNIDs[n] = struct{}{}
 	}
-	roomNIDList := make([]types.RoomNID, 0, len(uniqueRoomNIDs))
+	roomVersions := make(map[types.RoomNID]gomatrixserverlib.RoomVersion)
+	fetchNIDList := make([]types.RoomNID, 0, len(uniqueRoomNIDs))
 	for n := range uniqueRoomNIDs {
-		roomNIDList = append(roomNIDList, n)
+		if roomID, ok := d.Cache.GetRoomServerRoomID(n); ok {
+			if roomInfo, ok := d.Cache.GetRoomInfo(roomID); ok {
+				roomVersions[n] = roomInfo.RoomVersion
+				continue
+			}
+		}
+		fetchNIDList = append(fetchNIDList, n)
 	}
-	roomVersions, err := d.RoomsTable.SelectRoomVersionsForRoomNIDs(ctx, roomNIDList)
+	dbRoomVersions, err := d.RoomsTable.SelectRoomVersionsForRoomNIDs(ctx, fetchNIDList)
 	if err != nil {
 		return nil, err
 	}
+	for n, v := range dbRoomVersions {
+		roomVersions[n] = v
+	}
 	results := make([]types.Event, len(eventJSONs))
 	for i, eventJSON := range eventJSONs {
 		result := &results[i]
@@ -556,8 +574,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 +587,6 @@ func (d *Database) assignRoomNID(
 			roomNID, err = d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
 		}
 	}
-	if err == nil {
-		d.Cache.StoreRoomServerRoomNID(roomID, roomNID)
-	}
 	return roomNID, err
 }