Hit the database far less in Events to find room NIDs and room versions (#1643)

* Hit the database far less to find room NIDs for event NIDs

* Close the rows

* Fix SQLite selectRoomNIDsForEventNIDsSQL

* Give same treatment to room version lookups

6 files changed, 138 insertions, 65 deletions

diff --git a/roomserver/storage/postgres/events_table.go b/roomserver/storage/postgres/events_table.go
index c8eb8e2d..0cf0bd22 100644
--- a/roomserver/storage/postgres/events_table.go
+++ b/roomserver/storage/postgres/events_table.go
@@ -120,8 +120,8 @@ const bulkSelectEventNIDSQL = "" +
 const selectMaxEventDepthSQL = "" +
 	"SELECT COALESCE(MAX(depth) + 1, 0) FROM roomserver_events WHERE event_nid = ANY($1)"
 
-const selectRoomNIDForEventNIDSQL = "" +
-	"SELECT room_nid FROM roomserver_events WHERE event_nid = $1"
+const selectRoomNIDsForEventNIDsSQL = "" +
+	"SELECT event_nid, room_nid FROM roomserver_events WHERE event_nid = ANY($1)"
 
 type eventStatements struct {
 	insertEventStmt                        *sql.Stmt
@@ -137,7 +137,7 @@ type eventStatements struct {
 	bulkSelectEventIDStmt                  *sql.Stmt
 	bulkSelectEventNIDStmt                 *sql.Stmt
 	selectMaxEventDepthStmt                *sql.Stmt
-	selectRoomNIDForEventNIDStmt           *sql.Stmt
+	selectRoomNIDsForEventNIDsStmt         *sql.Stmt
 }
 
 func NewPostgresEventsTable(db *sql.DB) (tables.Events, error) {
@@ -161,7 +161,7 @@ func NewPostgresEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.bulkSelectEventIDStmt, bulkSelectEventIDSQL},
 		{&s.bulkSelectEventNIDStmt, bulkSelectEventNIDSQL},
 		{&s.selectMaxEventDepthStmt, selectMaxEventDepthSQL},
-		{&s.selectRoomNIDForEventNIDStmt, selectRoomNIDForEventNIDSQL},
+		{&s.selectRoomNIDsForEventNIDsStmt, selectRoomNIDsForEventNIDsSQL},
 	}.Prepare(db)
 }
 
@@ -432,11 +432,24 @@ func (s *eventStatements) SelectMaxEventDepth(ctx context.Context, txn *sql.Tx,
 	return result, nil
 }
 
-func (s *eventStatements) SelectRoomNIDForEventNID(
-	ctx context.Context, eventNID types.EventNID,
-) (roomNID types.RoomNID, err error) {
-	err = s.selectRoomNIDForEventNIDStmt.QueryRowContext(ctx, int64(eventNID)).Scan(&roomNID)
-	return
+func (s *eventStatements) SelectRoomNIDsForEventNIDs(
+	ctx context.Context, eventNIDs []types.EventNID,
+) (map[types.EventNID]types.RoomNID, error) {
+	rows, err := s.selectRoomNIDsForEventNIDsStmt.QueryContext(ctx, eventNIDsAsArray(eventNIDs))
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectRoomNIDsForEventNIDsStmt: rows.close() failed")
+	result := make(map[types.EventNID]types.RoomNID)
+	for rows.Next() {
+		var eventNID types.EventNID
+		var roomNID types.RoomNID
+		if err = rows.Scan(&eventNID, &roomNID); err != nil {
+			return nil, err
+		}
+		result[eventNID] = roomNID
+	}
+	return result, nil
 }
 
 func eventNIDsAsArray(eventNIDs []types.EventNID) pq.Int64Array {
diff --git a/roomserver/storage/postgres/rooms_table.go b/roomserver/storage/postgres/rooms_table.go
index ce635210..637680bd 100644
--- a/roomserver/storage/postgres/rooms_table.go
+++ b/roomserver/storage/postgres/rooms_table.go
@@ -18,7 +18,6 @@ package postgres
 import (
 	"context"
 	"database/sql"
-	"errors"
 
 	"github.com/lib/pq"
 	"github.com/matrix-org/dendrite/internal"
@@ -69,8 +68,8 @@ const selectLatestEventNIDsForUpdateSQL = "" +
 const updateLatestEventNIDsSQL = "" +
 	"UPDATE roomserver_rooms SET latest_event_nids = $2, last_event_sent_nid = $3, state_snapshot_nid = $4 WHERE room_nid = $1"
 
-const selectRoomVersionForRoomNIDSQL = "" +
-	"SELECT room_version FROM roomserver_rooms WHERE room_nid = $1"
+const selectRoomVersionsForRoomNIDsSQL = "" +
+	"SELECT room_nid, room_version FROM roomserver_rooms WHERE room_nid = ANY($1)"
 
 const selectRoomInfoSQL = "" +
 	"SELECT room_version, room_nid, state_snapshot_nid, latest_event_nids FROM roomserver_rooms WHERE room_id = $1"
@@ -90,7 +89,7 @@ type roomStatements struct {
 	selectLatestEventNIDsStmt          *sql.Stmt
 	selectLatestEventNIDsForUpdateStmt *sql.Stmt
 	updateLatestEventNIDsStmt          *sql.Stmt
-	selectRoomVersionForRoomNIDStmt    *sql.Stmt
+	selectRoomVersionsForRoomNIDsStmt  *sql.Stmt
 	selectRoomInfoStmt                 *sql.Stmt
 	selectRoomIDsStmt                  *sql.Stmt
 	bulkSelectRoomIDsStmt              *sql.Stmt
@@ -109,7 +108,7 @@ func NewPostgresRoomsTable(db *sql.DB) (tables.Rooms, error) {
 		{&s.selectLatestEventNIDsStmt, selectLatestEventNIDsSQL},
 		{&s.selectLatestEventNIDsForUpdateStmt, selectLatestEventNIDsForUpdateSQL},
 		{&s.updateLatestEventNIDsStmt, updateLatestEventNIDsSQL},
-		{&s.selectRoomVersionForRoomNIDStmt, selectRoomVersionForRoomNIDSQL},
+		{&s.selectRoomVersionsForRoomNIDsStmt, selectRoomVersionsForRoomNIDsSQL},
 		{&s.selectRoomInfoStmt, selectRoomInfoSQL},
 		{&s.selectRoomIDsStmt, selectRoomIDsSQL},
 		{&s.bulkSelectRoomIDsStmt, bulkSelectRoomIDsSQL},
@@ -219,15 +218,24 @@ func (s *roomStatements) UpdateLatestEventNIDs(
 	return err
 }
 
-func (s *roomStatements) SelectRoomVersionForRoomNID(
-	ctx context.Context, roomNID types.RoomNID,
-) (gomatrixserverlib.RoomVersion, error) {
-	var roomVersion gomatrixserverlib.RoomVersion
-	err := s.selectRoomVersionForRoomNIDStmt.QueryRowContext(ctx, roomNID).Scan(&roomVersion)
-	if err == sql.ErrNoRows {
-		return roomVersion, errors.New("room not found")
+func (s *roomStatements) SelectRoomVersionsForRoomNIDs(
+	ctx context.Context, roomNIDs []types.RoomNID,
+) (map[types.RoomNID]gomatrixserverlib.RoomVersion, error) {
+	rows, err := s.selectRoomVersionsForRoomNIDsStmt.QueryContext(ctx, roomNIDsAsArray(roomNIDs))
+	if err != nil {
+		return nil, err
 	}
-	return roomVersion, err
+	defer internal.CloseAndLogIfError(ctx, rows, "selectRoomVersionsForRoomNIDsStmt: rows.close() failed")
+	result := make(map[types.RoomNID]gomatrixserverlib.RoomVersion)
+	for rows.Next() {
+		var roomNID types.RoomNID
+		var roomVersion gomatrixserverlib.RoomVersion
+		if err = rows.Scan(&roomNID, &roomVersion); err != nil {
+			return nil, err
+		}
+		result[roomNID] = roomVersion
+	}
+	return result, nil
 }
 
 func (s *roomStatements) BulkSelectRoomIDs(ctx context.Context, roomNIDs []types.RoomNID) ([]string, error) {
@@ -271,3 +279,11 @@ func (s *roomStatements) BulkSelectRoomNIDs(ctx context.Context, roomIDs []strin
 	}
 	return roomNIDs, nil
 }
+
+func roomNIDsAsArray(roomNIDs []types.RoomNID) pq.Int64Array {
+	nids := make([]int64, len(roomNIDs))
+	for i := range roomNIDs {
+		nids[i] = int64(roomNIDs[i])
+	}
+	return nids
+}
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 83982299..d61fa61d 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -313,25 +313,29 @@ func (d *Database) Events(
 	if err != nil {
 		eventIDs = map[types.EventNID]string{}
 	}
+	var roomNIDs map[types.EventNID]types.RoomNID
+	roomNIDs, err = d.EventsTable.SelectRoomNIDsForEventNIDs(ctx, eventNIDs)
+	if err != nil {
+		return nil, err
+	}
+	uniqueRoomNIDs := make(map[types.RoomNID]struct{})
+	for _, n := range roomNIDs {
+		uniqueRoomNIDs[n] = struct{}{}
+	}
+	roomNIDList := make([]types.RoomNID, 0, len(uniqueRoomNIDs))
+	for n := range uniqueRoomNIDs {
+		roomNIDList = append(roomNIDList, n)
+	}
+	roomVersions, err := d.RoomsTable.SelectRoomVersionsForRoomNIDs(ctx, roomNIDList)
+	if err != nil {
+		return nil, err
+	}
 	results := make([]types.Event, len(eventJSONs))
 	for i, eventJSON := range eventJSONs {
-		var roomNID types.RoomNID
-		var roomVersion gomatrixserverlib.RoomVersion
 		result := &results[i]
 		result.EventNID = eventJSON.EventNID
-		roomNID, err = d.EventsTable.SelectRoomNIDForEventNID(ctx, eventJSON.EventNID)
-		if err != nil {
-			return nil, err
-		}
-		if roomID, ok := d.Cache.GetRoomServerRoomID(roomNID); ok {
-			roomVersion, _ = d.Cache.GetRoomVersion(roomID)
-		}
-		if roomVersion == "" {
-			roomVersion, err = d.RoomsTable.SelectRoomVersionForRoomNID(ctx, roomNID)
-			if err != nil {
-				return nil, err
-			}
-		}
+		roomNID := roomNIDs[result.EventNID]
+		roomVersion := roomVersions[roomNID]
 		result.Event, err = gomatrixserverlib.NewEventFromTrustedJSONWithEventID(
 			eventIDs[eventJSON.EventNID], eventJSON.EventJSON, false, roomVersion,
 		)
diff --git a/roomserver/storage/sqlite3/events_table.go b/roomserver/storage/sqlite3/events_table.go
index 773e9ade..53269657 100644
--- a/roomserver/storage/sqlite3/events_table.go
+++ b/roomserver/storage/sqlite3/events_table.go
@@ -95,8 +95,8 @@ const bulkSelectEventNIDSQL = "" +
 const selectMaxEventDepthSQL = "" +
 	"SELECT COALESCE(MAX(depth) + 1, 0) FROM roomserver_events WHERE event_nid IN ($1)"
 
-const selectRoomNIDForEventNIDSQL = "" +
-	"SELECT room_nid FROM roomserver_events WHERE event_nid = $1"
+const selectRoomNIDsForEventNIDsSQL = "" +
+	"SELECT event_nid, room_nid FROM roomserver_events WHERE event_nid IN ($1)"
 
 type eventStatements struct {
 	db                                     *sql.DB
@@ -112,7 +112,7 @@ type eventStatements struct {
 	bulkSelectEventReferenceStmt           *sql.Stmt
 	bulkSelectEventIDStmt                  *sql.Stmt
 	bulkSelectEventNIDStmt                 *sql.Stmt
-	selectRoomNIDForEventNIDStmt           *sql.Stmt
+	//selectRoomNIDsForEventNIDsStmt           *sql.Stmt
 }
 
 func NewSqliteEventsTable(db *sql.DB) (tables.Events, error) {
@@ -137,7 +137,7 @@ func NewSqliteEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.bulkSelectEventReferenceStmt, bulkSelectEventReferenceSQL},
 		{&s.bulkSelectEventIDStmt, bulkSelectEventIDSQL},
 		{&s.bulkSelectEventNIDStmt, bulkSelectEventNIDSQL},
-		{&s.selectRoomNIDForEventNIDStmt, selectRoomNIDForEventNIDSQL},
+		//{&s.selectRoomNIDForEventNIDStmt, selectRoomNIDForEventNIDSQL},
 	}.Prepare(db)
 }
 
@@ -480,11 +480,33 @@ func (s *eventStatements) SelectMaxEventDepth(ctx context.Context, txn *sql.Tx,
 	return result, nil
 }
 
-func (s *eventStatements) SelectRoomNIDForEventNID(
-	ctx context.Context, eventNID types.EventNID,
-) (roomNID types.RoomNID, err error) {
-	err = s.selectRoomNIDForEventNIDStmt.QueryRowContext(ctx, int64(eventNID)).Scan(&roomNID)
-	return
+func (s *eventStatements) SelectRoomNIDsForEventNIDs(
+	ctx context.Context, eventNIDs []types.EventNID,
+) (map[types.EventNID]types.RoomNID, error) {
+	sqlStr := strings.Replace(selectRoomNIDsForEventNIDsSQL, "($1)", sqlutil.QueryVariadic(len(eventNIDs)), 1)
+	sqlPrep, err := s.db.Prepare(sqlStr)
+	if err != nil {
+		return nil, err
+	}
+	iEventNIDs := make([]interface{}, len(eventNIDs))
+	for i, v := range eventNIDs {
+		iEventNIDs[i] = v
+	}
+	rows, err := sqlPrep.QueryContext(ctx, iEventNIDs...)
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectRoomNIDsForEventNIDsStmt: rows.close() failed")
+	result := make(map[types.EventNID]types.RoomNID)
+	for rows.Next() {
+		var eventNID types.EventNID
+		var roomNID types.RoomNID
+		if err = rows.Scan(&eventNID, &roomNID); err != nil {
+			return nil, err
+		}
+		result[eventNID] = roomNID
+	}
+	return result, nil
 }
 
 func eventNIDsAsArray(eventNIDs []types.EventNID) string {
diff --git a/roomserver/storage/sqlite3/rooms_table.go b/roomserver/storage/sqlite3/rooms_table.go
index b4564aff..fe8e601f 100644
--- a/roomserver/storage/sqlite3/rooms_table.go
+++ b/roomserver/storage/sqlite3/rooms_table.go
@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"strings"
 
@@ -60,8 +59,8 @@ const selectLatestEventNIDsForUpdateSQL = "" +
 const updateLatestEventNIDsSQL = "" +
 	"UPDATE roomserver_rooms SET latest_event_nids = $1, last_event_sent_nid = $2, state_snapshot_nid = $3 WHERE room_nid = $4"
 
-const selectRoomVersionForRoomNIDSQL = "" +
-	"SELECT room_version FROM roomserver_rooms WHERE room_nid = $1"
+const selectRoomVersionsForRoomNIDsSQL = "" +
+	"SELECT room_nid, room_version FROM roomserver_rooms WHERE room_nid IN ($1)"
 
 const selectRoomInfoSQL = "" +
 	"SELECT room_version, room_nid, state_snapshot_nid, latest_event_nids FROM roomserver_rooms WHERE room_id = $1"
@@ -82,9 +81,9 @@ type roomStatements struct {
 	selectLatestEventNIDsStmt          *sql.Stmt
 	selectLatestEventNIDsForUpdateStmt *sql.Stmt
 	updateLatestEventNIDsStmt          *sql.Stmt
-	selectRoomVersionForRoomNIDStmt    *sql.Stmt
-	selectRoomInfoStmt                 *sql.Stmt
-	selectRoomIDsStmt                  *sql.Stmt
+	//selectRoomVersionForRoomNIDStmt    *sql.Stmt
+	selectRoomInfoStmt *sql.Stmt
+	selectRoomIDsStmt  *sql.Stmt
 }
 
 func NewSqliteRoomsTable(db *sql.DB) (tables.Rooms, error) {
@@ -101,7 +100,7 @@ func NewSqliteRoomsTable(db *sql.DB) (tables.Rooms, error) {
 		{&s.selectLatestEventNIDsStmt, selectLatestEventNIDsSQL},
 		{&s.selectLatestEventNIDsForUpdateStmt, selectLatestEventNIDsForUpdateSQL},
 		{&s.updateLatestEventNIDsStmt, updateLatestEventNIDsSQL},
-		{&s.selectRoomVersionForRoomNIDStmt, selectRoomVersionForRoomNIDSQL},
+		//{&s.selectRoomVersionForRoomNIDsStmt, selectRoomVersionForRoomNIDsSQL},
 		{&s.selectRoomInfoStmt, selectRoomInfoSQL},
 		{&s.selectRoomIDsStmt, selectRoomIDsSQL},
 	}.Prepare(db)
@@ -223,15 +222,33 @@ func (s *roomStatements) UpdateLatestEventNIDs(
 	return err
 }
 
-func (s *roomStatements) SelectRoomVersionForRoomNID(
-	ctx context.Context, roomNID types.RoomNID,
-) (gomatrixserverlib.RoomVersion, error) {
-	var roomVersion gomatrixserverlib.RoomVersion
-	err := s.selectRoomVersionForRoomNIDStmt.QueryRowContext(ctx, roomNID).Scan(&roomVersion)
-	if err == sql.ErrNoRows {
-		return roomVersion, errors.New("room not found")
+func (s *roomStatements) SelectRoomVersionsForRoomNIDs(
+	ctx context.Context, roomNIDs []types.RoomNID,
+) (map[types.RoomNID]gomatrixserverlib.RoomVersion, error) {
+	sqlStr := strings.Replace(selectRoomVersionsForRoomNIDsSQL, "($1)", sqlutil.QueryVariadic(len(roomNIDs)), 1)
+	sqlPrep, err := s.db.Prepare(sqlStr)
+	if err != nil {
+		return nil, err
+	}
+	iRoomNIDs := make([]interface{}, len(roomNIDs))
+	for i, v := range roomNIDs {
+		iRoomNIDs[i] = v
+	}
+	rows, err := sqlPrep.QueryContext(ctx, iRoomNIDs...)
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectRoomVersionsForRoomNIDsStmt: rows.close() failed")
+	result := make(map[types.RoomNID]gomatrixserverlib.RoomVersion)
+	for rows.Next() {
+		var roomNID types.RoomNID
+		var roomVersion gomatrixserverlib.RoomVersion
+		if err = rows.Scan(&roomNID, &roomVersion); err != nil {
+			return nil, err
+		}
+		result[roomNID] = roomVersion
 	}
-	return roomVersion, err
+	return result, nil
 }
 
 func (s *roomStatements) BulkSelectRoomIDs(ctx context.Context, roomNIDs []types.RoomNID) ([]string, error) {
diff --git a/roomserver/storage/tables/interface.go b/roomserver/storage/tables/interface.go
index d7344584..26bf5cf0 100644
--- a/roomserver/storage/tables/interface.go
+++ b/roomserver/storage/tables/interface.go
@@ -10,8 +10,9 @@ import (
 )
 
 type EventJSONPair struct {
-	EventNID  types.EventNID
-	EventJSON []byte
+	EventNID    types.EventNID
+	RoomVersion gomatrixserverlib.RoomVersion
+	EventJSON   []byte
 }
 
 type EventJSON interface {
@@ -58,7 +59,7 @@ type Events interface {
 	// If an event ID is not in the database then it is omitted from the map.
 	BulkSelectEventNID(ctx context.Context, eventIDs []string) (map[string]types.EventNID, error)
 	SelectMaxEventDepth(ctx context.Context, txn *sql.Tx, eventNIDs []types.EventNID) (int64, error)
-	SelectRoomNIDForEventNID(ctx context.Context, eventNID types.EventNID) (roomNID types.RoomNID, err error)
+	SelectRoomNIDsForEventNIDs(ctx context.Context, eventNIDs []types.EventNID) (roomNIDs map[types.EventNID]types.RoomNID, err error)
 }
 
 type Rooms interface {
@@ -67,7 +68,7 @@ type Rooms interface {
 	SelectLatestEventNIDs(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID) ([]types.EventNID, types.StateSnapshotNID, error)
 	SelectLatestEventsNIDsForUpdate(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID) ([]types.EventNID, types.EventNID, types.StateSnapshotNID, error)
 	UpdateLatestEventNIDs(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, eventNIDs []types.EventNID, lastEventSentNID types.EventNID, stateSnapshotNID types.StateSnapshotNID) error
-	SelectRoomVersionForRoomNID(ctx context.Context, roomNID types.RoomNID) (gomatrixserverlib.RoomVersion, error)
+	SelectRoomVersionsForRoomNIDs(ctx context.Context, roomNID []types.RoomNID) (map[types.RoomNID]gomatrixserverlib.RoomVersion, error)
 	SelectRoomInfo(ctx context.Context, roomID string) (*types.RoomInfo, error)
 	SelectRoomIDs(ctx context.Context) ([]string, error)
 	BulkSelectRoomIDs(ctx context.Context, roomNIDs []types.RoomNID) ([]string, error)