Merge Updater structs (#1069)

* Move Updater structs to shared and use it for postgres

* Add constructors for NewXXXUpdater and a useTxns flag

In sqlite, we set useTxns=false and comment why.

* Handle nil txn

* Handle nil in transaction

* Missed one

* Close the txn at the right time

* Don't close the transaction as we reuse it between calls

roomserver/storage/shared/membership_updater.go

@@ -0,0 +1,183 @@
+package shared
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+type membershipUpdater struct {
+	transaction
+	d             *Database
+	roomNID       types.RoomNID
+	targetUserNID types.EventStateKeyNID
+	membership    tables.MembershipState
+}
+
+func NewMembershipUpdater(
+	ctx context.Context, d *Database, roomID, targetUserID string,
+	targetLocal bool, roomVersion gomatrixserverlib.RoomVersion,
+	useTxns bool,
+) (types.MembershipUpdater, error) {
+	txn, err := d.DB.Begin()
+	if err != nil {
+		return nil, err
+	}
+	succeeded := false
+	defer func() {
+		if !succeeded {
+			txn.Rollback() // nolint: errcheck
+		}
+	}()
+
+	roomNID, err := d.assignRoomNID(ctx, txn, roomID, roomVersion)
+	if err != nil {
+		return nil, err
+	}
+
+	targetUserNID, err := d.assignStateKeyNID(ctx, txn, targetUserID)
+	if err != nil {
+		return nil, err
+	}
+
+	updater, err := d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID, targetLocal)
+	if err != nil {
+		return nil, err
+	}
+
+	succeeded = true
+	if !useTxns {
+		txn.Commit() // nolint: errcheck
+		updater.transaction.txn = nil
+	}
+	return updater, nil
+}
+
+func (d *Database) membershipUpdaterTxn(
+	ctx context.Context,
+	txn *sql.Tx,
+	roomNID types.RoomNID,
+	targetUserNID types.EventStateKeyNID,
+	targetLocal bool,
+) (*membershipUpdater, error) {
+
+	if err := d.MembershipTable.InsertMembership(ctx, txn, roomNID, targetUserNID, targetLocal); err != nil {
+		return nil, err
+	}
+
+	membership, err := d.MembershipTable.SelectMembershipForUpdate(ctx, txn, roomNID, targetUserNID)
+	if err != nil {
+		return nil, err
+	}
+
+	return &membershipUpdater{
+		transaction{ctx, txn}, d, roomNID, targetUserNID, membership,
+	}, nil
+}
+
+// IsInvite implements types.MembershipUpdater
+func (u *membershipUpdater) IsInvite() bool {
+	return u.membership == tables.MembershipStateInvite
+}
+
+// IsJoin implements types.MembershipUpdater
+func (u *membershipUpdater) IsJoin() bool {
+	return u.membership == tables.MembershipStateJoin
+}
+
+// IsLeave implements types.MembershipUpdater
+func (u *membershipUpdater) IsLeave() bool {
+	return u.membership == tables.MembershipStateLeaveOrBan
+}
+
+// SetToInvite implements types.MembershipUpdater
+func (u *membershipUpdater) SetToInvite(event gomatrixserverlib.Event) (bool, error) {
+	senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, event.Sender())
+	if err != nil {
+		return false, err
+	}
+	inserted, err := u.d.InvitesTable.InsertInviteEvent(
+		u.ctx, u.txn, event.EventID(), u.roomNID, u.targetUserNID, senderUserNID, event.JSON(),
+	)
+	if err != nil {
+		return false, err
+	}
+	if u.membership != tables.MembershipStateInvite {
+		if err = u.d.MembershipTable.UpdateMembership(
+			u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateInvite, 0,
+		); err != nil {
+			return false, err
+		}
+	}
+	return inserted, nil
+}
+
+// SetToJoin implements types.MembershipUpdater
+func (u *membershipUpdater) SetToJoin(senderUserID string, eventID string, isUpdate bool) ([]string, error) {
+	var inviteEventIDs []string
+
+	senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, senderUserID)
+	if err != nil {
+		return nil, err
+	}
+
+	// If this is a join event update, there is no invite to update
+	if !isUpdate {
+		inviteEventIDs, err = u.d.InvitesTable.UpdateInviteRetired(
+			u.ctx, u.txn, u.roomNID, u.targetUserNID,
+		)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// Look up the NID of the new join event
+	nIDs, err := u.d.EventNIDs(u.ctx, []string{eventID})
+	if err != nil {
+		return nil, err
+	}
+
+	if u.membership != tables.MembershipStateJoin || isUpdate {
+		if err = u.d.MembershipTable.UpdateMembership(
+			u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID,
+			tables.MembershipStateJoin, nIDs[eventID],
+		); err != nil {
+			return nil, err
+		}
+	}
+
+	return inviteEventIDs, nil
+}
+
+// SetToLeave implements types.MembershipUpdater
+func (u *membershipUpdater) SetToLeave(senderUserID string, eventID string) ([]string, error) {
+	senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, senderUserID)
+	if err != nil {
+		return nil, err
+	}
+	inviteEventIDs, err := u.d.InvitesTable.UpdateInviteRetired(
+		u.ctx, u.txn, u.roomNID, u.targetUserNID,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Look up the NID of the new leave event
+	nIDs, err := u.d.EventNIDs(u.ctx, []string{eventID})
+	if err != nil {
+		return nil, err
+	}
+
+	if u.membership != tables.MembershipStateLeaveOrBan {
+		if err = u.d.MembershipTable.UpdateMembership(
+			u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID,
+			tables.MembershipStateLeaveOrBan, nIDs[eventID],
+		); err != nil {
+			return nil, err
+		}
+	}
+	return inviteEventIDs, nil
+}

roomserver/storage/shared/prepare.go

@@ -16,6 +16,7 @@
 package shared
 
 import (
+	"context"
 	"database/sql"
 )
 
@@ -34,3 +35,26 @@ func (s StatementList) Prepare(db *sql.DB) (err error) {
 	}
 	return
 }
+
+type transaction struct {
+	ctx context.Context
+	txn *sql.Tx
+}
+
+// Commit implements types.Transaction
+func (t *transaction) Commit() error {
+	if t.txn == nil {
+		// The Updater structs can operate in useTxns=false mode. The code will still call this though.
+		return nil
+	}
+	return t.txn.Commit()
+}
+
+// Rollback implements types.Transaction
+func (t *transaction) Rollback() error {
+	if t.txn == nil {
+		// The Updater structs can operate in useTxns=false mode. The code will still call this though.
+		return nil
+	}
+	return t.txn.Rollback()
+}

roomserver/storage/shared/room_recent_events_updater.go

@@ -0,0 +1,120 @@
+package shared
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+type roomRecentEventsUpdater struct {
+	transaction
+	d                       *Database
+	roomNID                 types.RoomNID
+	latestEvents            []types.StateAtEventAndReference
+	lastEventIDSent         string
+	currentStateSnapshotNID types.StateSnapshotNID
+}
+
+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
+	}
+	eventNIDs, lastEventNIDSent, currentStateSnapshotNID, err :=
+		d.RoomsTable.SelectLatestEventsNIDsForUpdate(ctx, txn, roomNID)
+	if err != nil {
+		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
+	}
+	var lastEventIDSent string
+	if lastEventNIDSent != 0 {
+		lastEventIDSent, err = d.EventsTable.SelectEventID(ctx, txn, lastEventNIDSent)
+		if err != nil {
+			txn.Rollback() // nolint: errcheck
+			return nil, err
+		}
+	}
+	if !useTxns {
+		txn.Commit() // nolint: errcheck
+		txn = nil
+	}
+	return &roomRecentEventsUpdater{
+		transaction{ctx, txn}, d, roomNID, stateAndRefs, lastEventIDSent, currentStateSnapshotNID,
+	}, nil
+}
+
+// RoomVersion implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) RoomVersion() (version gomatrixserverlib.RoomVersion) {
+	version, _ = u.d.GetRoomVersionForRoomNID(u.ctx, u.roomNID)
+	return
+}
+
+// LatestEvents implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) LatestEvents() []types.StateAtEventAndReference {
+	return u.latestEvents
+}
+
+// LastEventIDSent implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) LastEventIDSent() string {
+	return u.lastEventIDSent
+}
+
+// CurrentStateSnapshotNID implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) CurrentStateSnapshotNID() types.StateSnapshotNID {
+	return u.currentStateSnapshotNID
+}
+
+// StorePreviousEvents implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) 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 err
+		}
+	}
+	return nil
+}
+
+// IsReferenced implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) 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, err
+}
+
+// SetLatestEvents implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) SetLatestEvents(
+	roomNID types.RoomNID, latest []types.StateAtEventAndReference, lastEventNIDSent types.EventNID,
+	currentStateSnapshotNID types.StateSnapshotNID,
+) error {
+	eventNIDs := make([]types.EventNID, len(latest))
+	for i := range latest {
+		eventNIDs[i] = latest[i].EventNID
+	}
+	return u.d.RoomsTable.UpdateLatestEventNIDs(u.ctx, u.txn, roomNID, eventNIDs, lastEventNIDSent, currentStateSnapshotNID)
+}
+
+// HasEventBeenSent implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) HasEventBeenSent(eventNID types.EventNID) (bool, error) {
+	return u.d.EventsTable.SelectEventSentToOutput(u.ctx, u.txn, eventNID)
+}
+
+// MarkEventAsSent implements types.RoomRecentEventsUpdater
+func (u *roomRecentEventsUpdater) MarkEventAsSent(eventNID types.EventNID) error {
+	return u.d.EventsTable.UpdateEventSentToOutput(u.ctx, u.txn, eventNID)
+}
+
+func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID, targetLocal bool) (types.MembershipUpdater, error) {
+	return u.d.membershipUpdaterTxn(u.ctx, u.txn, u.roomNID, targetUserNID, targetLocal)
+}

roomserver/storage/shared/storage.go

@@ -311,6 +311,19 @@ func (d *Database) GetTransactionEventID(
 	return eventID, err
 }
 
+func (d *Database) MembershipUpdater(
+	ctx context.Context, roomID, targetUserID string,
+	targetLocal bool, roomVersion gomatrixserverlib.RoomVersion,
+) (types.MembershipUpdater, error) {
+	return NewMembershipUpdater(ctx, d, roomID, targetUserID, targetLocal, roomVersion, true)
+}
+
+func (d *Database) GetLatestEventsForUpdate(
+	ctx context.Context, roomNID types.RoomNID,
+) (types.RoomRecentEventsUpdater, error) {
+	return NewRoomRecentEventsUpdater(d, ctx, roomNID, true)
+}
+
 func (d *Database) StoreEvent(
 	ctx context.Context, event gomatrixserverlib.Event,
 	txnAndSessionID *api.TransactionID, authEventNIDs []types.EventNID,