Optimize history visibility checks (#2848)

This optimizes history visibility checks by (mostly) avoiding database
hits.
Possibly solves https://github.com/matrix-org/dendrite/issues/2777

Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>

roomserver/internal/helpers/helpers.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"sort"
 	"strings"
 
 	"github.com/matrix-org/gomatrixserverlib"
@@ -159,7 +160,7 @@ func GetMembershipsAtState(
 	ctx context.Context, db storage.Database, stateEntries []types.StateEntry, joinedOnly bool,
 ) ([]types.Event, error) {
 
-	var eventNIDs []types.EventNID
+	var eventNIDs types.EventNIDs
 	for _, entry := range stateEntries {
 		// Filter the events to retrieve to only keep the membership events
 		if entry.EventTypeNID == types.MRoomMemberNID {
@@ -167,6 +168,14 @@ func GetMembershipsAtState(
 		}
 	}
 
+	// There are no events to get, don't bother asking the database
+	if len(eventNIDs) == 0 {
+		return []types.Event{}, nil
+	}
+
+	sort.Sort(eventNIDs)
+	util.Unique(eventNIDs)
+
 	// Get all of the events in this state
 	stateEvents, err := db.Events(ctx, eventNIDs)
 	if err != nil {

roomserver/internal/query/query.go

@@ -239,16 +239,42 @@ func (r *Queryer) QueryMembershipAtEvent(
 		return fmt.Errorf("unable to get state before event: %w", err)
 	}
 
+	// If we only have one or less state entries, we can short circuit the below
+	// loop and avoid hitting the database
+	allStateEventNIDs := make(map[types.EventNID]types.StateEntry)
+	for _, eventID := range request.EventIDs {
+		stateEntry := stateEntries[eventID]
+		for _, s := range stateEntry {
+			allStateEventNIDs[s.EventNID] = s
+		}
+	}
+
+	var canShortCircuit bool
+	if len(allStateEventNIDs) <= 1 {
+		canShortCircuit = true
+	}
+
+	var memberships []types.Event
 	for _, eventID := range request.EventIDs {
 		stateEntry, ok := stateEntries[eventID]
-		if !ok {
+		if !ok || len(stateEntry) == 0 {
 			response.Memberships[eventID] = []*gomatrixserverlib.HeaderedEvent{}
 			continue
 		}
-		memberships, err := helpers.GetMembershipsAtState(ctx, r.DB, stateEntry, false)
+
+		// If we can short circuit, e.g. we only have 0 or 1 membership events, we only get the memberships
+		// once. If we have more than one membership event, we need to get the state for each state entry.
+		if canShortCircuit {
+			if len(memberships) == 0 {
+				memberships, err = helpers.GetMembershipsAtState(ctx, r.DB, stateEntry, false)
+			}
+		} else {
+			memberships, err = helpers.GetMembershipsAtState(ctx, r.DB, stateEntry, false)
+		}
 		if err != nil {
 			return fmt.Errorf("unable to get memberships at state: %w", err)
 		}
+
 		res := make([]*gomatrixserverlib.HeaderedEvent, 0, len(memberships))
 
 		for i := range memberships {