From d7ea814fa80ea2aba671be17c7d985d1191fbf6a Mon Sep 17 00:00:00 2001
From: Kegsay <kegan@matrix.org>
Date: Tue, 13 Oct 2020 10:20:27 +0100
Subject: [PATCH] Wrap NewMembershipUpdater in a db writer (#1515)

---
 internal/sqlutil/trace.go                     |  2 +-
 .../storage/shared/membership_updater.go      | 19 ++++++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/internal/sqlutil/trace.go b/internal/sqlutil/trace.go
index 23359b500..0684e92e1 100644
--- a/internal/sqlutil/trace.go
+++ b/internal/sqlutil/trace.go
@@ -93,7 +93,7 @@ func trackGoID(query string) {
 	if strings.HasPrefix(q, "SELECT") {
 		return // SELECTs can go on other goroutines
 	}
-	logrus.Warnf("unsafe goid: SQL executed not on an ExclusiveWriter: %s", q)
+	logrus.Warnf("unsafe goid %d: SQL executed not on an ExclusiveWriter: %s", thisGoID, q)
 }
 
 // Open opens a database specified by its database driver name and a driver-specific data source name,
diff --git a/roomserver/storage/shared/membership_updater.go b/roomserver/storage/shared/membership_updater.go
index 834af6069..7abddd018 100644
--- a/roomserver/storage/shared/membership_updater.go
+++ b/roomserver/storage/shared/membership_updater.go
@@ -22,12 +22,21 @@ func NewMembershipUpdater(
 	ctx context.Context, d *Database, txn *sql.Tx, roomID, targetUserID string,
 	targetLocal bool, roomVersion gomatrixserverlib.RoomVersion,
 ) (*MembershipUpdater, error) {
-	roomNID, err := d.assignRoomNID(ctx, txn, roomID, roomVersion)
-	if err != nil {
-		return nil, err
-	}
+	var roomNID types.RoomNID
+	var targetUserNID types.EventStateKeyNID
+	var err error
+	err = d.Writer.Do(d.DB, txn, func(txn *sql.Tx) error {
+		roomNID, err = d.assignRoomNID(ctx, txn, roomID, roomVersion)
+		if err != nil {
+			return err
+		}
 
-	targetUserNID, err := d.assignStateKeyNID(ctx, txn, targetUserID)
+		targetUserNID, err = d.assignStateKeyNID(ctx, txn, targetUserID)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}