Tweak `LoadMembershipAtEvent` behaviour when state not known (#2716)

Previously `LoadMembershipAtEvent` would fail if the state before one of
the events was not known, i.e. because it was an outlier. This modifies
it so that it gracefully handles not knowing the state and returns no
memberships instead, so that history visibility doesn't freak out and
kill `/sync` requests dead.

3 files changed, 16 insertions, 4 deletions

diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index 32d63bb5..aa7dc473 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -439,6 +439,7 @@ type QueryMembershipAtEventRequest struct {
 
 // QueryMembershipAtEventResponse is the response to QueryMembershipAtEventRequest.
 type QueryMembershipAtEventResponse struct {
-	// Memberships is a map from eventID to a list of events (if any).
+	// Memberships is a map from eventID to a list of events (if any). Events that
+	// do not have known state will return an empty array here.
 	Memberships map[string][]*gomatrixserverlib.HeaderedEvent `json:"memberships"`
 }
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index d08c5c49..b41a92e9 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -208,6 +208,9 @@ func (r *Queryer) QueryMembershipForUser(
 	return err
 }
 
+// QueryMembershipAtEvent returns the known memberships at a given event.
+// If the state before an event is not known, an empty list will be returned
+// for that event instead.
 func (r *Queryer) QueryMembershipAtEvent(
 	ctx context.Context,
 	request *api.QueryMembershipAtEventRequest,
@@ -237,7 +240,11 @@ func (r *Queryer) QueryMembershipAtEvent(
 	}
 
 	for _, eventID := range request.EventIDs {
-		stateEntry := stateEntries[eventID]
+		stateEntry, ok := stateEntries[eventID]
+		if !ok {
+			response.Memberships[eventID] = []*gomatrixserverlib.HeaderedEvent{}
+			continue
+		}
 		memberships, err := helpers.GetMembershipsAtState(ctx, r.DB, stateEntry, false)
 		if err != nil {
 			return fmt.Errorf("unable to get memberships at state: %w", err)
diff --git a/roomserver/state/state.go b/roomserver/state/state.go
index a40a2e9b..cb96d83e 100644
--- a/roomserver/state/state.go
+++ b/roomserver/state/state.go
@@ -18,6 +18,7 @@ package state
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"sort"
 	"sync"
@@ -134,11 +135,14 @@ func (v *StateResolution) LoadMembershipAtEvent(
 	for i := range eventIDs {
 		eventID := eventIDs[i]
 		snapshotNID, err := v.db.SnapshotNIDFromEventID(ctx, eventID)
-		if err != nil {
+		if err != nil && err != sql.ErrNoRows {
 			return nil, fmt.Errorf("LoadStateAtEvent.SnapshotNIDFromEventID failed for event %s : %w", eventID, err)
 		}
 		if snapshotNID == 0 {
-			return nil, fmt.Errorf("LoadStateAtEvent.SnapshotNIDFromEventID(%s) returned 0 NID, was this event stored?", eventID)
+			// If we don't know a state snapshot for this event then we can't calculate
+			// memberships at the time of the event, so skip over it. This means that
+			// it isn't guaranteed that the response map will contain every single event.
+			continue
 		}
 		snapshotNIDMap[snapshotNID] = append(snapshotNIDMap[snapshotNID], eventID)
 	}