Optimize `/sync` and history visibility (#2961)

Should fix the following issues or make a lot less worse when using
Postgres:

The main issue behind #2911: The client gives up after a certain time,
causing a cascade of context errors, because the response couldn't be
built up fast enough. This mostly happens on accounts with many rooms,
due to the inefficient way we're getting recent events and current state

For #2777: The queries for getting the membership events for history
visibility were being executed for each room (I think 185?), resulting
in a whooping 2k queries for membership events. (Getting the
statesnapshot -> block nids -> actual wanted membership event)

Both should now be better by:
- Using a LATERAL join to get all recent events for all joined rooms in
one go (TODO: maybe do the same for room summary and current state etc)
- If we're lazy loading on initial syncs, we're now not getting the
whole current state, just to drop the majority of it because we're lazy
loading members - we add a filter to exclude membership events on the
first call to `CurrentState`.
- Using an optimized query to get the membership events needed to
calculate history visibility

---------

Co-authored-by: kegsay <kegan@matrix.org>

7 files changed, 91 insertions, 5 deletions

diff --git a/roomserver/storage/interface.go b/roomserver/storage/interface.go
index e0b9c56b..99891512 100644
--- a/roomserver/storage/interface.go
+++ b/roomserver/storage/interface.go
@@ -175,4 +175,11 @@ type Database interface {
 	GetLeftUsers(ctx context.Context, userIDs []string) ([]string, error)
 	PurgeRoom(ctx context.Context, roomID string) error
 	UpgradeRoom(ctx context.Context, oldRoomID, newRoomID, eventSender string) error
+
+	// GetMembershipForHistoryVisibility queries the membership events for the given eventIDs.
+	// Returns a map from (input) eventID -> membership event. If no membership event is found, returns an empty event, resulting in
+	// a membership of "leave" when calculating history visibility.
+	GetMembershipForHistoryVisibility(
+		ctx context.Context, userNID types.EventStateKeyNID, info *types.RoomInfo, eventIDs ...string,
+	) (map[string]*gomatrixserverlib.HeaderedEvent, error)
 }
diff --git a/roomserver/storage/postgres/events_table.go b/roomserver/storage/postgres/events_table.go
index 9b5ed6ed..f4a21c8a 100644
--- a/roomserver/storage/postgres/events_table.go
+++ b/roomserver/storage/postgres/events_table.go
@@ -69,6 +69,9 @@ CREATE TABLE IF NOT EXISTS roomserver_events (
 	auth_event_nids BIGINT[] NOT NULL,
 	is_rejected BOOLEAN NOT NULL DEFAULT FALSE
 );
+
+-- Create an index which helps in resolving membership events (event_type_nid = 5) - (used for history visibility)
+CREATE INDEX IF NOT EXISTS roomserver_events_memberships_idx ON roomserver_events (room_nid, event_state_key_nid) WHERE (event_type_nid = 5);
 `
 
 const insertEventSQL = "" +
diff --git a/roomserver/storage/postgres/state_snapshot_table.go b/roomserver/storage/postgres/state_snapshot_table.go
index a00c026f..0e83cfc2 100644
--- a/roomserver/storage/postgres/state_snapshot_table.go
+++ b/roomserver/storage/postgres/state_snapshot_table.go
@@ -21,10 +21,10 @@ import (
 	"fmt"
 
 	"github.com/lib/pq"
+	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
 
 	"github.com/matrix-org/dendrite/internal/sqlutil"
-	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
 )
 
@@ -99,10 +99,26 @@ const bulkSelectStateForHistoryVisibilitySQL = `
 	  AND (event_type_nid = 7 OR event_state_key LIKE '%:' || $2);
 `
 
+// bulkSelectMembershipForHistoryVisibilitySQL is an optimization to get membership events for a specific user for defined set of events.
+// Returns the event_id of the event we want the membership event for, the event_id of the membership event and the membership event JSON.
+const bulkSelectMembershipForHistoryVisibilitySQL = `
+SELECT re.event_id, re2.event_id, rej.event_json
+FROM roomserver_events re
+LEFT JOIN roomserver_state_snapshots rss on re.state_snapshot_nid = rss.state_snapshot_nid
+CROSS JOIN unnest(rss.state_block_nids) AS blocks(block_nid)
+LEFT JOIN roomserver_state_block rsb ON rsb.state_block_nid = blocks.block_nid
+CROSS JOIN unnest(rsb.event_nids) AS rsb2(event_nid)
+JOIN roomserver_events re2 ON re2.room_nid = $3 AND re2.event_type_nid = 5 AND re2.event_nid = rsb2.event_nid AND re2.event_state_key_nid = $1
+LEFT JOIN roomserver_event_json rej ON rej.event_nid = re2.event_nid
+WHERE re.event_id = ANY($2)
+
+`
+
 type stateSnapshotStatements struct {
-	insertStateStmt                         *sql.Stmt
-	bulkSelectStateBlockNIDsStmt            *sql.Stmt
-	bulkSelectStateForHistoryVisibilityStmt *sql.Stmt
+	insertStateStmt                               *sql.Stmt
+	bulkSelectStateBlockNIDsStmt                  *sql.Stmt
+	bulkSelectStateForHistoryVisibilityStmt       *sql.Stmt
+	bulktSelectMembershipForHistoryVisibilityStmt *sql.Stmt
 }
 
 func CreateStateSnapshotTable(db *sql.DB) error {
@@ -110,13 +126,14 @@ func CreateStateSnapshotTable(db *sql.DB) error {
 	return err
 }
 
-func PrepareStateSnapshotTable(db *sql.DB) (tables.StateSnapshot, error) {
+func PrepareStateSnapshotTable(db *sql.DB) (*stateSnapshotStatements, error) {
 	s := &stateSnapshotStatements{}
 
 	return s, sqlutil.StatementList{
 		{&s.insertStateStmt, insertStateSQL},
 		{&s.bulkSelectStateBlockNIDsStmt, bulkSelectStateBlockNIDsSQL},
 		{&s.bulkSelectStateForHistoryVisibilityStmt, bulkSelectStateForHistoryVisibilitySQL},
+		{&s.bulktSelectMembershipForHistoryVisibilityStmt, bulkSelectMembershipForHistoryVisibilitySQL},
 	}.Prepare(db)
 }
 
@@ -185,3 +202,45 @@ func (s *stateSnapshotStatements) BulkSelectStateForHistoryVisibility(
 	}
 	return results, rows.Err()
 }
+
+func (s *stateSnapshotStatements) BulkSelectMembershipForHistoryVisibility(
+	ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomInfo *types.RoomInfo, eventIDs ...string,
+) (map[string]*gomatrixserverlib.HeaderedEvent, error) {
+	stmt := sqlutil.TxStmt(txn, s.bulktSelectMembershipForHistoryVisibilityStmt)
+	rows, err := stmt.QueryContext(ctx, userNID, pq.Array(eventIDs), roomInfo.RoomNID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close() // nolint: errcheck
+	result := make(map[string]*gomatrixserverlib.HeaderedEvent, len(eventIDs))
+	var evJson []byte
+	var eventID string
+	var membershipEventID string
+
+	knownEvents := make(map[string]*gomatrixserverlib.HeaderedEvent, len(eventIDs))
+
+	for rows.Next() {
+		if err = rows.Scan(&eventID, &membershipEventID, &evJson); err != nil {
+			return nil, err
+		}
+		if len(evJson) == 0 {
+			result[eventID] = &gomatrixserverlib.HeaderedEvent{}
+			continue
+		}
+		// If we already know this event, don't try to marshal the json again
+		if ev, ok := knownEvents[membershipEventID]; ok {
+			result[eventID] = ev
+			continue
+		}
+		event, err := gomatrixserverlib.NewEventFromTrustedJSON(evJson, false, roomInfo.RoomVersion)
+		if err != nil {
+			result[eventID] = &gomatrixserverlib.HeaderedEvent{}
+			// not fatal
+			continue
+		}
+		he := event.Headered(roomInfo.RoomVersion)
+		result[eventID] = he
+		knownEvents[membershipEventID] = he
+	}
+	return result, rows.Err()
+}
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 654b078d..f8672496 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -51,6 +51,12 @@ func (d *Database) SupportsConcurrentRoomInputs() bool {
 	return true
 }
 
+func (d *Database) GetMembershipForHistoryVisibility(
+	ctx context.Context, userNID types.EventStateKeyNID, roomInfo *types.RoomInfo, eventIDs ...string,
+) (map[string]*gomatrixserverlib.HeaderedEvent, error) {
+	return d.StateSnapshotTable.BulkSelectMembershipForHistoryVisibility(ctx, nil, userNID, roomInfo, eventIDs...)
+}
+
 func (d *Database) EventTypeNIDs(
 	ctx context.Context, eventTypes []string,
 ) (map[string]types.EventTypeNID, error) {
diff --git a/roomserver/storage/sqlite3/state_snapshot_table.go b/roomserver/storage/sqlite3/state_snapshot_table.go
index 930ad14d..e57e1a4b 100644
--- a/roomserver/storage/sqlite3/state_snapshot_table.go
+++ b/roomserver/storage/sqlite3/state_snapshot_table.go
@@ -26,6 +26,7 @@ import (
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
 )
 
@@ -152,6 +153,10 @@ func (s *stateSnapshotStatements) BulkSelectStateForHistoryVisibility(
 	return nil, tables.OptimisationNotSupportedError
 }
 
+func (s *stateSnapshotStatements) BulkSelectMembershipForHistoryVisibility(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomInfo *types.RoomInfo, eventIDs ...string) (map[string]*gomatrixserverlib.HeaderedEvent, error) {
+	return nil, tables.OptimisationNotSupportedError
+}
+
 func (s *stateSnapshotStatements) selectStateBlockNIDsForRoomNID(
 	ctx context.Context, txn *sql.Tx, roomNID types.RoomNID,
 ) ([]types.StateBlockNID, error) {
diff --git a/roomserver/storage/tables/interface.go b/roomserver/storage/tables/interface.go
index 64145f83..c7f1064d 100644
--- a/roomserver/storage/tables/interface.go
+++ b/roomserver/storage/tables/interface.go
@@ -91,6 +91,10 @@ type StateSnapshot interface {
 	// which users are in a room faster than having to load the entire room state. In the
 	// case of SQLite, this will return tables.OptimisationNotSupportedError.
 	BulkSelectStateForHistoryVisibility(ctx context.Context, txn *sql.Tx, stateSnapshotNID types.StateSnapshotNID, domain string) ([]types.EventNID, error)
+
+	BulkSelectMembershipForHistoryVisibility(
+		ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomInfo *types.RoomInfo, eventIDs ...string,
+	) (map[string]*gomatrixserverlib.HeaderedEvent, error)
 }
 
 type StateBlock interface {
diff --git a/roomserver/storage/tables/state_snapshot_table_test.go b/roomserver/storage/tables/state_snapshot_table_test.go
index b2e59377..c7c991b2 100644
--- a/roomserver/storage/tables/state_snapshot_table_test.go
+++ b/roomserver/storage/tables/state_snapshot_table_test.go
@@ -29,6 +29,8 @@ func mustCreateStateSnapshotTable(t *testing.T, dbType test.DBType) (tab tables.
 		assert.NoError(t, err)
 		err = postgres.CreateEventsTable(db)
 		assert.NoError(t, err)
+		err = postgres.CreateEventJSONTable(db)
+		assert.NoError(t, err)
 		err = postgres.CreateStateBlockTable(db)
 		assert.NoError(t, err)
 		// ... and then the snapshot table itself