From 1a7f4c8aa978d7e2f6046b6628ecf523460eee28 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 9 May 2022 15:22:33 +0100
Subject: [PATCH] Don't try to re-fetch the event if it is listed in
 `adds_state_event_ids` (#2437)

* Don't try to re-fetch the event in the output message

* Try that again

* Add the initial event into the set
---
 syncapi/consumers/roomserver.go | 75 ++++++++++++++++++++++++---------
 1 file changed, 55 insertions(+), 20 deletions(-)

diff --git a/syncapi/consumers/roomserver.go b/syncapi/consumers/roomserver.go
index 7712c840..e1c2ea82 100644
--- a/syncapi/consumers/roomserver.go
+++ b/syncapi/consumers/roomserver.go
@@ -154,41 +154,76 @@ func (s *OutputRoomEventConsumer) onNewRoomEvent(
 	ctx context.Context, msg api.OutputNewRoomEvent,
 ) error {
 	ev := msg.Event
-
 	addsStateEvents := []*gomatrixserverlib.HeaderedEvent{}
-	foundEventIDs := map[string]bool{}
-	if len(msg.AddsStateEventIDs) > 0 {
-		for _, eventID := range msg.AddsStateEventIDs {
-			foundEventIDs[eventID] = false
+
+	// Work out the list of events we need to find out about. Either
+	// they will be the event supplied in the request, we will find it
+	// in the sync API database or we'll need to ask the roomserver.
+	knownEventIDs := make(map[string]bool, len(msg.AddsStateEventIDs))
+	for _, eventID := range msg.AddsStateEventIDs {
+		if eventID == ev.EventID() {
+			knownEventIDs[eventID] = true
+			addsStateEvents = append(addsStateEvents, ev)
+		} else {
+			knownEventIDs[eventID] = false
 		}
-		foundEvents, err := s.db.Events(ctx, msg.AddsStateEventIDs)
+	}
+
+	// Work out which events we want to look up in the sync API database.
+	// At this stage the only event that should be excluded is the event
+	// supplied in the request, if it appears in the adds_state_event_ids.
+	missingEventIDs := make([]string, 0, len(msg.AddsStateEventIDs))
+	for eventID, known := range knownEventIDs {
+		if !known {
+			missingEventIDs = append(missingEventIDs, eventID)
+		}
+	}
+
+	// Look the events up in the database. If we know them, add them into
+	// the set of adds state events.
+	if len(missingEventIDs) > 0 {
+		alreadyKnown, err := s.db.Events(ctx, msg.AddsStateEventIDs)
 		if err != nil {
 			return fmt.Errorf("s.db.Events: %w", err)
 		}
-		for _, event := range foundEvents {
-			foundEventIDs[event.EventID()] = true
+		for _, knownEvent := range alreadyKnown {
+			knownEventIDs[knownEvent.EventID()] = true
+			addsStateEvents = append(addsStateEvents, knownEvent)
+		}
+	}
+
+	// Now work out if there are any remaining events we don't know. For
+	// these we will need to ask the roomserver for help.
+	missingEventIDs = missingEventIDs[:0]
+	for eventID, known := range knownEventIDs {
+		if !known {
+			missingEventIDs = append(missingEventIDs, eventID)
+		}
+	}
+
+	// Ask the roomserver and add in the rest of the results into the set.
+	// Finally, work out if there are any more events missing.
+	if len(missingEventIDs) > 0 {
+		eventsReq := &api.QueryEventsByIDRequest{
+			EventIDs: missingEventIDs,
 		}
-		eventsReq := &api.QueryEventsByIDRequest{}
 		eventsRes := &api.QueryEventsByIDResponse{}
-		for eventID, found := range foundEventIDs {
-			if !found {
-				eventsReq.EventIDs = append(eventsReq.EventIDs, eventID)
-			}
-		}
-		if err = s.rsAPI.QueryEventsByID(ctx, eventsReq, eventsRes); err != nil {
+		if err := s.rsAPI.QueryEventsByID(ctx, eventsReq, eventsRes); err != nil {
 			return fmt.Errorf("s.rsAPI.QueryEventsByID: %w", err)
 		}
 		for _, event := range eventsRes.Events {
-			eventID := event.EventID()
-			foundEvents = append(foundEvents, event)
-			foundEventIDs[eventID] = true
+			addsStateEvents = append(addsStateEvents, event)
+			knownEventIDs[event.EventID()] = true
 		}
-		for eventID, found := range foundEventIDs {
+
+		// This should never happen because this would imply that the
+		// roomserver has sent us adds_state_event_ids for events that it
+		// also doesn't know about, but let's just be sure.
+		for eventID, found := range knownEventIDs {
 			if !found {
 				return fmt.Errorf("event %s is missing", eventID)
 			}
 		}
-		addsStateEvents = foundEvents
 	}
 
 	ev, err := s.updateStateEvent(ev)