From f269e437914497e5aeb2f5cc58d48e0e18adcf98 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 7 Jul 2021 16:45:30 +0100
Subject: [PATCH] Add hadEvent function

---
 federationapi/routing/send.go | 32 +++++++++++++-------------------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/federationapi/routing/send.go b/federationapi/routing/send.go
index 989224a49..1c9e72bff 100644
--- a/federationapi/routing/send.go
+++ b/federationapi/routing/send.go
@@ -241,6 +241,12 @@ type txnReq struct {
 	work            string // metrics
 }
 
+func (t *txnReq) hadEvent(eventID string, had bool) {
+	t.hadEventsMutex.Lock()
+	defer t.hadEventsMutex.Unlock()
+	t.hadEvents[eventID] = had
+}
+
 // A subset of FederationClient functionality that txn requires. Useful for testing.
 type txnFederationClient interface {
 	LookupState(ctx context.Context, s gomatrixserverlib.ServerName, roomID string, eventID string, roomVersion gomatrixserverlib.RoomVersion) (
@@ -595,14 +601,12 @@ func (t *txnReq) processEvent(ctx context.Context, e *gomatrixserverlib.Event) e
 
 	// Prepare a map of all the events we already had before this point, so
 	// that we don't send them to the roomserver again.
-	t.hadEventsMutex.Lock()
 	for _, eventID := range append(e.AuthEventIDs(), e.PrevEventIDs()...) {
-		t.hadEvents[eventID] = true
+		t.hadEvent(eventID, true)
 	}
 	for _, eventID := range append(stateResp.MissingAuthEventIDs, stateResp.MissingPrevEventIDs...) {
-		t.hadEvents[eventID] = false
+		t.hadEvent(eventID, false)
 	}
-	t.hadEventsMutex.Unlock()
 
 	if len(stateResp.MissingAuthEventIDs) > 0 {
 		t.work = MetricsWorkMissingAuthEvents
@@ -676,9 +680,7 @@ withNextEvent:
 			); err != nil {
 				return fmt.Errorf("api.SendEvents: %w", err)
 			}
-			t.hadEventsMutex.Lock()
-			t.hadEvents[ev.EventID()] = true // if the roomserver didn't know about the event before, it does now
-			t.hadEventsMutex.Unlock()
+			t.hadEvent(ev.EventID(), true) // if the roomserver didn't know about the event before, it does now
 			t.cacheAndReturn(ev.Headered(stateResp.RoomVersion))
 			delete(missingAuthEvents, missingAuthEventID)
 			continue withNextEvent
@@ -918,9 +920,7 @@ func (t *txnReq) lookupStateAfterEventLocally(ctx context.Context, roomID, event
 		// set the event from the haveEvents cache - this means we will share pointers with other prev_event branches for this
 		// processEvent request, which is better for memory.
 		stateEvents[i] = t.cacheAndReturn(ev)
-		t.hadEventsMutex.Lock()
-		t.hadEvents[ev.EventID()] = true
-		t.hadEventsMutex.Unlock()
+		t.hadEvent(ev.EventID(), true)
 	}
 	// we should never access res.StateEvents again so we delete it here to make GC faster
 	res.StateEvents = nil
@@ -955,9 +955,7 @@ func (t *txnReq) lookupStateAfterEventLocally(ctx context.Context, roomID, event
 		}
 		for i, ev := range queryRes.Events {
 			authEvents = append(authEvents, t.cacheAndReturn(queryRes.Events[i]).Unwrap())
-			t.hadEventsMutex.Lock()
-			t.hadEvents[ev.EventID()] = true
-			t.hadEventsMutex.Unlock()
+			t.hadEvent(ev.EventID(), true)
 		}
 		queryRes.Events = nil
 	}
@@ -1034,9 +1032,7 @@ func (t *txnReq) getMissingEvents(ctx context.Context, e *gomatrixserverlib.Even
 	latestEvents := make([]string, len(res.LatestEvents))
 	for i, ev := range res.LatestEvents {
 		latestEvents[i] = res.LatestEvents[i].EventID
-		t.hadEventsMutex.Lock()
-		t.hadEvents[ev.EventID] = true
-		t.hadEventsMutex.Unlock()
+		t.hadEvent(ev.EventID, true)
 	}
 
 	var missingResp *gomatrixserverlib.RespMissingEvents
@@ -1172,9 +1168,7 @@ func (t *txnReq) lookupMissingStateViaStateIDs(ctx context.Context, roomID, even
 	}
 	for i, ev := range queryRes.Events {
 		queryRes.Events[i] = t.cacheAndReturn(queryRes.Events[i])
-		t.hadEventsMutex.Lock()
-		t.hadEvents[ev.EventID()] = true
-		t.hadEventsMutex.Unlock()
+		t.hadEvent(ev.EventID(), true)
 		evID := queryRes.Events[i].EventID()
 		if missing[evID] {
 			delete(missing, evID)