Optimise QueryServerJoinedToRoom (#1933)

* Optimise checking if a server is in a room

* Fix queries

* Fix queries

1 files changed, 25 insertions, 0 deletions

diff --git a/roomserver/storage/sqlite3/membership_table.go b/roomserver/storage/sqlite3/membership_table.go
index 82babe0d..911a2516 100644
--- a/roomserver/storage/sqlite3/membership_table.go
+++ b/roomserver/storage/sqlite3/membership_table.go
@@ -26,6 +26,7 @@ import (
 	"github.com/matrix-org/dendrite/roomserver/storage/shared"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
 )
 
 const membershipSchema = `
@@ -108,6 +109,16 @@ var selectKnownUsersSQL = "" +
 const selectLocalServerInRoomSQL = "" +
 	"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 {
 	db                                              *sql.DB
 	insertMembershipStmt                            *sql.Stmt
@@ -122,6 +133,7 @@ type membershipStatements struct {
 	selectKnownUsersStmt                            *sql.Stmt
 	updateMembershipForgetRoomStmt                  *sql.Stmt
 	selectLocalServerInRoomStmt                     *sql.Stmt
+	selectServerInRoomStmt                          *sql.Stmt
 }
 
 func createMembershipTable(db *sql.DB) error {
@@ -147,6 +159,7 @@ func prepareMembershipTable(db *sql.DB) (tables.Membership, error) {
 		{&s.selectKnownUsersStmt, selectKnownUsersSQL},
 		{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
 		{&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL},
+		{&s.selectServerInRoomStmt, selectServerInRoomSQL},
 	}.Prepare(db)
 }
 
@@ -327,3 +340,15 @@ func (s *membershipStatements) SelectLocalServerInRoom(ctx context.Context, room
 	found := nid > 0
 	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
+}