Offset updates take place using TransactionWriter

appservice/storage/postgres/storage.go

@@ -32,6 +32,7 @@ type Database struct {
 	events eventsStatements
 	txnID  txnStatements
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 }
 
 // NewDatabase opens a new database
@@ -41,10 +42,11 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 	if result.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
+	result.writer = sqlutil.NewDummyTransactionWriter()
 	if err = result.prepare(); err != nil {
 		return nil, err
 	}
-	if err = result.PartitionOffsetStatements.Prepare(result.db, "appservice"); err != nil {
+	if err = result.PartitionOffsetStatements.Prepare(result.db, result.writer, "appservice"); err != nil {
 		return nil, err
 	}
 	return &result, nil

appservice/storage/sqlite3/storage.go

@@ -32,6 +32,7 @@ type Database struct {
 	events eventsStatements
 	txnID  txnStatements
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 }
 
 // NewDatabase opens a new database
@@ -41,10 +42,11 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 	if result.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
+	result.writer = sqlutil.NewTransactionWriter()
 	if err = result.prepare(); err != nil {
 		return nil, err
 	}
-	if err = result.PartitionOffsetStatements.Prepare(result.db, "appservice"); err != nil {
+	if err = result.PartitionOffsetStatements.Prepare(result.db, result.writer, "appservice"); err != nil {
 		return nil, err
 	}
 	return &result, nil

currentstateserver/storage/postgres/storage.go

@@ -11,6 +11,7 @@ import (
 type Database struct {
 	shared.Database
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 	sqlutil.PartitionOffsetStatements
 }
 
@@ -21,7 +22,8 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
-	if err = d.PartitionOffsetStatements.Prepare(d.db, "currentstate"); err != nil {
+	d.writer = sqlutil.NewDummyTransactionWriter()
+	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "currentstate"); err != nil {
 		return nil, err
 	}
 	currRoomState, err := NewPostgresCurrentRoomStateTable(d.db)

currentstateserver/storage/sqlite3/storage.go

@@ -11,6 +11,7 @@ import (
 type Database struct {
 	shared.Database
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 	sqlutil.PartitionOffsetStatements
 }
 
@@ -22,7 +23,8 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
-	if err = d.PartitionOffsetStatements.Prepare(d.db, "currentstate"); err != nil {
+	d.writer = sqlutil.NewTransactionWriter()
+	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "currentstate"); err != nil {
 		return nil, err
 	}
 	currRoomState, err := NewSqliteCurrentRoomStateTable(d.db)

federationsender/storage/postgres/storage.go

@@ -28,6 +28,7 @@ type Database struct {
 	shared.Database
 	sqlutil.PartitionOffsetStatements
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 }
 
 // NewDatabase opens a new database
@@ -37,6 +38,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
+	d.writer = sqlutil.NewDummyTransactionWriter()
 	joinedHosts, err := NewPostgresJoinedHostsTable(d.db)
 	if err != nil {
 		return nil, err
@@ -70,7 +72,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 		FederationSenderRooms:       rooms,
 		FederationSenderBlacklist:   blacklist,
 	}
-	if err = d.PartitionOffsetStatements.Prepare(d.db, "federationsender"); err != nil {
+	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "federationsender"); err != nil {
 		return nil, err
 	}
 	return &d, nil

federationsender/storage/sqlite3/storage.go

@@ -30,6 +30,7 @@ type Database struct {
 	shared.Database
 	sqlutil.PartitionOffsetStatements
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 }
 
 // NewDatabase opens a new database
@@ -39,6 +40,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
+	d.writer = sqlutil.NewTransactionWriter()
 	joinedHosts, err := NewSQLiteJoinedHostsTable(d.db)
 	if err != nil {
 		return nil, err
@@ -72,7 +74,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*Database, error) {
 		FederationSenderRooms:       rooms,
 		FederationSenderBlacklist:   blacklist,
 	}
-	if err = d.PartitionOffsetStatements.Prepare(d.db, "federationsender"); err != nil {
+	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "federationsender"); err != nil {
 		return nil, err
 	}
 	return &d, nil

internal/sqlutil/partition_offset_table.go

@@ -53,6 +53,8 @@ const upsertPartitionOffsetsSQL = "" +
 
 // PartitionOffsetStatements represents a set of statements that can be run on a partition_offsets table.
 type PartitionOffsetStatements struct {
+	db                         *sql.DB
+	writer                     TransactionWriter
 	selectPartitionOffsetsStmt *sql.Stmt
 	upsertPartitionOffsetStmt  *sql.Stmt
 }
@@ -60,7 +62,9 @@ type PartitionOffsetStatements struct {
 // Prepare converts the raw SQL statements into prepared statements.
 // Takes a prefix to prepend to the table name used to store the partition offsets.
 // This allows multiple components to share the same database schema.
-func (s *PartitionOffsetStatements) Prepare(db *sql.DB, prefix string) (err error) {
+func (s *PartitionOffsetStatements) Prepare(db *sql.DB, writer TransactionWriter, prefix string) (err error) {
+	s.db = db
+	s.writer = writer
 	_, err = db.Exec(strings.Replace(partitionOffsetsSchema, "${prefix}", prefix, -1))
 	if err != nil {
 		return
@@ -121,6 +125,9 @@ func (s *PartitionOffsetStatements) selectPartitionOffsets(
 func (s *PartitionOffsetStatements) upsertPartitionOffset(
 	ctx context.Context, topic string, partition int32, offset int64,
 ) error {
-	_, err := s.upsertPartitionOffsetStmt.ExecContext(ctx, topic, partition, offset)
+	return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
+		stmt := TxStmt(txn, s.upsertPartitionOffsetStmt)
+		_, err := stmt.ExecContext(ctx, topic, partition, offset)
 		return err
+	})
 }

syncapi/storage/postgres/syncserver.go

@@ -31,6 +31,7 @@ import (
 type SyncServerDatasource struct {
 	shared.Database
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 	sqlutil.PartitionOffsetStatements
 }
 
@@ -41,7 +42,8 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*SyncServerDatasource, e
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
-	if err = d.PartitionOffsetStatements.Prepare(d.db, "syncapi"); err != nil {
+	d.writer = sqlutil.NewDummyTransactionWriter()
+	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "syncapi"); err != nil {
 		return nil, err
 	}
 	accountData, err := NewPostgresAccountDataTable(d.db)

syncapi/storage/sqlite3/syncserver.go

@@ -32,6 +32,7 @@ import (
 type SyncServerDatasource struct {
 	shared.Database
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 	sqlutil.PartitionOffsetStatements
 	streamID streamIDStatements
 }
@@ -44,6 +45,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*SyncServerDatasource, e
 	if d.db, err = sqlutil.Open(dbProperties); err != nil {
 		return nil, err
 	}
+	d.writer = sqlutil.NewTransactionWriter()
 	if err = d.prepare(); err != nil {
 		return nil, err
 	}
@@ -51,7 +53,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*SyncServerDatasource, e
 }
 
 func (d *SyncServerDatasource) prepare() (err error) {
-	if err = d.PartitionOffsetStatements.Prepare(d.db, "syncapi"); err != nil {
+	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "syncapi"); err != nil {
 		return err
 	}
 	if err = d.streamID.prepare(d.db); err != nil {

userapi/storage/accounts/postgres/storage.go

@@ -35,6 +35,7 @@ import (
 // Database represents an account database
 type Database struct {
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
 	sqlutil.PartitionOffsetStatements
 	accounts     accountsStatements
 	profiles     profilesStatements
@@ -49,27 +50,27 @@ func NewDatabase(dbProperties *config.DatabaseOptions, serverName gomatrixserver
 	if err != nil {
 		return nil, err
 	}
-	partitions := sqlutil.PartitionOffsetStatements{}
+	d := &Database{
-	if err = partitions.Prepare(db, "account"); err != nil {
+		serverName: serverName,
+		db:         db,
+		writer:     sqlutil.NewDummyTransactionWriter(),
+	}
+	if err = d.PartitionOffsetStatements.Prepare(db, d.writer, "account"); err != nil {
 		return nil, err
 	}
-	a := accountsStatements{}
+	if err = d.accounts.prepare(db, serverName); err != nil {
-	if err = a.prepare(db, serverName); err != nil {
 		return nil, err
 	}
-	p := profilesStatements{}
+	if err = d.profiles.prepare(db); err != nil {
-	if err = p.prepare(db); err != nil {
 		return nil, err
 	}
-	ac := accountDataStatements{}
+	if err = d.accountDatas.prepare(db); err != nil {
-	if err = ac.prepare(db); err != nil {
 		return nil, err
 	}
-	t := threepidStatements{}
+	if err = d.threepids.prepare(db); err != nil {
-	if err = t.prepare(db); err != nil {
 		return nil, err
 	}
-	return &Database{db, partitions, a, p, ac, t, serverName}, nil
+	return d, nil
 }
 
 // GetAccountByPassword returns the account associated with the given localpart and password.

userapi/storage/accounts/sqlite3/storage.go

@@ -34,6 +34,8 @@ import (
 // Database represents an account database
 type Database struct {
 	db     *sql.DB
+	writer sqlutil.TransactionWriter
+
 	sqlutil.PartitionOffsetStatements
 	accounts     accountsStatements
 	profiles     profilesStatements
@@ -53,35 +55,28 @@ func NewDatabase(dbProperties *config.DatabaseOptions, serverName gomatrixserver
 	if err != nil {
 		return nil, err
 	}
-	partitions := sqlutil.PartitionOffsetStatements{}
+	d := &Database{
-	if err = partitions.Prepare(db, "account"); err != nil {
-		return nil, err
-	}
-	a := accountsStatements{}
-	if err = a.prepare(db, serverName); err != nil {
-		return nil, err
-	}
-	p := profilesStatements{}
-	if err = p.prepare(db); err != nil {
-		return nil, err
-	}
-	ac := accountDataStatements{}
-	if err = ac.prepare(db); err != nil {
-		return nil, err
-	}
-	t := threepidStatements{}
-	if err = t.prepare(db); err != nil {
-		return nil, err
-	}
-	return &Database{
-		db:                        db,
-		PartitionOffsetStatements: partitions,
-		accounts:                  a,
-		profiles:                  p,
-		accountDatas:              ac,
-		threepids:                 t,
 		serverName: serverName,
-	}, nil
+		db:         db,
+		writer:     sqlutil.NewTransactionWriter(),
+	}
+	partitions := sqlutil.PartitionOffsetStatements{}
+	if err = partitions.Prepare(db, d.writer, "account"); err != nil {
+		return nil, err
+	}
+	if err = d.accounts.prepare(db, serverName); err != nil {
+		return nil, err
+	}
+	if err = d.profiles.prepare(db); err != nil {
+		return nil, err
+	}
+	if err = d.accountDatas.prepare(db); err != nil {
+		return nil, err
+	}
+	if err = d.threepids.prepare(db); err != nil {
+		return nil, err
+	}
+	return d, nil
 }
 
 // GetAccountByPassword returns the account associated with the given localpart and password.