From fdbb2406bef9594b1ea93f6f01e32c21092a3e42 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 1 Feb 2022 14:54:28 +0000
Subject: [PATCH] More tweaks

---
 roomserver/storage/shared/room_updater.go | 14 +++++++-----
 roomserver/storage/shared/storage.go      | 27 ++++++++---------------
 2 files changed, 17 insertions(+), 24 deletions(-)

diff --git a/roomserver/storage/shared/room_updater.go b/roomserver/storage/shared/room_updater.go
index 59cabf5da..5726ca5e9 100644
--- a/roomserver/storage/shared/room_updater.go
+++ b/roomserver/storage/shared/room_updater.go
@@ -83,19 +83,21 @@ func (u *RoomUpdater) CurrentStateSnapshotNID() types.StateSnapshotNID {
 
 // StorePreviousEvents implements types.RoomRecentEventsUpdater - This must be called from a Writer
 func (u *RoomUpdater) StorePreviousEvents(eventNID types.EventNID, previousEventReferences []gomatrixserverlib.EventReference) error {
-	for _, ref := range previousEventReferences {
-		if err := u.d.PrevEventsTable.InsertPreviousEvent(u.ctx, u.txn, ref.EventID, ref.EventSHA256, eventNID); err != nil {
-			return fmt.Errorf("u.d.PrevEventsTable.InsertPreviousEvent: %w", err)
+	return u.d.Writer.Do(u.d.DB, u.txn, func(txn *sql.Tx) error {
+		for _, ref := range previousEventReferences {
+			if err := u.d.PrevEventsTable.InsertPreviousEvent(u.ctx, txn, ref.EventID, ref.EventSHA256, eventNID); err != nil {
+				return fmt.Errorf("u.d.PrevEventsTable.InsertPreviousEvent: %w", err)
+			}
 		}
-	}
-	return nil
+		return nil
+	})
 }
 
 func (u *RoomUpdater) StoreEvent(
 	ctx context.Context, event *gomatrixserverlib.Event,
 	authEventNIDs []types.EventNID, isRejected bool,
 ) (types.EventNID, types.RoomNID, types.StateAtEvent, *gomatrixserverlib.Event, string, error) {
-	return u.d.storeEvent(ctx, u.txn, u, event, authEventNIDs, isRejected)
+	return u.d.storeEvent(ctx, u, event, authEventNIDs, isRejected)
 }
 
 func (u *RoomUpdater) AddState(
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index de78dd9f4..787cde73b 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -497,11 +497,11 @@ func (d *Database) StoreEvent(
 	ctx context.Context, event *gomatrixserverlib.Event,
 	authEventNIDs []types.EventNID, isRejected bool,
 ) (types.EventNID, types.RoomNID, types.StateAtEvent, *gomatrixserverlib.Event, string, error) {
-	return d.storeEvent(ctx, nil, nil, event, authEventNIDs, isRejected)
+	return d.storeEvent(ctx, nil, event, authEventNIDs, isRejected)
 }
 
 func (d *Database) storeEvent(
-	ctx context.Context, txn *sql.Tx, updater *RoomUpdater, event *gomatrixserverlib.Event,
+	ctx context.Context, updater *RoomUpdater, event *gomatrixserverlib.Event,
 	authEventNIDs []types.EventNID, isRejected bool,
 ) (types.EventNID, types.RoomNID, types.StateAtEvent, *gomatrixserverlib.Event, string, error) {
 	var (
@@ -514,7 +514,10 @@ func (d *Database) storeEvent(
 		redactedEventID  string
 		err              error
 	)
-
+	var txn *sql.Tx
+	if updater != nil {
+		txn = updater.txn
+	}
 	err = d.Writer.Do(d.DB, txn, func(txn *sql.Tx) error {
 		// TODO: Here we should aim to have two different code paths for new rooms
 		// vs existing ones.
@@ -610,22 +613,10 @@ func (d *Database) storeEvent(
 			}
 			defer sqlutil.EndTransactionWithCheck(updater.txn, &succeeded, &err)
 		}
-		// Ensure that we atomically store prev events AND commit them. If we don't wrap StorePreviousEvents
-		// and EndTransaction in a writer then it's possible for a new write txn to be made between the two
-		// function calls which will then fail with 'database is locked'. This new write txn would HAVE to be
-		// something like SetRoomAlias/RemoveRoomAlias as normal input events are already done sequentially due to
-		// SupportsConcurrentRoomInputs() == false on sqlite, though this does not apply to setting room aliases
-		// as they don't go via InputRoomEvents
-		err = d.Writer.Do(d.DB, updater.txn, func(txn *sql.Tx) error {
-			if err = updater.StorePreviousEvents(eventNID, prevEvents); err != nil {
-				return fmt.Errorf("updater.StorePreviousEvents: %w", err)
-			}
-			succeeded = true
-			return nil
-		})
-		if err != nil {
-			return 0, 0, types.StateAtEvent{}, nil, "", err
+		if err = updater.StorePreviousEvents(eventNID, prevEvents); err != nil {
+			return 0, 0, types.StateAtEvent{}, nil, "", fmt.Errorf("updater.StorePreviousEvents: %w", err)
 		}
+		succeeded = true
 	}
 
 	return eventNID, roomNID, types.StateAtEvent{