Roomserver NID caches (#1335)

* Initial work on roomserver NID caches

* Give caches to roomserver storage

* Populate caches

* Fix bugs

* Fix WASM build

* Don't hit cache twice in RoomNIDExcludingStubs

* Store reverse room ID-room NID mapping, consult caches when assigning NIDs

internal/caching/cache_roomservernids.go

@@ -0,0 +1,101 @@
+package caching
+
+import (
+	"github.com/matrix-org/dendrite/roomserver/types"
+)
+
+const (
+	RoomServerStateKeyNIDsCacheName       = "roomserver_statekey_nids"
+	RoomServerStateKeyNIDsCacheMaxEntries = 1024
+	RoomServerStateKeyNIDsCacheMutable    = false
+
+	RoomServerEventTypeNIDsCacheName       = "roomserver_eventtype_nids"
+	RoomServerEventTypeNIDsCacheMaxEntries = 64
+	RoomServerEventTypeNIDsCacheMutable    = false
+
+	RoomServerRoomNIDsCacheName       = "roomserver_room_nids"
+	RoomServerRoomNIDsCacheMaxEntries = 1024
+	RoomServerRoomNIDsCacheMutable    = false
+
+	RoomServerRoomIDsCacheName       = "roomserver_room_ids"
+	RoomServerRoomIDsCacheMaxEntries = 1024
+	RoomServerRoomIDsCacheMutable    = false
+)
+
+type RoomServerCaches interface {
+	RoomServerNIDsCache
+	RoomVersionCache
+}
+
+// RoomServerNIDsCache contains the subset of functions needed for
+// a roomserver NID cache.
+type RoomServerNIDsCache interface {
+	GetRoomServerStateKeyNID(stateKey string) (types.EventStateKeyNID, bool)
+	StoreRoomServerStateKeyNID(stateKey string, nid types.EventStateKeyNID)
+
+	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)
+}
+
+func (c Caches) GetRoomServerStateKeyNID(stateKey string) (types.EventStateKeyNID, bool) {
+	val, found := c.RoomServerStateKeyNIDs.Get(stateKey)
+	if found && val != nil {
+		if stateKeyNID, ok := val.(types.EventStateKeyNID); ok {
+			return stateKeyNID, true
+		}
+	}
+	return 0, false
+}
+
+func (c Caches) StoreRoomServerStateKeyNID(stateKey string, nid types.EventStateKeyNID) {
+	c.RoomServerStateKeyNIDs.Set(stateKey, nid)
+}
+
+func (c Caches) GetRoomServerEventTypeNID(eventType string) (types.EventTypeNID, bool) {
+	val, found := c.RoomServerEventTypeNIDs.Get(eventType)
+	if found && val != nil {
+		if eventTypeNID, ok := val.(types.EventTypeNID); ok {
+			return eventTypeNID, true
+		}
+	}
+	return 0, false
+}
+
+func (c Caches) StoreRoomServerEventTypeNID(eventType string, nid types.EventTypeNID) {
+	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(string(roomNID), roomID)
+}
+
+func (c Caches) GetRoomServerRoomID(roomNID types.RoomNID) (string, bool) {
+	val, found := c.RoomServerRoomIDs.Get(string(roomNID))
+	if found && val != nil {
+		if roomID, ok := val.(string); ok {
+			return roomID, true
+		}
+	}
+	return "", false
+}
+
+func (c Caches) StoreRoomServerRoomID(roomNID types.RoomNID, roomID string) {
+	c.StoreRoomServerRoomNID(roomID, roomNID)
+}

internal/caching/caches.go

@@ -4,8 +4,12 @@ package caching
 // different implementations as long as they satisfy the Cache
 // interface.
 type Caches struct {
-	RoomVersions Cache // implements RoomVersionCache
+	RoomVersions            Cache // RoomVersionCache
-	ServerKeys   Cache // implements ServerKeyCache
+	ServerKeys              Cache // ServerKeyCache
+	RoomServerStateKeyNIDs  Cache // RoomServerNIDsCache
+	RoomServerEventTypeNIDs Cache // RoomServerNIDsCache
+	RoomServerRoomNIDs      Cache // RoomServerNIDsCache
+	RoomServerRoomIDs       Cache // RoomServerNIDsCache
 }
 
 // Cache is the interface that an implementation must satisfy.

internal/caching/impl_inmemorylru.go

@@ -27,9 +27,49 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 	if err != nil {
 		return nil, err
 	}
+	roomServerStateKeyNIDs, err := NewInMemoryLRUCachePartition(
+		RoomServerStateKeyNIDsCacheName,
+		RoomServerStateKeyNIDsCacheMutable,
+		RoomServerStateKeyNIDsCacheMaxEntries,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	roomServerEventTypeNIDs, err := NewInMemoryLRUCachePartition(
+		RoomServerEventTypeNIDsCacheName,
+		RoomServerEventTypeNIDsCacheMutable,
+		RoomServerEventTypeNIDsCacheMaxEntries,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	roomServerRoomNIDs, err := NewInMemoryLRUCachePartition(
+		RoomServerRoomNIDsCacheName,
+		RoomServerRoomNIDsCacheMutable,
+		RoomServerRoomNIDsCacheMaxEntries,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
+	roomServerRoomIDs, err := NewInMemoryLRUCachePartition(
+		RoomServerRoomIDsCacheName,
+		RoomServerRoomIDsCacheMutable,
+		RoomServerRoomIDsCacheMaxEntries,
+		enablePrometheus,
+	)
+	if err != nil {
+		return nil, err
+	}
 	return &Caches{
 		RoomVersions:            roomVersions,
 		ServerKeys:              serverKeys,
+		RoomServerStateKeyNIDs:  roomServerStateKeyNIDs,
+		RoomServerEventTypeNIDs: roomServerEventTypeNIDs,
+		RoomServerRoomNIDs:      roomServerRoomNIDs,
+		RoomServerRoomIDs:       roomServerRoomIDs,
 	}, nil
 }
 

roomserver/internal/api.go

@@ -16,7 +16,7 @@ type RoomserverInternalAPI struct {
 	DB                   storage.Database
 	Cfg                  *config.RoomServer
 	Producer             sarama.SyncProducer
-	Cache                caching.RoomVersionCache
+	Cache                caching.RoomServerCaches
 	ServerName           gomatrixserverlib.ServerName
 	KeyRing              gomatrixserverlib.JSONVerifier
 	FedClient            *gomatrixserverlib.FederationClient

roomserver/roomserver.go

@@ -42,7 +42,7 @@ func NewInternalAPI(
 ) api.RoomserverInternalAPI {
 	cfg := &base.Cfg.RoomServer
 
-	roomserverDB, err := storage.Open(&cfg.Database)
+	roomserverDB, err := storage.Open(&cfg.Database, base.Caches)
 	if err != nil {
 		logrus.WithError(err).Panicf("failed to connect to room server db")
 	}

roomserver/storage/postgres/storage.go

@@ -18,6 +18,7 @@ package postgres
 import (
 	"database/sql"
 
+	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 
@@ -33,7 +34,7 @@ type Database struct {
 
 // Open a postgres database.
 // nolint: gocyclo
-func Open(dbProperties *config.DatabaseOptions) (*Database, error) {
+func Open(dbProperties *config.DatabaseOptions, cache caching.RoomServerCaches) (*Database, error) {
 	var d Database
 	var db *sql.DB
 	var err error
@@ -98,6 +99,7 @@ func Open(dbProperties *config.DatabaseOptions) (*Database, error) {
 	}
 	d.Database = shared.Database{
 		DB:                  db,
+		Cache:               cache,
 		Writer:              sqlutil.NewDummyWriter(),
 		EventTypesTable:     eventTypes,
 		EventStateKeysTable: eventStateKeys,

roomserver/storage/shared/storage.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
@@ -27,6 +28,7 @@ const redactionsArePermanent = false
 
 type Database struct {
 	DB                  *sql.DB
+	Cache               caching.RoomServerCaches
 	Writer              sqlutil.Writer
 	EventsTable         tables.Events
 	EventJSONTable      tables.EventJSON
@@ -51,7 +53,26 @@ func (d *Database) SupportsConcurrentRoomInputs() bool {
 func (d *Database) EventTypeNIDs(
 	ctx context.Context, eventTypes []string,
 ) (map[string]types.EventTypeNID, error) {
-	return d.EventTypesTable.BulkSelectEventTypeNID(ctx, eventTypes)
+	result := make(map[string]types.EventTypeNID)
+	remaining := []string{}
+	for _, eventType := range eventTypes {
+		if nid, ok := d.Cache.GetRoomServerEventTypeNID(eventType); ok {
+			result[eventType] = nid
+		} else {
+			remaining = append(remaining, eventType)
+		}
+	}
+	if len(remaining) > 0 {
+		nids, err := d.EventTypesTable.BulkSelectEventTypeNID(ctx, remaining)
+		if err != nil {
+			return nil, err
+		}
+		for eventType, nid := range nids {
+			result[eventType] = nid
+			d.Cache.StoreRoomServerEventTypeNID(eventType, nid)
+		}
+	}
+	return result, nil
 }
 
 func (d *Database) EventStateKeys(
@@ -63,7 +84,26 @@ func (d *Database) EventStateKeys(
 func (d *Database) EventStateKeyNIDs(
 	ctx context.Context, eventStateKeys []string,
 ) (map[string]types.EventStateKeyNID, error) {
-	return d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, eventStateKeys)
+	result := make(map[string]types.EventStateKeyNID)
+	remaining := []string{}
+	for _, eventStateKey := range eventStateKeys {
+		if nid, ok := d.Cache.GetRoomServerStateKeyNID(eventStateKey); ok {
+			result[eventStateKey] = nid
+		} else {
+			remaining = append(remaining, eventStateKey)
+		}
+	}
+	if len(remaining) > 0 {
+		nids, err := d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, remaining)
+		if err != nil {
+			return nil, err
+		}
+		for eventStateKey, nid := range nids {
+			result[eventStateKey] = nid
+			d.Cache.StoreRoomServerStateKeyNID(eventStateKey, nid)
+		}
+	}
+	return result, nil
 }
 
 func (d *Database) StateEntriesForEventIDs(
@@ -157,10 +197,14 @@ func (d *Database) EventsFromIDs(ctx context.Context, eventIDs []string) ([]type
 }
 
 func (d *Database) RoomNID(ctx context.Context, roomID string) (types.RoomNID, error) {
+	if nid, ok := d.Cache.GetRoomServerRoomNID(roomID); ok {
+		return nid, nil
+	}
 	roomNID, err := d.RoomsTable.SelectRoomNID(ctx, nil, roomID)
 	if err == sql.ErrNoRows {
 		return 0, nil
 	}
+	d.Cache.StoreRoomServerRoomNID(roomID, roomNID)
 	return roomNID, err
 }
 
@@ -214,6 +258,9 @@ func (d *Database) StateEntries(
 func (d *Database) GetRoomVersionForRoom(
 	ctx context.Context, roomID string,
 ) (gomatrixserverlib.RoomVersion, error) {
+	if roomVersion, ok := d.Cache.GetRoomVersion(roomID); ok {
+		return roomVersion, nil
+	}
 	return d.RoomsTable.SelectRoomVersionForRoomID(
 		ctx, nil, roomID,
 	)
@@ -222,6 +269,11 @@ func (d *Database) GetRoomVersionForRoom(
 func (d *Database) GetRoomVersionForRoomNID(
 	ctx context.Context, roomNID types.RoomNID,
 ) (gomatrixserverlib.RoomVersion, error) {
+	if roomID, ok := d.Cache.GetRoomServerRoomID(roomNID); ok {
+		if roomVersion, ok := d.Cache.GetRoomVersion(roomID); ok {
+			return roomVersion, nil
+		}
+	}
 	return d.RoomsTable.SelectRoomVersionForRoomNID(
 		ctx, roomNID,
 	)
@@ -488,6 +540,9 @@ 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
+	}
 	// Check if we already have a numeric ID in the database.
 	roomNID, err := d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
 	if err == sql.ErrNoRows {
@@ -498,14 +553,20 @@ func (d *Database) assignRoomNID(
 			roomNID, err = d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
 		}
 	}
+	if err == nil {
+		d.Cache.StoreRoomServerRoomNID(roomID, roomNID)
+	}
 	return roomNID, err
 }
 
 func (d *Database) assignEventTypeNID(
 	ctx context.Context, txn *sql.Tx, eventType string,
-) (eventTypeNID types.EventTypeNID, err error) {
+) (types.EventTypeNID, error) {
+	if eventTypeNID, ok := d.Cache.GetRoomServerEventTypeNID(eventType); ok {
+		return eventTypeNID, nil
+	}
 	// Check if we already have a numeric ID in the database.
-	eventTypeNID, err = d.EventTypesTable.SelectEventTypeNID(ctx, txn, eventType)
+	eventTypeNID, err := d.EventTypesTable.SelectEventTypeNID(ctx, txn, eventType)
 	if err == sql.ErrNoRows {
 		// We don't have a numeric ID so insert one into the database.
 		eventTypeNID, err = d.EventTypesTable.InsertEventTypeNID(ctx, txn, eventType)
@@ -514,12 +575,18 @@ func (d *Database) assignEventTypeNID(
 			eventTypeNID, err = d.EventTypesTable.SelectEventTypeNID(ctx, txn, eventType)
 		}
 	}
-	return
+	if err == nil {
+		d.Cache.StoreRoomServerEventTypeNID(eventType, eventTypeNID)
+	}
+	return eventTypeNID, err
 }
 
 func (d *Database) assignStateKeyNID(
 	ctx context.Context, txn *sql.Tx, eventStateKey string,
 ) (types.EventStateKeyNID, error) {
+	if eventStateKeyNID, ok := d.Cache.GetRoomServerStateKeyNID(eventStateKey); ok {
+		return eventStateKeyNID, nil
+	}
 	// Check if we already have a numeric ID in the database.
 	eventStateKeyNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, txn, eventStateKey)
 	if err == sql.ErrNoRows {
@@ -530,6 +597,9 @@ func (d *Database) assignStateKeyNID(
 			eventStateKeyNID, err = d.EventStateKeysTable.SelectEventStateKeyNID(ctx, txn, eventStateKey)
 		}
 	}
+	if err == nil {
+		d.Cache.StoreRoomServerStateKeyNID(eventStateKey, eventStateKeyNID)
+	}
 	return eventStateKeyNID, err
 }
 

roomserver/storage/sqlite3/storage.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"database/sql"
 
+	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/storage/shared"
@@ -46,7 +47,7 @@ type Database struct {
 
 // Open a sqlite database.
 // nolint: gocyclo
-func Open(dbProperties *config.DatabaseOptions) (*Database, error) {
+func Open(dbProperties *config.DatabaseOptions, cache caching.RoomServerCaches) (*Database, error) {
 	var d Database
 	var err error
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
@@ -120,6 +121,7 @@ func Open(dbProperties *config.DatabaseOptions) (*Database, error) {
 	}
 	d.Database = shared.Database{
 		DB:                  d.db,
+		Cache:               cache,
 		Writer:              sqlutil.NewExclusiveWriter(),
 		EventsTable:         d.events,
 		EventTypesTable:     d.eventTypes,

roomserver/storage/storage.go

@@ -19,18 +19,19 @@ package storage
 import (
 	"fmt"
 
+	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/roomserver/storage/postgres"
 	"github.com/matrix-org/dendrite/roomserver/storage/sqlite3"
 )
 
 // Open opens a database connection.
-func Open(dbProperties *config.DatabaseOptions) (Database, error) {
+func Open(dbProperties *config.DatabaseOptions, cache caching.RoomServerCaches) (Database, error) {
 	switch {
 	case dbProperties.ConnectionString.IsSQLite():
-		return sqlite3.Open(dbProperties)
+		return sqlite3.Open(dbProperties, cache)
 	case dbProperties.ConnectionString.IsPostgres():
-		return postgres.Open(dbProperties)
+		return postgres.Open(dbProperties, cache)
 	default:
 		return nil, fmt.Errorf("unexpected database type")
 	}

roomserver/storage/storage_wasm.go

@@ -17,15 +17,16 @@ package storage
 import (
 	"fmt"
 
+	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/roomserver/storage/sqlite3"
 )
 
 // NewPublicRoomsServerDatabase opens a database connection.
-func Open(dbProperties *config.DatabaseOptions) (Database, error) {
+func Open(dbProperties *config.DatabaseOptions, cache caching.RoomServerCaches) (Database, error) {
 	switch {
 	case dbProperties.ConnectionString.IsSQLite():
-		return sqlite3.Open(dbProperties)
+		return sqlite3.Open(dbProperties, cache)
 	case dbProperties.ConnectionString.IsPostgres():
 		return nil, fmt.Errorf("can't use Postgres implementation")
 	default: