Optimise QueryServerJoinedToRoom (#1933)

* Optimise checking if a server is in a room

* Fix queries

* Fix queries
main
Neil Alexander 2021-07-21 13:06:32 +01:00 committed by GitHub
parent 2f5e4f5991
commit f0f8c7f055
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 64 additions and 38 deletions

View File

@ -181,11 +181,8 @@ type QueryServerJoinedToRoomRequest struct {
type QueryServerJoinedToRoomResponse struct { type QueryServerJoinedToRoomResponse struct {
// True if the room exists on the server // True if the room exists on the server
RoomExists bool `json:"room_exists"` RoomExists bool `json:"room_exists"`
// True if we still believe that we are participating in the room // True if we still believe that the server is participating in the room
IsInRoom bool `json:"is_in_room"` IsInRoom bool `json:"is_in_room"`
// List of servers that are also in the room. This will not be populated
// if the queried ServerName is the local server name.
ServerNames []gomatrixserverlib.ServerName `json:"server_names"`
} }
// QueryServerAllowedToSeeEventRequest is a request to QueryServerAllowedToSeeEvent // QueryServerAllowedToSeeEventRequest is a request to QueryServerAllowedToSeeEvent

View File

@ -330,46 +330,17 @@ func (r *Queryer) QueryServerJoinedToRoom(
response.RoomExists = true response.RoomExists = true
if request.ServerName == r.ServerName || request.ServerName == "" { if request.ServerName == r.ServerName || request.ServerName == "" {
var joined bool response.IsInRoom, err = r.DB.GetLocalServerInRoom(ctx, info.RoomNID)
joined, err = r.DB.GetLocalServerInRoom(ctx, info.RoomNID)
if err != nil { if err != nil {
return fmt.Errorf("r.DB.GetLocalServerInRoom: %w", err) return fmt.Errorf("r.DB.GetLocalServerInRoom: %w", err)
} }
response.IsInRoom = joined } else {
return nil response.IsInRoom, err = r.DB.GetServerInRoom(ctx, info.RoomNID, request.ServerName)
} if err != nil {
return fmt.Errorf("r.DB.GetServerInRoom: %w", err)
eventNIDs, err := r.DB.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, true, false)
if err != nil {
return fmt.Errorf("r.DB.GetMembershipEventNIDsForRoom: %w", err)
}
if len(eventNIDs) == 0 {
return nil
}
events, err := r.DB.Events(ctx, eventNIDs)
if err != nil {
return fmt.Errorf("r.DB.Events: %w", err)
}
servers := map[gomatrixserverlib.ServerName]struct{}{}
for _, e := range events {
if e.Type() == gomatrixserverlib.MRoomMember && e.StateKey() != nil {
_, serverName, err := gomatrixserverlib.SplitID('@', *e.StateKey())
if err != nil {
continue
}
servers[serverName] = struct{}{}
if serverName == request.ServerName {
response.IsInRoom = true
}
} }
} }
for server := range servers {
response.ServerNames = append(response.ServerNames, server)
}
return nil return nil
} }

View File

@ -156,6 +156,8 @@ type Database interface {
JoinedUsersSetInRooms(ctx context.Context, roomIDs []string) (map[string]int, error) JoinedUsersSetInRooms(ctx context.Context, roomIDs []string) (map[string]int, error)
// GetLocalServerInRoom returns true if we think we're in a given room or false otherwise. // GetLocalServerInRoom returns true if we think we're in a given room or false otherwise.
GetLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error) GetLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error)
// GetServerInRoom returns true if we think a server is in a given room or false otherwise.
GetServerInRoom(ctx context.Context, roomNID types.RoomNID, serverName gomatrixserverlib.ServerName) (bool, error)
// GetKnownUsers searches all users that userID knows about. // GetKnownUsers searches all users that userID knows about.
GetKnownUsers(ctx context.Context, userID, searchString string, limit int) ([]string, error) GetKnownUsers(ctx context.Context, userID, searchString string, limit int) ([]string, error)
// GetKnownRooms returns a list of all rooms we know about. // GetKnownRooms returns a list of all rooms we know about.

View File

@ -26,6 +26,7 @@ import (
"github.com/matrix-org/dendrite/roomserver/storage/shared" "github.com/matrix-org/dendrite/roomserver/storage/shared"
"github.com/matrix-org/dendrite/roomserver/storage/tables" "github.com/matrix-org/dendrite/roomserver/storage/tables"
"github.com/matrix-org/dendrite/roomserver/types" "github.com/matrix-org/dendrite/roomserver/types"
"github.com/matrix-org/gomatrixserverlib"
) )
const membershipSchema = ` const membershipSchema = `
@ -132,6 +133,16 @@ var selectKnownUsersSQL = "" +
const selectLocalServerInRoomSQL = "" + const selectLocalServerInRoomSQL = "" +
"SELECT room_nid FROM roomserver_membership WHERE target_local = true AND membership_nid = $1 AND room_nid = $2 LIMIT 1" "SELECT room_nid FROM roomserver_membership WHERE target_local = true AND membership_nid = $1 AND room_nid = $2 LIMIT 1"
// selectServerMembersInRoomSQL is an optimised case for checking for server members in a room.
// The JOIN is significantly leaner than the previous case of looking up event NIDs and reading the
// membership events from the database, as the JOIN query amounts to little more than two index
// scans which are very fast. The presence of a single row from this query suggests the server is
// in the room, no rows returned suggests they aren't.
const selectServerInRoomSQL = "" +
"SELECT room_nid FROM roomserver_membership" +
" JOIN roomserver_event_state_keys ON roomserver_membership.target_nid = roomserver_event_state_keys.event_state_key_nid" +
" WHERE membership_nid = $1 AND room_nid = $2 AND event_state_key LIKE '%:' || $3 LIMIT 1"
type membershipStatements struct { type membershipStatements struct {
insertMembershipStmt *sql.Stmt insertMembershipStmt *sql.Stmt
selectMembershipForUpdateStmt *sql.Stmt selectMembershipForUpdateStmt *sql.Stmt
@ -146,6 +157,7 @@ type membershipStatements struct {
selectKnownUsersStmt *sql.Stmt selectKnownUsersStmt *sql.Stmt
updateMembershipForgetRoomStmt *sql.Stmt updateMembershipForgetRoomStmt *sql.Stmt
selectLocalServerInRoomStmt *sql.Stmt selectLocalServerInRoomStmt *sql.Stmt
selectServerInRoomStmt *sql.Stmt
} }
func createMembershipTable(db *sql.DB) error { func createMembershipTable(db *sql.DB) error {
@ -170,6 +182,7 @@ func prepareMembershipTable(db *sql.DB) (tables.Membership, error) {
{&s.selectKnownUsersStmt, selectKnownUsersSQL}, {&s.selectKnownUsersStmt, selectKnownUsersSQL},
{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom}, {&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
{&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL}, {&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL},
{&s.selectServerInRoomStmt, selectServerInRoomSQL},
}.Prepare(db) }.Prepare(db)
} }
@ -347,3 +360,15 @@ func (s *membershipStatements) SelectLocalServerInRoom(ctx context.Context, room
found := nid > 0 found := nid > 0
return found, nil return found, nil
} }
func (s *membershipStatements) SelectServerInRoom(ctx context.Context, roomNID types.RoomNID, serverName gomatrixserverlib.ServerName) (bool, error) {
var nid types.RoomNID
err := s.selectServerInRoomStmt.QueryRowContext(ctx, tables.MembershipStateJoin, roomNID, serverName).Scan(&nid)
if err != nil {
if err == sql.ErrNoRows {
return false, nil
}
return false, err
}
return roomNID == nid, nil
}

View File

@ -1068,6 +1068,11 @@ func (d *Database) GetLocalServerInRoom(ctx context.Context, roomNID types.RoomN
return d.MembershipTable.SelectLocalServerInRoom(ctx, roomNID) return d.MembershipTable.SelectLocalServerInRoom(ctx, roomNID)
} }
// GetServerInRoom returns true if we think a server is in a given room or false otherwise.
func (d *Database) GetServerInRoom(ctx context.Context, roomNID types.RoomNID, serverName gomatrixserverlib.ServerName) (bool, error) {
return d.MembershipTable.SelectServerInRoom(ctx, roomNID, serverName)
}
// GetKnownUsers searches all users that userID knows about. // GetKnownUsers searches all users that userID knows about.
func (d *Database) GetKnownUsers(ctx context.Context, userID, searchString string, limit int) ([]string, error) { func (d *Database) GetKnownUsers(ctx context.Context, userID, searchString string, limit int) ([]string, error) {
stateKeyNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, nil, userID) stateKeyNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, nil, userID)

View File

@ -26,6 +26,7 @@ import (
"github.com/matrix-org/dendrite/roomserver/storage/shared" "github.com/matrix-org/dendrite/roomserver/storage/shared"
"github.com/matrix-org/dendrite/roomserver/storage/tables" "github.com/matrix-org/dendrite/roomserver/storage/tables"
"github.com/matrix-org/dendrite/roomserver/types" "github.com/matrix-org/dendrite/roomserver/types"
"github.com/matrix-org/gomatrixserverlib"
) )
const membershipSchema = ` const membershipSchema = `
@ -108,6 +109,16 @@ var selectKnownUsersSQL = "" +
const selectLocalServerInRoomSQL = "" + const selectLocalServerInRoomSQL = "" +
"SELECT room_nid FROM roomserver_membership WHERE target_local = 1 AND membership_nid = $1 AND room_nid = $2 LIMIT 1" "SELECT room_nid FROM roomserver_membership WHERE target_local = 1 AND membership_nid = $1 AND room_nid = $2 LIMIT 1"
// selectServerMembersInRoomSQL is an optimised case for checking for server members in a room.
// The JOIN is significantly leaner than the previous case of looking up event NIDs and reading the
// membership events from the database, as the JOIN query amounts to little more than two index
// scans which are very fast. The presence of a single row from this query suggests the server is
// in the room, no rows returned suggests they aren't.
const selectServerInRoomSQL = "" +
"SELECT room_nid FROM roomserver_membership" +
" JOIN roomserver_event_state_keys ON roomserver_membership.target_nid = roomserver_event_state_keys.event_state_key_nid" +
" WHERE membership_nid = $1 AND room_nid = $2 AND event_state_key LIKE '%:' || $3 LIMIT 1"
type membershipStatements struct { type membershipStatements struct {
db *sql.DB db *sql.DB
insertMembershipStmt *sql.Stmt insertMembershipStmt *sql.Stmt
@ -122,6 +133,7 @@ type membershipStatements struct {
selectKnownUsersStmt *sql.Stmt selectKnownUsersStmt *sql.Stmt
updateMembershipForgetRoomStmt *sql.Stmt updateMembershipForgetRoomStmt *sql.Stmt
selectLocalServerInRoomStmt *sql.Stmt selectLocalServerInRoomStmt *sql.Stmt
selectServerInRoomStmt *sql.Stmt
} }
func createMembershipTable(db *sql.DB) error { func createMembershipTable(db *sql.DB) error {
@ -147,6 +159,7 @@ func prepareMembershipTable(db *sql.DB) (tables.Membership, error) {
{&s.selectKnownUsersStmt, selectKnownUsersSQL}, {&s.selectKnownUsersStmt, selectKnownUsersSQL},
{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom}, {&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
{&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL}, {&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL},
{&s.selectServerInRoomStmt, selectServerInRoomSQL},
}.Prepare(db) }.Prepare(db)
} }
@ -327,3 +340,15 @@ func (s *membershipStatements) SelectLocalServerInRoom(ctx context.Context, room
found := nid > 0 found := nid > 0
return found, nil return found, nil
} }
func (s *membershipStatements) SelectServerInRoom(ctx context.Context, roomNID types.RoomNID, serverName gomatrixserverlib.ServerName) (bool, error) {
var nid types.RoomNID
err := s.selectServerInRoomStmt.QueryRowContext(ctx, tables.MembershipStateJoin, roomNID, serverName).Scan(&nid)
if err != nil {
if err == sql.ErrNoRows {
return false, nil
}
return false, err
}
return roomNID == nid, nil
}

View File

@ -136,6 +136,7 @@ type Membership interface {
SelectKnownUsers(ctx context.Context, userID types.EventStateKeyNID, searchString string, limit int) ([]string, error) SelectKnownUsers(ctx context.Context, userID types.EventStateKeyNID, searchString string, limit int) ([]string, error)
UpdateForgetMembership(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, forget bool) error UpdateForgetMembership(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, forget bool) error
SelectLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error) SelectLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error)
SelectServerInRoom(ctx context.Context, roomNID types.RoomNID, serverName gomatrixserverlib.ServerName) (bool, error)
} }
type Published interface { type Published interface {