From 93f996bda3a8d15aeac92d297d9dacbef113e5ea Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 18 Aug 2020 14:02:27 +0100
Subject: [PATCH] Juggle around things a bit

---
 .../shared/room_recent_events_updater.go      | 26 ++++++++-------
 roomserver/storage/shared/storage.go          | 32 +++++++++----------
 2 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/roomserver/storage/shared/room_recent_events_updater.go b/roomserver/storage/shared/room_recent_events_updater.go
index 8131f712d..2798df037 100644
--- a/roomserver/storage/shared/room_recent_events_updater.go
+++ b/roomserver/storage/shared/room_recent_events_updater.go
@@ -3,6 +3,7 @@ package shared
 import (
 	"context"
 	"database/sql"
+	"fmt"
 
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -20,25 +21,28 @@ type roomRecentEventsUpdater struct {
 func NewRoomRecentEventsUpdater(d *Database, ctx context.Context, roomNID types.RoomNID, useTxns bool) (types.RoomRecentEventsUpdater, error) {
 	txn, err := d.DB.Begin()
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("d.DB.Begin: %w", err)
 	}
 	eventNIDs, lastEventNIDSent, currentStateSnapshotNID, err :=
 		d.RoomsTable.SelectLatestEventsNIDsForUpdate(ctx, txn, roomNID)
-	if err != nil {
+	if err != nil && err != sql.ErrNoRows {
 		txn.Rollback() // nolint: errcheck
-		return nil, err
-	}
-	stateAndRefs, err := d.EventsTable.BulkSelectStateAtEventAndReference(ctx, txn, eventNIDs)
-	if err != nil {
-		txn.Rollback() // nolint: errcheck
-		return nil, err
+		return nil, fmt.Errorf("d.RoomsTable.SelectLatestEventsNIDsForUpdate: %w", err)
 	}
+	var stateAndRefs []types.StateAtEventAndReference
 	var lastEventIDSent string
-	if lastEventNIDSent != 0 {
-		lastEventIDSent, err = d.EventsTable.SelectEventID(ctx, txn, lastEventNIDSent)
+	if err == nil {
+		stateAndRefs, err = d.EventsTable.BulkSelectStateAtEventAndReference(ctx, txn, eventNIDs)
 		if err != nil {
 			txn.Rollback() // nolint: errcheck
-			return nil, err
+			return nil, fmt.Errorf("d.EventsTable.BulkSelectStateAtEventAndReference: %w", err)
+		}
+		if lastEventNIDSent != 0 {
+			lastEventIDSent, err = d.EventsTable.SelectEventID(ctx, txn, lastEventNIDSent)
+			if err != nil {
+				txn.Rollback() // nolint: errcheck
+				return nil, fmt.Errorf("d.EventsTable.SelectEventID: %w", err)
+			}
 		}
 	}
 	if !useTxns {
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index c3c7567a8..d0bd7f076 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -357,16 +357,16 @@ func (d *Database) StoreEvent(
 		err              error
 	)
 
-	err = sqlutil.WithTransaction(d.DB, func(txn *sql.Tx) error {
-		if txnAndSessionID != nil {
-			if err = d.TransactionsTable.InsertTransaction(
-				ctx, txn, txnAndSessionID.TransactionID,
-				txnAndSessionID.SessionID, event.Sender(), event.EventID(),
-			); err != nil {
-				return fmt.Errorf("d.TransactionsTable.InsertTransaction: %w", err)
-			}
+	if txnAndSessionID != nil {
+		if err = d.TransactionsTable.InsertTransaction(
+			ctx, nil, txnAndSessionID.TransactionID,
+			txnAndSessionID.SessionID, event.Sender(), event.EventID(),
+		); err != nil {
+			return 0, types.StateAtEvent{}, nil, "", fmt.Errorf("d.TransactionsTable.InsertTransaction: %w", err)
 		}
+	}
 
+	err = sqlutil.WithTransaction(d.DB, func(txn *sql.Tx) error {
 		// TODO: Here we should aim to have two different code paths for new rooms
 		// vs existing ones.
 
@@ -426,6 +426,9 @@ func (d *Database) StoreEvent(
 		}
 
 		redactionEvent, redactedEventID, err = d.handleRedactions(ctx, txn, eventNID, event)
+		if err != nil {
+			return fmt.Errorf("d.handleRedactions: %w", err)
+		}
 		return nil
 	})
 	if err != nil {
@@ -456,15 +459,12 @@ func (d *Database) assignRoomNID(
 	ctx context.Context, txn *sql.Tx,
 	roomID string, roomVersion gomatrixserverlib.RoomVersion,
 ) (types.RoomNID, error) {
-	// Check if we already have a numeric ID in the database.
-	roomNID, err := d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
+	// Try to insert the room. If it's ineffectual then we have the
+	// room in the database already.
+	roomNID, err := d.RoomsTable.InsertRoomNID(ctx, txn, roomID, roomVersion)
 	if err == sql.ErrNoRows {
-		// We don't have a numeric ID so insert one into the database.
-		roomNID, err = d.RoomsTable.InsertRoomNID(ctx, txn, roomID, roomVersion)
-		if err == sql.ErrNoRows {
-			// We raced with another insert so run the select again.
-			roomNID, err = d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
-		}
+		// Look up the room NID that we already have.
+		roomNID, err = d.RoomsTable.SelectRoomNID(ctx, txn, roomID)
 	}
 	return roomNID, err
 }