Don't load state entries of rejected events

roomserver/internal/helpers/auth.go

@@ -39,7 +39,7 @@ func CheckForSoftFail(
 	var authStateEntries []types.StateEntry
 	var err error
 	if rewritesState {
-		authStateEntries, err = db.StateEntriesForEventIDs(ctx, stateEventIDs)
+		authStateEntries, err = db.StateEntriesForEventIDs(ctx, stateEventIDs, false)
 		if err != nil {
 			return true, fmt.Errorf("StateEntriesForEventIDs failed: %w", err)
 		}
@@ -97,7 +97,7 @@ func CheckAuthEvents(
 	authEventIDs []string,
 ) ([]types.EventNID, error) {
 	// Grab the numeric IDs for the supplied auth state events from the database.
-	authStateEntries, err := db.StateEntriesForEventIDs(ctx, authEventIDs)
+	authStateEntries, err := db.StateEntriesForEventIDs(ctx, authEventIDs, false)
 	if err != nil {
 		return nil, fmt.Errorf("db.StateEntriesForEventIDs: %w", err)
 	}

roomserver/internal/input/input_events.go

@@ -661,7 +661,7 @@ func (r *Inputer) calculateAndSetState(
 		// We've been told what the state at the event is so we don't need to calculate it.
 		// Check that those state events are in the database and store the state.
 		var entries []types.StateEntry
-		if entries, err = r.DB.StateEntriesForEventIDs(ctx, input.StateEventIDs); err != nil {
+		if entries, err = r.DB.StateEntriesForEventIDs(ctx, input.StateEventIDs, true); err != nil {
 			return fmt.Errorf("updater.StateEntriesForEventIDs: %w", err)
 		}
 		entries = types.DeduplicateStateEntries(entries)

roomserver/internal/perform/perform_backfill.go

@@ -140,11 +140,11 @@ func (r *Backfiller) backfillViaFederation(ctx context.Context, req *api.Perform
 			continue
 		}
 		var entries []types.StateEntry
-		if entries, err = r.DB.StateEntriesForEventIDs(ctx, stateIDs); err != nil {
+		if entries, err = r.DB.StateEntriesForEventIDs(ctx, stateIDs, true); err != nil {
 			// attempt to fetch the missing events
 			r.fetchAndStoreMissingEvents(ctx, info.RoomVersion, requester, stateIDs)
 			// try again
-			entries, err = r.DB.StateEntriesForEventIDs(ctx, stateIDs)
+			entries, err = r.DB.StateEntriesForEventIDs(ctx, stateIDs, true)
 			if err != nil {
 				logrus.WithError(err).WithField("event_id", ev.EventID()).Error("backfillViaFederation: failed to get state entries for event")
 				return err

roomserver/storage/interface.go

@@ -79,7 +79,7 @@ type Database interface {
 	// Look up the state entries for a list of string event IDs
 	// Returns an error if the there is an error talking to the database
 	// Returns a types.MissingEventError if the event IDs aren't in the database.
-	StateEntriesForEventIDs(ctx context.Context, eventIDs []string) ([]types.StateEntry, error)
+	StateEntriesForEventIDs(ctx context.Context, eventIDs []string, excludeRejected bool) ([]types.StateEntry, error)
 	// Look up the string event state keys for a list of numeric event state keys
 	// Returns an error if there was a problem talking to the database.
 	EventStateKeys(ctx context.Context, eventStateKeyNIDs []types.EventStateKeyNID) (map[types.EventStateKeyNID]string, error)

roomserver/storage/postgres/events_table.go

@@ -88,6 +88,14 @@ const bulkSelectStateEventByIDSQL = "" +
 	" WHERE event_id = ANY($1)" +
 	" ORDER BY event_type_nid, event_state_key_nid ASC"
 
+// Bulk lookup of events by string ID that aren't excluded.
+// Sort by the numeric IDs for event type and state key.
+// This means we can use binary search to lookup entries by type and state key.
+const bulkSelectStateEventByIDExcludingRejectedSQL = "" +
+	"SELECT event_type_nid, event_state_key_nid, event_nid FROM roomserver_events" +
+	" WHERE event_id = ANY($1) AND is_rejected = FALSE" +
+	" ORDER BY event_type_nid, event_state_key_nid ASC"
+
 // Bulk look up of events by event NID, optionally filtering by the event type
 // or event state key NIDs if provided. (The CARDINALITY check will return true
 // if the provided arrays are empty, ergo no filtering).
@@ -140,23 +148,24 @@ const selectEventRejectedSQL = "" +
 	"SELECT is_rejected FROM roomserver_events WHERE room_nid = $1 AND event_id = $2"
 
 type eventStatements struct {
-	insertEventStmt                        *sql.Stmt
+	insertEventStmt                               *sql.Stmt
-	selectEventStmt                        *sql.Stmt
+	selectEventStmt                               *sql.Stmt
-	bulkSelectStateEventByIDStmt           *sql.Stmt
+	bulkSelectStateEventByIDStmt                  *sql.Stmt
-	bulkSelectStateEventByNIDStmt          *sql.Stmt
+	bulkSelectStateEventByIDExcludingRejectedStmt *sql.Stmt
-	bulkSelectStateAtEventByIDStmt         *sql.Stmt
+	bulkSelectStateEventByNIDStmt                 *sql.Stmt
-	updateEventStateStmt                   *sql.Stmt
+	bulkSelectStateAtEventByIDStmt                *sql.Stmt
-	selectEventSentToOutputStmt            *sql.Stmt
+	updateEventStateStmt                          *sql.Stmt
-	updateEventSentToOutputStmt            *sql.Stmt
+	selectEventSentToOutputStmt                   *sql.Stmt
-	selectEventIDStmt                      *sql.Stmt
+	updateEventSentToOutputStmt                   *sql.Stmt
-	bulkSelectStateAtEventAndReferenceStmt *sql.Stmt
+	selectEventIDStmt                             *sql.Stmt
-	bulkSelectEventReferenceStmt           *sql.Stmt
+	bulkSelectStateAtEventAndReferenceStmt        *sql.Stmt
-	bulkSelectEventIDStmt                  *sql.Stmt
+	bulkSelectEventReferenceStmt                  *sql.Stmt
-	bulkSelectEventNIDStmt                 *sql.Stmt
+	bulkSelectEventIDStmt                         *sql.Stmt
-	bulkSelectUnsentEventNIDStmt           *sql.Stmt
+	bulkSelectEventNIDStmt                        *sql.Stmt
-	selectMaxEventDepthStmt                *sql.Stmt
+	bulkSelectUnsentEventNIDStmt                  *sql.Stmt
-	selectRoomNIDsForEventNIDsStmt         *sql.Stmt
+	selectMaxEventDepthStmt                       *sql.Stmt
-	selectEventRejectedStmt                *sql.Stmt
+	selectRoomNIDsForEventNIDsStmt                *sql.Stmt
+	selectEventRejectedStmt                       *sql.Stmt
 }
 
 func CreateEventsTable(db *sql.DB) error {
@@ -171,6 +180,7 @@ func PrepareEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.insertEventStmt, insertEventSQL},
 		{&s.selectEventStmt, selectEventSQL},
 		{&s.bulkSelectStateEventByIDStmt, bulkSelectStateEventByIDSQL},
+		{&s.bulkSelectStateEventByIDExcludingRejectedStmt, bulkSelectStateEventByIDExcludingRejectedSQL},
 		{&s.bulkSelectStateEventByNIDStmt, bulkSelectStateEventByNIDSQL},
 		{&s.bulkSelectStateAtEventByIDStmt, bulkSelectStateAtEventByIDSQL},
 		{&s.updateEventStateStmt, updateEventStateSQL},
@@ -223,9 +233,14 @@ func (s *eventStatements) SelectEvent(
 // bulkSelectStateEventByID lookups a list of state events by event ID.
 // If any of the requested events are missing from the database it returns a types.MissingEventError
 func (s *eventStatements) BulkSelectStateEventByID(
-	ctx context.Context, txn *sql.Tx, eventIDs []string,
+	ctx context.Context, txn *sql.Tx, eventIDs []string, excludeRejected bool,
 ) ([]types.StateEntry, error) {
-	stmt := sqlutil.TxStmt(txn, s.bulkSelectStateEventByIDStmt)
+	var stmt *sql.Stmt
+	if excludeRejected {
+		stmt = sqlutil.TxStmt(txn, s.bulkSelectStateEventByIDExcludingRejectedStmt)
+	} else {
+		stmt = sqlutil.TxStmt(txn, s.bulkSelectStateEventByIDStmt)
+	}
 	rows, err := stmt.QueryContext(ctx, pq.StringArray(eventIDs))
 	if err != nil {
 		return nil, err
@@ -250,7 +265,7 @@ func (s *eventStatements) BulkSelectStateEventByID(
 	if err = rows.Err(); err != nil {
 		return nil, err
 	}
-	if i != len(eventIDs) {
+	if !excludeRejected && i != len(eventIDs) {
 		// If there are fewer rows returned than IDs then we were asked to lookup event IDs we don't have.
 		// We don't know which ones were missing because we don't return the string IDs in the query.
 		// However it should be possible debug this by replaying queries or entries from the input kafka logs.

roomserver/storage/shared/storage.go

@@ -113,9 +113,9 @@ func (d *Database) eventStateKeyNIDs(
 }
 
 func (d *Database) StateEntriesForEventIDs(
-	ctx context.Context, eventIDs []string,
+	ctx context.Context, eventIDs []string, excludeRejected bool,
 ) ([]types.StateEntry, error) {
-	return d.EventsTable.BulkSelectStateEventByID(ctx, nil, eventIDs)
+	return d.EventsTable.BulkSelectStateEventByID(ctx, nil, eventIDs, excludeRejected)
 }
 
 func (d *Database) StateEntriesForTuples(

roomserver/storage/sqlite3/events_table.go

@@ -65,6 +65,14 @@ const bulkSelectStateEventByIDSQL = "" +
 	" WHERE event_id IN ($1)" +
 	" ORDER BY event_type_nid, event_state_key_nid ASC"
 
+// Bulk lookup of events by string ID that aren't rejected.
+// Sort by the numeric IDs for event type and state key.
+// This means we can use binary search to lookup entries by type and state key.
+const bulkSelectStateEventByIDExcludingRejectedSQL = "" +
+	"SELECT event_type_nid, event_state_key_nid, event_nid FROM roomserver_events" +
+	" WHERE event_id IN ($1) AND Is_rejected = 0" +
+	" ORDER BY event_type_nid, event_state_key_nid ASC"
+
 const bulkSelectStateEventByNIDSQL = "" +
 	"SELECT event_type_nid, event_state_key_nid, event_nid FROM roomserver_events" +
 	" WHERE event_nid IN ($1)"
@@ -113,19 +121,20 @@ const selectEventRejectedSQL = "" +
 	"SELECT is_rejected FROM roomserver_events WHERE room_nid = $1 AND event_id = $2"
 
 type eventStatements struct {
-	db                                     *sql.DB
+	db                                            *sql.DB
-	insertEventStmt                        *sql.Stmt
+	insertEventStmt                               *sql.Stmt
-	selectEventStmt                        *sql.Stmt
+	selectEventStmt                               *sql.Stmt
-	bulkSelectStateEventByIDStmt           *sql.Stmt
+	bulkSelectStateEventByIDStmt                  *sql.Stmt
-	bulkSelectStateAtEventByIDStmt         *sql.Stmt
+	bulkSelectStateEventByIDExcludingRejectedStmt *sql.Stmt
-	updateEventStateStmt                   *sql.Stmt
+	bulkSelectStateAtEventByIDStmt                *sql.Stmt
-	selectEventSentToOutputStmt            *sql.Stmt
+	updateEventStateStmt                          *sql.Stmt
-	updateEventSentToOutputStmt            *sql.Stmt
+	selectEventSentToOutputStmt                   *sql.Stmt
-	selectEventIDStmt                      *sql.Stmt
+	updateEventSentToOutputStmt                   *sql.Stmt
-	bulkSelectStateAtEventAndReferenceStmt *sql.Stmt
+	selectEventIDStmt                             *sql.Stmt
-	bulkSelectEventReferenceStmt           *sql.Stmt
+	bulkSelectStateAtEventAndReferenceStmt        *sql.Stmt
-	bulkSelectEventIDStmt                  *sql.Stmt
+	bulkSelectEventReferenceStmt                  *sql.Stmt
-	selectEventRejectedStmt                *sql.Stmt
+	bulkSelectEventIDStmt                         *sql.Stmt
+	selectEventRejectedStmt                       *sql.Stmt
 	//bulkSelectEventNIDStmt               *sql.Stmt
 	//bulkSelectUnsentEventNIDStmt         *sql.Stmt
 	//selectRoomNIDsForEventNIDsStmt       *sql.Stmt
@@ -145,6 +154,7 @@ func PrepareEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.insertEventStmt, insertEventSQL},
 		{&s.selectEventStmt, selectEventSQL},
 		{&s.bulkSelectStateEventByIDStmt, bulkSelectStateEventByIDSQL},
+		{&s.bulkSelectStateEventByIDExcludingRejectedStmt, bulkSelectStateEventByIDExcludingRejectedSQL},
 		{&s.bulkSelectStateAtEventByIDStmt, bulkSelectStateAtEventByIDSQL},
 		{&s.updateEventStateStmt, updateEventStateSQL},
 		{&s.updateEventSentToOutputStmt, updateEventSentToOutputSQL},
@@ -196,14 +206,20 @@ func (s *eventStatements) SelectEvent(
 // bulkSelectStateEventByID lookups a list of state events by event ID.
 // If any of the requested events are missing from the database it returns a types.MissingEventError
 func (s *eventStatements) BulkSelectStateEventByID(
-	ctx context.Context, txn *sql.Tx, eventIDs []string,
+	ctx context.Context, txn *sql.Tx, eventIDs []string, excludeRejected bool,
 ) ([]types.StateEntry, error) {
 	///////////////
+	var sql string
+	if excludeRejected {
+		sql = bulkSelectStateEventByIDExcludingRejectedSQL
+	} else {
+		sql = bulkSelectStateEventByIDSQL
+	}
 	iEventIDs := make([]interface{}, len(eventIDs))
 	for k, v := range eventIDs {
 		iEventIDs[k] = v
 	}
-	selectOrig := strings.Replace(bulkSelectStateEventByIDSQL, "($1)", sqlutil.QueryVariadic(len(iEventIDs)), 1)
+	selectOrig := strings.Replace(sql, "($1)", sqlutil.QueryVariadic(len(iEventIDs)), 1)
 	selectPrep, err := s.db.Prepare(selectOrig)
 	if err != nil {
 		return nil, err

roomserver/storage/tables/events_table_test.go

@@ -102,7 +102,7 @@ func Test_EventsTable(t *testing.T) {
 			})
 		}
 
-		stateEvents, err := tab.BulkSelectStateEventByID(ctx, nil, eventIDs)
+		stateEvents, err := tab.BulkSelectStateEventByID(ctx, nil, eventIDs, false)
 		assert.NoError(t, err)
 		assert.Equal(t, len(stateEvents), len(eventIDs))
 		nids := make([]types.EventNID, 0, len(stateEvents))

roomserver/storage/tables/interface.go

@@ -46,7 +46,7 @@ type Events interface {
 	SelectEvent(ctx context.Context, txn *sql.Tx, eventID string) (types.EventNID, types.StateSnapshotNID, error)
 	// bulkSelectStateEventByID lookups a list of state events by event ID.
 	// If any of the requested events are missing from the database it returns a types.MissingEventError
-	BulkSelectStateEventByID(ctx context.Context, txn *sql.Tx, eventIDs []string) ([]types.StateEntry, error)
+	BulkSelectStateEventByID(ctx context.Context, txn *sql.Tx, eventIDs []string, excludeRejected bool) ([]types.StateEntry, error)
 	BulkSelectStateEventByNID(ctx context.Context, txn *sql.Tx, eventNIDs []types.EventNID, stateKeyTuples []types.StateKeyTuple) ([]types.StateEntry, error)
 	// BulkSelectStateAtEventByID lookups the state at a list of events by event ID.
 	// If any of the requested events are missing from the database it returns a types.MissingEventError.