Use a writer to assign state key NIDs (#2877)

roomserver/storage/shared/storage.go

@@ -103,6 +103,7 @@ func (d *Database) eventStateKeyNIDs(
 	ctx context.Context, txn *sql.Tx, eventStateKeys []string,
 ) (map[string]types.EventStateKeyNID, error) {
 	result := make(map[string]types.EventStateKeyNID)
+	eventStateKeys = util.UniqueStrings(eventStateKeys)
 	nids, err := d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, txn, eventStateKeys)
 	if err != nil {
 		return nil, err
@@ -112,15 +113,23 @@ func (d *Database) eventStateKeyNIDs(
 	}
 	// We received some nids, but are still missing some, work out which and create them
 	if len(eventStateKeys) > len(result) {
-		for _, eventStateKey := range eventStateKeys {
+		var nid types.EventStateKeyNID
-			if _, ok := result[eventStateKey]; ok {
+		err = d.Writer.Do(d.DB, txn, func(txn *sql.Tx) error {
-				continue
+			for _, eventStateKey := range eventStateKeys {
+				if _, ok := result[eventStateKey]; ok {
+					continue
+				}
+
+				nid, err = d.assignStateKeyNID(ctx, txn, eventStateKey)
+				if err != nil {
+					return err
+				}
+				result[eventStateKey] = nid
 			}
-			nid, err := d.assignStateKeyNID(ctx, txn, eventStateKey)
+			return nil
-			if err != nil {
+		})
-				return result, err
+		if err != nil {
-			}
+			return nil, err
-			result[eventStateKey] = nid
 		}
 	}
 	return result, nil