Modify sync transaction behaviour (#2758)

This now uses a transaction per stream, so that errors in one stream
don't propagate to another, and we therefore no longer need to do hacks
to reopen a new transaction after aborting a failed one.

11 files changed, 155 insertions, 91 deletions

diff --git a/syncapi/storage/interface.go b/syncapi/storage/interface.go
index be75f8ad..4a03aca7 100644
--- a/syncapi/storage/interface.go
+++ b/syncapi/storage/interface.go
@@ -29,7 +29,6 @@ import (
 
 type DatabaseTransaction interface {
 	sqlutil.Transaction
-	Reset() (err error)
 	SharedUsers
 
 	MaxStreamPositionForPDUs(ctx context.Context) (types.StreamPosition, error)
diff --git a/syncapi/storage/shared/storage_sync.go b/syncapi/storage/shared/storage_sync.go
index 6cc83ebc..0e19d97d 100644
--- a/syncapi/storage/shared/storage_sync.go
+++ b/syncapi/storage/shared/storage_sync.go
@@ -31,19 +31,6 @@ func (d *DatabaseTransaction) Rollback() error {
 	return d.txn.Rollback()
 }
 
-func (d *DatabaseTransaction) Reset() (err error) {
-	if d.txn == nil {
-		return nil
-	}
-	if err = d.txn.Rollback(); err != nil {
-		return err
-	}
-	if d.txn, err = d.DB.BeginTx(d.ctx, nil); err != nil {
-		return err
-	}
-	return
-}
-
 func (d *DatabaseTransaction) MaxStreamPositionForPDUs(ctx context.Context) (types.StreamPosition, error) {
 	id, err := d.OutputEvents.SelectMaxEventID(ctx, d.txn)
 	if err != nil {
diff --git a/syncapi/streams/stream_accountdata.go b/syncapi/streams/stream_accountdata.go
index f3e7fbda..3f2f7d13 100644
--- a/syncapi/streams/stream_accountdata.go
+++ b/syncapi/streams/stream_accountdata.go
@@ -54,7 +54,6 @@ func (p *AccountDataStreamProvider) IncrementalSync(
 	)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.GetAccountDataInRange failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/streams/stream_devicelist.go b/syncapi/streams/stream_devicelist.go
index 307099b8..7996c203 100644
--- a/syncapi/streams/stream_devicelist.go
+++ b/syncapi/streams/stream_devicelist.go
@@ -34,13 +34,11 @@ func (p *DeviceListStreamProvider) IncrementalSync(
 	to, _, err = internal.DeviceListCatchup(context.Background(), snapshot, p.keyAPI, p.rsAPI, req.Device.UserID, req.Response, from, to)
 	if err != nil {
 		req.Log.WithError(err).Error("internal.DeviceListCatchup failed")
-		_ = snapshot.Reset()
 		return from
 	}
 	err = internal.DeviceOTKCounts(req.Context, p.keyAPI, req.Device.UserID, req.Device.ID, req.Response)
 	if err != nil {
 		req.Log.WithError(err).Error("internal.DeviceListCatchup failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/streams/stream_invite.go b/syncapi/streams/stream_invite.go
index 4c889b8f..17b3b843 100644
--- a/syncapi/streams/stream_invite.go
+++ b/syncapi/streams/stream_invite.go
@@ -56,7 +56,6 @@ func (p *InviteStreamProvider) IncrementalSync(
 	)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.InviteEventsInRange failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/streams/stream_notificationdata.go b/syncapi/streams/stream_notificationdata.go
index 5154dd33..5a81fd09 100644
--- a/syncapi/streams/stream_notificationdata.go
+++ b/syncapi/streams/stream_notificationdata.go
@@ -46,7 +46,6 @@ func (p *NotificationDataStreamProvider) IncrementalSync(
 	countsByRoom, err := snapshot.GetUserUnreadNotificationCountsForRooms(ctx, req.Device.UserID, req.Rooms)
 	if err != nil {
 		req.Log.WithError(err).Error("GetUserUnreadNotificationCountsForRooms failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/streams/stream_pdu.go b/syncapi/streams/stream_pdu.go
index 01ddf9ac..d252265f 100644
--- a/syncapi/streams/stream_pdu.go
+++ b/syncapi/streams/stream_pdu.go
@@ -75,7 +75,6 @@ func (p *PDUStreamProvider) CompleteSync(
 	joinedRoomIDs, err := snapshot.RoomIDsWithMembership(ctx, req.Device.UserID, gomatrixserverlib.Join)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.RoomIDsWithMembership failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
@@ -102,10 +101,10 @@ func (p *PDUStreamProvider) CompleteSync(
 		)
 		if jerr != nil {
 			req.Log.WithError(jerr).Error("p.getJoinResponseForCompleteSync failed")
-			if err = snapshot.Reset(); err != nil {
+			if err == context.DeadlineExceeded || err == context.Canceled || err == sql.ErrTxDone {
 				return from
 			}
-			continue // return from
+			continue
 		}
 		req.Response.Rooms.Join[roomID] = *jr
 		req.Rooms[roomID] = gomatrixserverlib.Join
@@ -115,7 +114,6 @@ func (p *PDUStreamProvider) CompleteSync(
 	peeks, err := snapshot.PeeksInRange(ctx, req.Device.UserID, req.Device.ID, r)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.PeeksInRange failed")
-		_ = snapshot.Reset()
 		return from
 	}
 	for _, peek := range peeks {
@@ -126,10 +124,10 @@ func (p *PDUStreamProvider) CompleteSync(
 			)
 			if err != nil {
 				req.Log.WithError(err).Error("p.getJoinResponseForCompleteSync failed")
-				if err = snapshot.Reset(); err != nil {
+				if err == context.DeadlineExceeded || err == context.Canceled || err == sql.ErrTxDone {
 					return from
 				}
-				continue // return from
+				continue
 			}
 			req.Response.Rooms.Peek[peek.RoomID] = *jr
 		}
@@ -160,14 +158,12 @@ func (p *PDUStreamProvider) IncrementalSync(
 	if req.WantFullState {
 		if stateDeltas, syncJoinedRooms, err = snapshot.GetStateDeltasForFullStateSync(ctx, req.Device, r, req.Device.UserID, &stateFilter); err != nil {
 			req.Log.WithError(err).Error("p.DB.GetStateDeltasForFullStateSync failed")
-			_ = snapshot.Reset()
-			return
+			return from
 		}
 	} else {
 		if stateDeltas, syncJoinedRooms, err = snapshot.GetStateDeltas(ctx, req.Device, r, req.Device.UserID, &stateFilter); err != nil {
 			req.Log.WithError(err).Error("p.DB.GetStateDeltas failed")
-			_ = snapshot.Reset()
-			return
+			return from
 		}
 	}
 
@@ -181,7 +177,6 @@ func (p *PDUStreamProvider) IncrementalSync(
 
 	if err = p.addIgnoredUsersToFilter(ctx, snapshot, req, &eventFilter); err != nil {
 		req.Log.WithError(err).Error("unable to update event filter with ignored users")
-		_ = snapshot.Reset()
 	}
 
 	newPos = from
@@ -201,13 +196,10 @@ func (p *PDUStreamProvider) IncrementalSync(
 		var pos types.StreamPosition
 		if pos, err = p.addRoomDeltaToResponse(ctx, snapshot, req.Device, newRange, delta, &eventFilter, &stateFilter, req.Response); err != nil {
 			req.Log.WithError(err).Error("d.addRoomDeltaToResponse failed")
-			if err == context.DeadlineExceeded || err == context.Canceled {
+			if err == context.DeadlineExceeded || err == context.Canceled || err == sql.ErrTxDone {
 				return newPos
 			}
-			if err = snapshot.Reset(); err != nil {
-				return from
-			}
-			continue // return to
+			continue
 		}
 		// Reset the position, as it is only for the special case of newly joined rooms
 		if delta.NewlyJoined {
@@ -307,7 +299,6 @@ func (p *PDUStreamProvider) addRoomDeltaToResponse(
 	events, err := applyHistoryVisibilityFilter(ctx, snapshot, p.rsAPI, delta.RoomID, device.UserID, eventFilter.Limit, recentEvents)
 	if err != nil {
 		logrus.WithError(err).Error("unable to apply history visibility filter")
-		_ = snapshot.Reset()
 	}
 
 	if len(delta.StateEvents) > 0 {
diff --git a/syncapi/streams/stream_presence.go b/syncapi/streams/stream_presence.go
index 8a3f01c2..8b87af45 100644
--- a/syncapi/streams/stream_presence.go
+++ b/syncapi/streams/stream_presence.go
@@ -67,7 +67,6 @@ func (p *PresenceStreamProvider) IncrementalSync(
 	presences, err := snapshot.PresenceAfter(ctx, from, gomatrixserverlib.EventFilter{Limit: 1000})
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.PresenceAfter failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/streams/stream_receipt.go b/syncapi/streams/stream_receipt.go
index 79fd65bf..8818a553 100644
--- a/syncapi/streams/stream_receipt.go
+++ b/syncapi/streams/stream_receipt.go
@@ -52,7 +52,6 @@ func (p *ReceiptStreamProvider) IncrementalSync(
 	lastPos, receipts, err := snapshot.RoomReceiptsAfter(ctx, joinedRooms, from)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.RoomReceiptsAfter failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/streams/stream_sendtodevice.go b/syncapi/streams/stream_sendtodevice.go
index c79efad0..00b67cc4 100644
--- a/syncapi/streams/stream_sendtodevice.go
+++ b/syncapi/streams/stream_sendtodevice.go
@@ -44,7 +44,6 @@ func (p *SendToDeviceStreamProvider) IncrementalSync(
 	lastPos, events, err := snapshot.SendToDeviceUpdatesForSync(req.Context, req.Device.UserID, req.Device.ID, from, to)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.SendToDeviceUpdatesForSync failed")
-		_ = snapshot.Reset()
 		return from
 	}
 
diff --git a/syncapi/sync/requestpool.go b/syncapi/sync/requestpool.go
index bef0a568..a71d32ab 100644
--- a/syncapi/sync/requestpool.go
+++ b/syncapi/sync/requestpool.go
@@ -306,13 +306,19 @@ func (rp *RequestPool) OnIncomingSyncRequest(req *http.Request, device *userapi.
 			syncReq.Log.WithField("currentPos", currentPos).Debugln("Responding to sync immediately")
 		}
 
-		snapshot, err := rp.db.NewDatabaseSnapshot(req.Context())
-		if err != nil {
-			logrus.WithError(err).Error("Failed to acquire database snapshot for sync request")
-			return jsonerror.InternalServerError()
+		withTransaction := func(from types.StreamPosition, f func(snapshot storage.DatabaseTransaction) types.StreamPosition) types.StreamPosition {
+			var succeeded bool
+			snapshot, err := rp.db.NewDatabaseSnapshot(req.Context())
+			if err != nil {
+				logrus.WithError(err).Error("Failed to acquire database snapshot for sync request")
+				return from
+			}
+			defer func() {
+				succeeded = err == nil
+				sqlutil.EndTransactionWithCheck(snapshot, &succeeded, &err)
+			}()
+			return f(snapshot)
 		}
-		var succeeded bool
-		defer sqlutil.EndTransactionWithCheck(snapshot, &succeeded, &err)
 
 		if syncReq.Since.IsEmpty() {
 			// Complete sync
@@ -320,72 +326,162 @@ func (rp *RequestPool) OnIncomingSyncRequest(req *http.Request, device *userapi.
 				// Get the current DeviceListPosition first, as the currentPosition
 				// might advance while processing other streams, resulting in flakey
 				// tests.
-				DeviceListPosition: rp.streams.DeviceListStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				DeviceListPosition: withTransaction(
+					syncReq.Since.DeviceListPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.DeviceListStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				PDUPosition: rp.streams.PDUStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				PDUPosition: withTransaction(
+					syncReq.Since.PDUPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.PDUStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				TypingPosition: rp.streams.TypingStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				TypingPosition: withTransaction(
+					syncReq.Since.TypingPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.TypingStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				ReceiptPosition: rp.streams.ReceiptStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				ReceiptPosition: withTransaction(
+					syncReq.Since.ReceiptPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.ReceiptStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				InvitePosition: rp.streams.InviteStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				InvitePosition: withTransaction(
+					syncReq.Since.InvitePosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.InviteStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				SendToDevicePosition: rp.streams.SendToDeviceStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				SendToDevicePosition: withTransaction(
+					syncReq.Since.SendToDevicePosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.SendToDeviceStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				AccountDataPosition: rp.streams.AccountDataStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				AccountDataPosition: withTransaction(
+					syncReq.Since.AccountDataPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.AccountDataStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				NotificationDataPosition: rp.streams.NotificationDataStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				NotificationDataPosition: withTransaction(
+					syncReq.Since.NotificationDataPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.NotificationDataStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
-				PresencePosition: rp.streams.PresenceStreamProvider.CompleteSync(
-					syncReq.Context, snapshot, syncReq,
+				PresencePosition: withTransaction(
+					syncReq.Since.PresencePosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.PresenceStreamProvider.CompleteSync(
+							syncReq.Context, txn, syncReq,
+						)
+					},
 				),
 			}
 		} else {
 			// Incremental sync
 			syncReq.Response.NextBatch = types.StreamingToken{
-				PDUPosition: rp.streams.PDUStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.PDUPosition, currentPos.PDUPosition,
+				PDUPosition: withTransaction(
+					syncReq.Since.PDUPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.PDUStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.PDUPosition, currentPos.PDUPosition,
+						)
+					},
 				),
-				TypingPosition: rp.streams.TypingStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.TypingPosition, currentPos.TypingPosition,
+				TypingPosition: withTransaction(
+					syncReq.Since.TypingPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.TypingStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.TypingPosition, currentPos.TypingPosition,
+						)
+					},
 				),
-				ReceiptPosition: rp.streams.ReceiptStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.ReceiptPosition, currentPos.ReceiptPosition,
+				ReceiptPosition: withTransaction(
+					syncReq.Since.ReceiptPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.ReceiptStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.ReceiptPosition, currentPos.ReceiptPosition,
+						)
+					},
 				),
-				InvitePosition: rp.streams.InviteStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.InvitePosition, currentPos.InvitePosition,
+				InvitePosition: withTransaction(
+					syncReq.Since.InvitePosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.InviteStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.InvitePosition, currentPos.InvitePosition,
+						)
+					},
 				),
-				SendToDevicePosition: rp.streams.SendToDeviceStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.SendToDevicePosition, currentPos.SendToDevicePosition,
+				SendToDevicePosition: withTransaction(
+					syncReq.Since.SendToDevicePosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.SendToDeviceStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.SendToDevicePosition, currentPos.SendToDevicePosition,
+						)
+					},
 				),
-				AccountDataPosition: rp.streams.AccountDataStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.AccountDataPosition, currentPos.AccountDataPosition,
+				AccountDataPosition: withTransaction(
+					syncReq.Since.AccountDataPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.AccountDataStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.AccountDataPosition, currentPos.AccountDataPosition,
+						)
+					},
 				),
-				NotificationDataPosition: rp.streams.NotificationDataStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.NotificationDataPosition, currentPos.NotificationDataPosition,
+				NotificationDataPosition: withTransaction(
+					syncReq.Since.NotificationDataPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.NotificationDataStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.NotificationDataPosition, currentPos.NotificationDataPosition,
+						)
+					},
 				),
-				DeviceListPosition: rp.streams.DeviceListStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.DeviceListPosition, currentPos.DeviceListPosition,
+				DeviceListPosition: withTransaction(
+					syncReq.Since.DeviceListPosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.DeviceListStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.DeviceListPosition, currentPos.DeviceListPosition,
+						)
+					},
 				),
-				PresencePosition: rp.streams.PresenceStreamProvider.IncrementalSync(
-					syncReq.Context, snapshot, syncReq,
-					syncReq.Since.PresencePosition, currentPos.PresencePosition,
+				PresencePosition: withTransaction(
+					syncReq.Since.PresencePosition,
+					func(txn storage.DatabaseTransaction) types.StreamPosition {
+						return rp.streams.PresenceStreamProvider.IncrementalSync(
+							syncReq.Context, txn, syncReq,
+							syncReq.Since.PresencePosition, currentPos.PresencePosition,
+						)
+					},
 				),
 			}
 			// it's possible for there to be no updates for this user even though since < current pos,
@@ -411,7 +507,6 @@ func (rp *RequestPool) OnIncomingSyncRequest(req *http.Request, device *userapi.
 			}
 		}
 
-		succeeded = true
 		return util.JSONResponse{
 			Code: http.StatusOK,
 			JSON: syncReq.Response,