From b0d24163c6478df5b8b95b8d4cf0f603a40e21df Mon Sep 17 00:00:00 2001 From: Neil Alexander Date: Tue, 15 Dec 2020 18:02:22 +0000 Subject: [PATCH] Give same treatment to room version lookups --- roomserver/storage/postgres/rooms_table.go | 42 +++++++++++++------ roomserver/storage/shared/storage.go | 26 ++++++------ roomserver/storage/sqlite3/rooms_table.go | 47 +++++++++++++++------- roomserver/storage/tables/interface.go | 2 +- 4 files changed, 76 insertions(+), 41 deletions(-) diff --git a/roomserver/storage/postgres/rooms_table.go b/roomserver/storage/postgres/rooms_table.go index ce635210e..637680bde 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 8301f56d3..d61fa61d3 100644 --- a/roomserver/storage/shared/storage.go +++ b/roomserver/storage/shared/storage.go @@ -318,22 +318,24 @@ func (d *Database) Events( 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 = roomNIDs[result.EventNID] - 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/rooms_table.go b/roomserver/storage/sqlite3/rooms_table.go index b4564aff9..fe8e601f5 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 } - return roomVersion, 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 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 cfd034f31..26bf5cf04 100644 --- a/roomserver/storage/tables/interface.go +++ b/roomserver/storage/tables/interface.go @@ -68,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)