Add more optimised code path for checking if we're in a room

federationapi/routing/send.go

@@ -599,6 +599,13 @@ func (t *txnReq) processEvent(ctx context.Context, e *gomatrixserverlib.Event) e
 		return roomNotFoundError{e.RoomID()}
 	}
 
+	if !stateResp.RoomJoined {
+		// We don't believe we're a member of this room anymore, therefore
+		// there's no point in wasting work trying to figure out what to do
+		// with missing auth or prev events. Silently drop the event.
+		return nil
+	}
+
 	// Prepare a map of all the events we already had before this point, so
 	// that we don't send them to the roomserver again.
 	for _, eventID := range append(e.AuthEventIDs(), e.PrevEventIDs()...) {

roomserver/api/query.go

@@ -96,6 +96,8 @@ type QueryMissingAuthPrevEventsResponse struct {
 	// Does the room exist on this roomserver?
 	// If the room doesn't exist all other fields will be empty.
 	RoomExists bool `json:"room_exists"`
+	// Are we joined to this room locally?
+	RoomJoined bool `json:"room_joined"`
 	// The room version of the room.
 	RoomVersion gomatrixserverlib.RoomVersion `json:"room_version"`
 	// The event IDs of the auth events that we don't know locally.
@@ -182,7 +184,8 @@ type QueryServerJoinedToRoomResponse struct {
 	RoomExists bool `json:"room_exists"`
 	// True if we still believe that we are participating in the room
 	IsInRoom bool `json:"is_in_room"`
-	// List of servers that are also in the 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"`
 }
 

roomserver/internal/api.go

@@ -59,6 +59,7 @@ func NewRoomserverAPI(
 		Queryer: &query.Queryer{
 			DB:         roomserverDB,
 			Cache:      caches,
+			ServerName: cfg.Matrix.ServerName,
 			ServerACLs: serverACLs,
 		},
 		Inputer: &input.Inputer{

roomserver/internal/query/query.go

@@ -36,6 +36,7 @@ import (
 type Queryer struct {
 	DB         storage.Database
 	Cache      caching.RoomServerCaches
+	ServerName gomatrixserverlib.ServerName
 	ServerACLs *acls.ServerACLs
 }
 
@@ -141,6 +142,18 @@ func (r *Queryer) QueryMissingAuthPrevEvents(
 	response.RoomExists = !info.IsStub
 	response.RoomVersion = info.RoomVersion
 
+	joined, err := r.DB.GetLocalServerInRoom(ctx, info.RoomNID)
+	if err != nil {
+		return fmt.Errorf("r.DB.GetLocalServerInRoom: %w", err)
+	}
+	response.RoomJoined = joined
+
+	// If we're not joined to the room then there's no point in hitting
+	// the database further to work out which events we're missing.
+	if !joined {
+		return nil
+	}
+
 	for _, authEventID := range request.AuthEventIDs {
 		if nids, err := r.DB.EventNIDs(ctx, []string{authEventID}); err != nil || len(nids) == 0 {
 			response.MissingAuthEventIDs = append(response.MissingAuthEventIDs, authEventID)
@@ -328,6 +341,16 @@ func (r *Queryer) QueryServerJoinedToRoom(
 	}
 	response.RoomExists = true
 
+	if request.ServerName == r.ServerName {
+		var joined bool
+		joined, err = r.DB.GetLocalServerInRoom(ctx, info.RoomNID)
+		if err != nil {
+			return fmt.Errorf("r.DB.GetLocalServerInRoom: %w", err)
+		}
+		response.IsInRoom = joined
+		return nil
+	}
+
 	eventNIDs, err := r.DB.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, true, false)
 	if err != nil {
 		return fmt.Errorf("r.DB.GetMembershipEventNIDsForRoom: %w", err)

roomserver/storage/interface.go

@@ -154,6 +154,8 @@ type Database interface {
 	GetBulkStateContent(ctx context.Context, roomIDs []string, tuples []gomatrixserverlib.StateKeyTuple, allowWildcards bool) ([]tables.StrippedEvent, error)
 	// JoinedUsersSetInRooms returns all joined users in the rooms given, along with the count of how many times they appear.
 	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(ctx context.Context, roomNID types.RoomNID) (bool, error)
 	// GetKnownUsers searches all users that userID knows about.
 	GetKnownUsers(ctx context.Context, userID, searchString string, limit int) ([]string, error)
 	// GetKnownRooms returns a list of all rooms we know about.

roomserver/storage/postgres/membership_table.go

@@ -124,6 +124,14 @@ var selectKnownUsersSQL = "" +
 	"  SELECT DISTINCT room_nid FROM roomserver_membership WHERE target_nid=$1 AND membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) +
 	") AND membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) + " AND event_state_key LIKE $2 LIMIT $3"
 
+// selectLocalServerInRoomSQL is an optimised case for checking if we, the local server,
+// are in the room by using the target_local column of the membership table. Normally when
+// we want to know if a server is in a room, we have to unmarshal the entire room state which
+// is expensive. The presence of a single row from this query suggests we're still in the
+// room, no rows returned suggests we aren't.
+const selectLocalServerInRoomSQL = "" +
+	"SELECT TOP 1 room_nid FROM roomserver_membership WHERE target_local = true AND membership_nid = 3 AND room_nid = $1"
+
 type membershipStatements struct {
 	insertMembershipStmt                            *sql.Stmt
 	selectMembershipForUpdateStmt                   *sql.Stmt
@@ -137,6 +145,7 @@ type membershipStatements struct {
 	selectJoinedUsersSetForRoomsStmt                *sql.Stmt
 	selectKnownUsersStmt                            *sql.Stmt
 	updateMembershipForgetRoomStmt                  *sql.Stmt
+	selectLocalServerInRoomStmt                     *sql.Stmt
 }
 
 func createMembershipTable(db *sql.DB) error {
@@ -160,6 +169,7 @@ func prepareMembershipTable(db *sql.DB) (tables.Membership, error) {
 		{&s.selectJoinedUsersSetForRoomsStmt, selectJoinedUsersSetForRoomsSQL},
 		{&s.selectKnownUsersStmt, selectKnownUsersSQL},
 		{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
+		{&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL},
 	}.Prepare(db)
 }
 
@@ -324,3 +334,19 @@ func (s *membershipStatements) UpdateForgetMembership(
 	)
 	return err
 }
+
+func (s *membershipStatements) SelectLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error) {
+	rows, err := s.selectLocalServerInRoomStmt.QueryContext(ctx, roomNID)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return false, nil
+		}
+		return false, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "SelectKnownUsers: rows.close() failed")
+	found := false
+	for rows.Next() {
+		found = true
+	}
+	return found, rows.Err()
+}

roomserver/storage/shared/storage.go

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

roomserver/storage/sqlite3/membership_table.go

@@ -100,6 +100,14 @@ var selectKnownUsersSQL = "" +
 	"  SELECT DISTINCT room_nid FROM roomserver_membership WHERE target_nid=$1 AND membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) +
 	") AND membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) + " AND event_state_key LIKE $2 LIMIT $3"
 
+// selectLocalServerInRoomSQL is an optimised case for checking if we, the local server,
+// are in the room by using the target_local column of the membership table. Normally when
+// we want to know if a server is in a room, we have to unmarshal the entire room state which
+// is expensive. The presence of a single row from this query suggests we're still in the
+// room, no rows returned suggests we aren't.
+const selectLocalServerInRoomSQL = "" +
+	"SELECT room_nid FROM roomserver_membership WHERE target_local = true AND membership_nid = 3 AND room_nid = $1 LIMIT 1"
+
 type membershipStatements struct {
 	db                                              *sql.DB
 	insertMembershipStmt                            *sql.Stmt
@@ -113,6 +121,7 @@ type membershipStatements struct {
 	updateMembershipStmt                            *sql.Stmt
 	selectKnownUsersStmt                            *sql.Stmt
 	updateMembershipForgetRoomStmt                  *sql.Stmt
+	selectLocalServerInRoomStmt                     *sql.Stmt
 }
 
 func createMembershipTable(db *sql.DB) error {
@@ -137,6 +146,7 @@ func prepareMembershipTable(db *sql.DB) (tables.Membership, error) {
 		{&s.selectRoomsWithMembershipStmt, selectRoomsWithMembershipSQL},
 		{&s.selectKnownUsersStmt, selectKnownUsersSQL},
 		{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
+		{&s.selectLocalServerInRoomStmt, selectLocalServerInRoomSQL},
 	}.Prepare(db)
 }
 
@@ -304,3 +314,19 @@ func (s *membershipStatements) UpdateForgetMembership(
 	)
 	return err
 }
+
+func (s *membershipStatements) SelectLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error) {
+	rows, err := s.selectLocalServerInRoomStmt.QueryContext(ctx, roomNID)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return false, nil
+		}
+		return false, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "SelectKnownUsers: rows.close() failed")
+	found := false
+	for rows.Next() {
+		found = true
+	}
+	return found, rows.Err()
+}

roomserver/storage/tables/interface.go

@@ -135,6 +135,7 @@ type Membership interface {
 	SelectJoinedUsersSetForRooms(ctx context.Context, roomNIDs []types.RoomNID) (map[types.EventStateKeyNID]int, 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
+	SelectLocalServerInRoom(ctx context.Context, roomNID types.RoomNID) (bool, error)
 }
 
 type Published interface {