From a883d180f1af4f5c5d5638d3b96571d2baeccac6 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 10 Oct 2022 16:38:27 +0100
Subject: [PATCH] Just use two different code paths actually

---
 federationapi/storage/shared/storage.go | 27 ++++++++++++++-----------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/federationapi/storage/shared/storage.go b/federationapi/storage/shared/storage.go
index 5e1ba2f5c..9e40f311c 100644
--- a/federationapi/storage/shared/storage.go
+++ b/federationapi/storage/shared/storage.go
@@ -70,25 +70,28 @@ func (d *Database) UpdateRoom(
 ) (joinedHosts []types.JoinedHost, err error) {
 	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
 		if purgeRoomFirst {
-			// If the event is a create event then we'll delete all of the existing
-			// data for the room. The only reason that a create event would be replayed
-			// to us in this way is if we're about to receive the entire room state.
 			if err = d.FederationJoinedHosts.DeleteJoinedHostsForRoom(ctx, txn, roomID); err != nil {
 				return fmt.Errorf("d.FederationJoinedHosts.DeleteJoinedHosts: %w", err)
 			}
-		} else if len(removeHosts) > 0 {
+			for _, add := range addHosts {
+				if err = d.FederationJoinedHosts.InsertJoinedHosts(ctx, txn, roomID, add.MemberEventID, add.ServerName); err != nil {
+					return err
+				}
+				joinedHosts = append(joinedHosts, add)
+			}
+		} else {
+			if joinedHosts, err = d.FederationJoinedHosts.SelectJoinedHostsWithTx(ctx, txn, roomID); err != nil {
+				return err
+			}
+			for _, add := range addHosts {
+				if err = d.FederationJoinedHosts.InsertJoinedHosts(ctx, txn, roomID, add.MemberEventID, add.ServerName); err != nil {
+					return err
+				}
+			}
 			if err = d.FederationJoinedHosts.DeleteJoinedHosts(ctx, txn, removeHosts); err != nil {
 				return err
 			}
 		}
-		for _, add := range addHosts {
-			if err = d.FederationJoinedHosts.InsertJoinedHosts(ctx, txn, roomID, add.MemberEventID, add.ServerName); err != nil {
-				return err
-			}
-		}
-		if joinedHosts, err = d.FederationJoinedHosts.SelectJoinedHostsWithTx(ctx, txn, roomID); err != nil {
-			return err
-		}
 		return nil
 	})
 	return