From fafee7060f718c0402069264bbd5f25e0e145144 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 7 Oct 2020 11:28:22 +0100
Subject: [PATCH] Use updater txn

---
 roomserver/internal/input/input_latest_events.go   |  4 ++--
 roomserver/storage/interface.go                    |  5 -----
 roomserver/storage/shared/latest_events_updater.go | 12 ++++++++++++
 roomserver/storage/shared/storage.go               | 10 ----------
 4 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/roomserver/internal/input/input_latest_events.go b/roomserver/internal/input/input_latest_events.go
index 883669538..441187fe7 100644
--- a/roomserver/internal/input/input_latest_events.go
+++ b/roomserver/internal/input/input_latest_events.go
@@ -279,7 +279,7 @@ func (u *latestEventsUpdater) calculateLatest(
 	// now have entries in the previous events table. If they do then they
 	// are no longer forward extremities.
 	for _, l := range oldLatest {
-		referenced, err := u.api.DB.EventIsReferenced(u.ctx, l.EventReference)
+		referenced, err := u.updater.IsReferenced(l.EventReference)
 		if err != nil {
 			logrus.WithError(err).Errorf("Failed to retrieve event reference for %q", l.EventID)
 		}
@@ -300,7 +300,7 @@ func (u *latestEventsUpdater) calculateLatest(
 
 	// If the new event isn't already in the set then we'll check if it
 	// really should be.
-	referenced, err := u.api.DB.EventIsReferenced(u.ctx, newEvent.EventReference)
+	referenced, err := u.updater.IsReferenced(newEvent.EventReference)
 	if err != nil {
 		logrus.WithError(err).Errorf("Failed to retrieve event reference for %q", newEvent.EventReference.EventID)
 	}
diff --git a/roomserver/storage/interface.go b/roomserver/storage/interface.go
index ae51a4a41..10a380e85 100644
--- a/roomserver/storage/interface.go
+++ b/roomserver/storage/interface.go
@@ -87,11 +87,6 @@ type Database interface {
 	// Lookup the event IDs for a batch of event numeric IDs.
 	// Returns an error if the retrieval went wrong.
 	EventIDs(ctx context.Context, eventNIDs []types.EventNID) (map[types.EventNID]string, error)
-	// EventIsReferenced returns true if the event is referenced by another event and false otherwise.
-	// This is used when working out if an event is a new forward extremity or not.
-	EventIsReferenced(
-		ctx context.Context, eventRef gomatrixserverlib.EventReference,
-	) (bool, error)
 	// Look up the latest events in a room in preparation for an update.
 	// The RoomRecentEventsUpdater must have Commit or Rollback called on it if this doesn't return an error.
 	// Returns the latest events in the room and the last eventID sent to the log along with an updater.
diff --git a/roomserver/storage/shared/latest_events_updater.go b/roomserver/storage/shared/latest_events_updater.go
index cd01b1271..29eab0c98 100644
--- a/roomserver/storage/shared/latest_events_updater.go
+++ b/roomserver/storage/shared/latest_events_updater.go
@@ -75,6 +75,18 @@ func (u *LatestEventsUpdater) StorePreviousEvents(eventNID types.EventNID, previ
 	})
 }
 
+// IsReferenced implements types.RoomRecentEventsUpdater
+func (u *LatestEventsUpdater) IsReferenced(eventReference gomatrixserverlib.EventReference) (bool, error) {
+	err := u.d.PrevEventsTable.SelectPreviousEventExists(u.ctx, u.txn, eventReference.EventID, eventReference.EventSHA256)
+	if err == nil {
+		return true, nil
+	}
+	if err == sql.ErrNoRows {
+		return false, nil
+	}
+	return false, fmt.Errorf("u.d.PrevEventsTable.SelectPreviousEventExists: %w", err)
+}
+
 // SetLatestEvents implements types.RoomRecentEventsUpdater
 func (u *LatestEventsUpdater) SetLatestEvents(
 	roomNID types.RoomNID, latest []types.StateAtEventAndReference, lastEventNIDSent types.EventNID,
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 0dbc0186a..f8e733ab7 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -186,16 +186,6 @@ func (d *Database) EventIDs(
 	return d.EventsTable.BulkSelectEventID(ctx, eventNIDs)
 }
 
-func (d *Database) EventIsReferenced(
-	ctx context.Context, eventRef gomatrixserverlib.EventReference,
-) (bool, error) {
-	err := d.PrevEventsTable.SelectPreviousEventExists(ctx, nil, eventRef.EventID, eventRef.EventSHA256)
-	if err != nil && err != sql.ErrNoRows {
-		return false, err
-	}
-	return err != sql.ErrNoRows, nil
-}
-
 func (d *Database) EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error) {
 	nidMap, err := d.EventNIDs(ctx, eventIDs)
 	if err != nil {