From e30565c029323944dd42d2c2d2251db5326b09a0 Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Fri, 22 Jul 2022 12:27:36 +0200
Subject: [PATCH] Try to optimize database requests

---
 roomserver/internal/helpers/helpers.go |  4 +-
 roomserver/internal/query/query.go     | 12 ++---
 roomserver/state/state.go              | 62 ++++++++++++++++----------
 3 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/roomserver/internal/helpers/helpers.go b/roomserver/internal/helpers/helpers.go
index 039f33f27..1e4545f4b 100644
--- a/roomserver/internal/helpers/helpers.go
+++ b/roomserver/internal/helpers/helpers.go
@@ -197,10 +197,10 @@ func StateBeforeEvent(ctx context.Context, db storage.Database, info *types.Room
 	return roomState.LoadCombinedStateAfterEvents(ctx, prevState)
 }
 
-func MembershipAtEvent(ctx context.Context, db storage.Database, info *types.RoomInfo, eventID string, stateKeyNID types.EventStateKeyNID) ([]types.StateEntry, error) {
+func MembershipAtEvent(ctx context.Context, db storage.Database, info *types.RoomInfo, eventIDs []string, stateKeyNID types.EventStateKeyNID) (map[string][]types.StateEntry, error) {
 	roomState := state.NewStateResolution(db, info)
 	// Fetch the state as it was when this event was fired
-	return roomState.LoadMembershipAtEvent(ctx, eventID, stateKeyNID)
+	return roomState.LoadMembershipAtEvent(ctx, eventIDs, stateKeyNID)
 }
 
 func LoadEvents(
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index ce3dfe2a2..fffa628c9 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -226,12 +226,14 @@ func (r *Queryer) QueryMembershipAtEvent(
 		return fmt.Errorf("requested stateKeyNID for %s was not found", request.UserID)
 	}
 
+	stateEntries, err := helpers.MembershipAtEvent(ctx, r.DB, info, request.EventIDs, stateKeyNIDs[request.UserID])
+	if err != nil {
+		return fmt.Errorf("unable to get state before event: %w", err)
+	}
+
 	for _, eventID := range request.EventIDs {
-		stateEntries, err := helpers.MembershipAtEvent(ctx, r.DB, info, eventID, stateKeyNIDs[request.UserID])
-		if err != nil {
-			return fmt.Errorf("unable to get state before event: %w", err)
-		}
-		memberships, err := helpers.GetMembershipsAtState(ctx, r.DB, stateEntries, false)
+		stateEntry := stateEntries[eventID]
+		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 ca79cd28d..8a03d41a4 100644
--- a/roomserver/state/state.go
+++ b/roomserver/state/state.go
@@ -124,43 +124,57 @@ func (v *StateResolution) LoadStateAtEvent(
 }
 
 func (v *StateResolution) LoadMembershipAtEvent(
-	ctx context.Context, eventID string, stateKeyNID types.EventStateKeyNID,
-) ([]types.StateEntry, error) {
+	ctx context.Context, eventIDs []string, stateKeyNID types.EventStateKeyNID,
+) (map[string][]types.StateEntry, error) {
 	span, ctx := opentracing.StartSpanFromContext(ctx, "StateResolution.LoadMembershipAtEvent")
 	defer span.Finish()
 
-	snapshotNID, err := v.db.SnapshotNIDFromEventID(ctx, eventID)
-	if err != nil {
-		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)
+	// De-dupe snapshotNIDs
+	snapshotNIDMap := make(map[types.StateSnapshotNID][]string) // map from snapshot NID to eventIDs
+	for i := range eventIDs {
+		eventID := eventIDs[i]
+		snapshotNID, err := v.db.SnapshotNIDFromEventID(ctx, eventID)
+		if err != nil {
+			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)
+		}
+		snapshotNIDMap[snapshotNID] = append(snapshotNIDMap[snapshotNID], eventID)
 	}
 
-	stateBlockNIDLists, err := v.db.StateBlockNIDs(ctx, []types.StateSnapshotNID{snapshotNID})
-	if err != nil {
-		return nil, err
+	snapshotNIDs := make([]types.StateSnapshotNID, 0, len(snapshotNIDMap))
+	for nid := range snapshotNIDMap {
+		snapshotNIDs = append(snapshotNIDs, nid)
 	}
-	// We've asked for exactly one snapshot from the db so we should have exactly one entry in the result.
-	stateBlockNIDList := stateBlockNIDLists[0]
 
-	// Query the membership event for the user at the given stateblocks
-	stateEntryLists, err := v.db.StateEntriesForTuples(ctx, stateBlockNIDList.StateBlockNIDs, []types.StateKeyTuple{
-		{
-			EventTypeNID:     types.MRoomMemberNID,
-			EventStateKeyNID: stateKeyNID,
-		},
-	})
+	stateBlockNIDLists, err := v.db.StateBlockNIDs(ctx, snapshotNIDs)
 	if err != nil {
 		return nil, err
 	}
 
-	var result []types.StateEntry
-	for _, x := range stateEntryLists {
-		if len(x.StateEntries) > 0 {
-			result = append(result, x.StateEntries...)
+	result := make(map[string][]types.StateEntry)
+	for _, stateBlockNIDList := range stateBlockNIDLists {
+		// Query the membership event for the user at the given stateblocks
+		stateEntryLists, err := v.db.StateEntriesForTuples(ctx, stateBlockNIDList.StateBlockNIDs, []types.StateKeyTuple{
+			{
+				EventTypeNID:     types.MRoomMemberNID,
+				EventStateKeyNID: stateKeyNID,
+			},
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		evIDs := snapshotNIDMap[stateBlockNIDList.StateSnapshotNID]
+
+		for _, evID := range evIDs {
+			for _, x := range stateEntryLists {
+				result[evID] = append(result[evID], x.StateEntries...)
+			}
 		}
 	}
+
 	return result, nil
 }