From 4b01f1cd122ed155ff8b05aa848f4daf2aa60c34 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Thu, 24 Feb 2022 11:09:01 +0000
Subject: [PATCH] State resolution v2 micro-optimisations (#2226)

* Don't populate duplicates into auth events

* Only append the single event

* Potentially reduce number of iterations in `isInAllAuthLists
---
 roomserver/state/state.go | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/roomserver/state/state.go b/roomserver/state/state.go
index e5f69521e..187b996cd 100644
--- a/roomserver/state/state.go
+++ b/roomserver/state/state.go
@@ -814,6 +814,7 @@ func (v *StateResolution) resolveConflictsV2(
 	// events may be duplicated across these sets but that's OK.
 	authSets := make(map[string][]*gomatrixserverlib.Event, len(conflicted))
 	authEvents := make([]*gomatrixserverlib.Event, 0, estimate*3)
+	gotAuthEvents := make(map[string]struct{}, estimate*3)
 	authDifference := make([]*gomatrixserverlib.Event, 0, estimate)
 
 	// For each conflicted event, let's try and get the needed auth events.
@@ -850,9 +851,22 @@ func (v *StateResolution) resolveConflictsV2(
 		if err != nil {
 			return nil, err
 		}
-		authEvents = append(authEvents, authSets[key]...)
+
+		// Only add auth events into the authEvents slice once, otherwise the
+		// check for the auth difference can become expensive and produce
+		// duplicate entries, which just waste memory and CPU time.
+		for _, event := range authSets[key] {
+			if _, ok := gotAuthEvents[event.EventID()]; !ok {
+				authEvents = append(authEvents, event)
+				gotAuthEvents[event.EventID()] = struct{}{}
+			}
+		}
 	}
 
+	// Kill the reference to this so that the GC may pick it up, since we no
+	// longer need this after this point.
+	gotAuthEvents = nil // nolint:ineffassign
+
 	// This function helps us to work out whether an event exists in one of the
 	// auth sets.
 	isInAuthList := func(k string, event *gomatrixserverlib.Event) bool {
@@ -866,11 +880,12 @@ func (v *StateResolution) resolveConflictsV2(
 
 	// This function works out if an event exists in all of the auth sets.
 	isInAllAuthLists := func(event *gomatrixserverlib.Event) bool {
-		found := true
 		for k := range authSets {
-			found = found && isInAuthList(k, event)
+			if !isInAuthList(k, event) {
+				return false
+			}
 		}
-		return found
+		return true
 	}
 
 	// Look through all of the auth events that we've been given and work out if