Simplify send-to-device messaging (#1702)

* Simplify send-to-device messaging

* Don't return error if there's no work to do

* Remove SQLite migrations for now

* Tweak Postgres migrations

* Tweaks

* Fixes

* Cleanup separately

* Fix SQLite migration

13 files changed, 248 insertions, 254 deletions

diff --git a/syncapi/storage/interface.go b/syncapi/storage/interface.go
index d66e9964..a51ab4e0 100644
--- a/syncapi/storage/interface.go
+++ b/syncapi/storage/interface.go
@@ -33,6 +33,7 @@ type Database interface {
 	MaxStreamPositionForReceipts(ctx context.Context) (types.StreamPosition, error)
 	MaxStreamPositionForInvites(ctx context.Context) (types.StreamPosition, error)
 	MaxStreamPositionForAccountData(ctx context.Context) (types.StreamPosition, error)
+	MaxStreamPositionForSendToDeviceMessages(ctx context.Context) (types.StreamPosition, error)
 
 	CurrentState(ctx context.Context, roomID string, stateFilterPart *gomatrixserverlib.StateFilter) ([]*gomatrixserverlib.HeaderedEvent, error)
 	GetStateDeltasForFullStateSync(ctx context.Context, device *userapi.Device, r types.Range, userID string, stateFilter *gomatrixserverlib.StateFilter) ([]types.StateDelta, []string, error)
@@ -117,26 +118,14 @@ type Database interface {
 	// matches the streamevent.transactionID device then the transaction ID gets
 	// added to the unsigned section of the output event.
 	StreamEventsToEvents(device *userapi.Device, in []types.StreamEvent) []*gomatrixserverlib.HeaderedEvent
-	// SendToDeviceUpdatesForSync returns a list of send-to-device updates. It returns three lists:
-	// - "events": a list of send-to-device events that should be included in the sync
-	// - "changes": a list of send-to-device events that should be updated in the database by
-	//      CleanSendToDeviceUpdates
-	// - "deletions": a list of send-to-device events which have been confirmed as sent and
-	//      can be deleted altogether by CleanSendToDeviceUpdates
-	// The token supplied should be the current requested sync token, e.g. from the "since"
-	// parameter.
-	SendToDeviceUpdatesForSync(ctx context.Context, userID, deviceID string, token types.StreamingToken) (pos types.StreamPosition, events []types.SendToDeviceEvent, changes []types.SendToDeviceNID, deletions []types.SendToDeviceNID, err error)
+	// SendToDeviceUpdatesForSync returns a list of send-to-device updates. It returns the
+	// relevant events within the given ranges for the supplied user ID and device ID.
+	SendToDeviceUpdatesForSync(ctx context.Context, userID, deviceID string, from, to types.StreamPosition) (pos types.StreamPosition, events []types.SendToDeviceEvent, err error)
 	// StoreNewSendForDeviceMessage stores a new send-to-device event for a user's device.
 	StoreNewSendForDeviceMessage(ctx context.Context, userID, deviceID string, event gomatrixserverlib.SendToDeviceEvent) (types.StreamPosition, error)
-	// CleanSendToDeviceUpdates will update or remove any send-to-device updates based on the
-	// result to a previous call to SendDeviceUpdatesForSync. This is separate as it allows
-	// SendToDeviceUpdatesForSync to be called multiple times if needed (e.g. before and after
-	// starting to wait for an incremental sync with timeout).
-	// The token supplied should be the current requested sync token, e.g. from the "since"
-	// parameter.
-	CleanSendToDeviceUpdates(ctx context.Context, toUpdate, toDelete []types.SendToDeviceNID, token types.StreamingToken) (err error)
-	// SendToDeviceUpdatesWaiting returns true if there are send-to-device updates waiting to be sent.
-	SendToDeviceUpdatesWaiting(ctx context.Context, userID, deviceID string) (bool, error)
+	// CleanSendToDeviceUpdates removes all send-to-device messages BEFORE the specified
+	// from position, preventing the send-to-device table from growing indefinitely.
+	CleanSendToDeviceUpdates(ctx context.Context, userID, deviceID string, before types.StreamPosition) (err error)
 	// GetFilter looks up the filter associated with a given local user and filter ID.
 	// Returns a filter structure. Otherwise returns an error if no such filter exists
 	// or if there was an error talking to the database.
diff --git a/syncapi/storage/postgres/deltas/20201211125500_sequences.go b/syncapi/storage/postgres/deltas/20201211125500_sequences.go
index a51df26f..7db524da 100644
--- a/syncapi/storage/postgres/deltas/20201211125500_sequences.go
+++ b/syncapi/storage/postgres/deltas/20201211125500_sequences.go
@@ -24,6 +24,7 @@ import (
 
 func LoadFromGoose() {
 	goose.AddMigration(UpFixSequences, DownFixSequences)
+	goose.AddMigration(UpRemoveSendToDeviceSentColumn, DownRemoveSendToDeviceSentColumn)
 }
 
 func LoadFixSequences(m *sqlutil.Migrations) {
diff --git a/syncapi/storage/postgres/deltas/20210112130000_sendtodevice_sentcolumn.go b/syncapi/storage/postgres/deltas/20210112130000_sendtodevice_sentcolumn.go
new file mode 100644
index 00000000..3690eca8
--- /dev/null
+++ b/syncapi/storage/postgres/deltas/20210112130000_sendtodevice_sentcolumn.go
@@ -0,0 +1,48 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package deltas
+
+import (
+	"database/sql"
+	"fmt"
+
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+)
+
+func LoadRemoveSendToDeviceSentColumn(m *sqlutil.Migrations) {
+	m.AddMigration(UpRemoveSendToDeviceSentColumn, DownRemoveSendToDeviceSentColumn)
+}
+
+func UpRemoveSendToDeviceSentColumn(tx *sql.Tx) error {
+	_, err := tx.Exec(`
+		ALTER TABLE syncapi_send_to_device
+		  DROP COLUMN IF EXISTS sent_by_token;
+	`)
+	if err != nil {
+		return fmt.Errorf("failed to execute upgrade: %w", err)
+	}
+	return nil
+}
+
+func DownRemoveSendToDeviceSentColumn(tx *sql.Tx) error {
+	_, err := tx.Exec(`
+		ALTER TABLE syncapi_send_to_device
+		  ADD COLUMN IF NOT EXISTS sent_by_token TEXT;
+	`)
+	if err != nil {
+		return fmt.Errorf("failed to execute downgrade: %w", err)
+	}
+	return nil
+}
diff --git a/syncapi/storage/postgres/send_to_device_table.go b/syncapi/storage/postgres/send_to_device_table.go
index ac60989c..47c1cdae 100644
--- a/syncapi/storage/postgres/send_to_device_table.go
+++ b/syncapi/storage/postgres/send_to_device_table.go
@@ -19,7 +19,6 @@ import (
 	"database/sql"
 	"encoding/json"
 
-	"github.com/lib/pq"
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/syncapi/storage/tables"
@@ -38,11 +37,7 @@ CREATE TABLE IF NOT EXISTS syncapi_send_to_device (
 	-- The device ID to send the message to.
 	device_id TEXT NOT NULL,
 	-- The event content JSON.
-	content TEXT NOT NULL,
-	-- The token that was supplied to the /sync at the time that this
-	-- message was included in a sync response, or NULL if we haven't
-	-- included it in a /sync response yet.
-	sent_by_token TEXT
+	content TEXT NOT NULL
 );
 `
 
@@ -52,34 +47,26 @@ const insertSendToDeviceMessageSQL = `
 	  RETURNING id
 `
 
-const countSendToDeviceMessagesSQL = `
-	SELECT COUNT(*)
-	  FROM syncapi_send_to_device
-	  WHERE user_id = $1 AND device_id = $2
-`
-
 const selectSendToDeviceMessagesSQL = `
-	SELECT id, user_id, device_id, content, sent_by_token
+	SELECT id, user_id, device_id, content
 	  FROM syncapi_send_to_device
-	  WHERE user_id = $1 AND device_id = $2
+	  WHERE user_id = $1 AND device_id = $2 AND id > $3 AND id <= $4
 	  ORDER BY id DESC
 `
 
-const updateSentSendToDeviceMessagesSQL = `
-	UPDATE syncapi_send_to_device SET sent_by_token = $1
-	  WHERE id = ANY($2)
-`
-
 const deleteSendToDeviceMessagesSQL = `
-	DELETE FROM syncapi_send_to_device WHERE id = ANY($1)
+	DELETE FROM syncapi_send_to_device
+	  WHERE user_id = $1 AND device_id = $2 AND id < $3
 `
 
+const selectMaxSendToDeviceIDSQL = "" +
+	"SELECT MAX(id) FROM syncapi_send_to_device"
+
 type sendToDeviceStatements struct {
-	insertSendToDeviceMessageStmt      *sql.Stmt
-	countSendToDeviceMessagesStmt      *sql.Stmt
-	selectSendToDeviceMessagesStmt     *sql.Stmt
-	updateSentSendToDeviceMessagesStmt *sql.Stmt
-	deleteSendToDeviceMessagesStmt     *sql.Stmt
+	insertSendToDeviceMessageStmt  *sql.Stmt
+	selectSendToDeviceMessagesStmt *sql.Stmt
+	deleteSendToDeviceMessagesStmt *sql.Stmt
+	selectMaxSendToDeviceIDStmt    *sql.Stmt
 }
 
 func NewPostgresSendToDeviceTable(db *sql.DB) (tables.SendToDevice, error) {
@@ -91,16 +78,13 @@ func NewPostgresSendToDeviceTable(db *sql.DB) (tables.SendToDevice, error) {
 	if s.insertSendToDeviceMessageStmt, err = db.Prepare(insertSendToDeviceMessageSQL); err != nil {
 		return nil, err
 	}
-	if s.countSendToDeviceMessagesStmt, err = db.Prepare(countSendToDeviceMessagesSQL); err != nil {
-		return nil, err
-	}
 	if s.selectSendToDeviceMessagesStmt, err = db.Prepare(selectSendToDeviceMessagesSQL); err != nil {
 		return nil, err
 	}
-	if s.updateSentSendToDeviceMessagesStmt, err = db.Prepare(updateSentSendToDeviceMessagesSQL); err != nil {
+	if s.deleteSendToDeviceMessagesStmt, err = db.Prepare(deleteSendToDeviceMessagesSQL); err != nil {
 		return nil, err
 	}
-	if s.deleteSendToDeviceMessagesStmt, err = db.Prepare(deleteSendToDeviceMessagesSQL); err != nil {
+	if s.selectMaxSendToDeviceIDStmt, err = db.Prepare(selectMaxSendToDeviceIDSQL); err != nil {
 		return nil, err
 	}
 	return s, nil
@@ -113,64 +97,55 @@ func (s *sendToDeviceStatements) InsertSendToDeviceMessage(
 	return
 }
 
-func (s *sendToDeviceStatements) CountSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, userID, deviceID string,
-) (count int, err error) {
-	row := sqlutil.TxStmt(txn, s.countSendToDeviceMessagesStmt).QueryRowContext(ctx, userID, deviceID)
-	if err = row.Scan(&count); err != nil {
-		return
-	}
-	return count, nil
-}
-
 func (s *sendToDeviceStatements) SelectSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, userID, deviceID string,
+	ctx context.Context, txn *sql.Tx, userID, deviceID string, from, to types.StreamPosition,
 ) (lastPos types.StreamPosition, events []types.SendToDeviceEvent, err error) {
-	rows, err := sqlutil.TxStmt(txn, s.selectSendToDeviceMessagesStmt).QueryContext(ctx, userID, deviceID)
+	rows, err := sqlutil.TxStmt(txn, s.selectSendToDeviceMessagesStmt).QueryContext(ctx, userID, deviceID, from, to)
 	if err != nil {
 		return
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "SelectSendToDeviceMessages: rows.close() failed")
 
 	for rows.Next() {
-		var id types.SendToDeviceNID
+		var id types.StreamPosition
 		var userID, deviceID, content string
-		var sentByToken *string
-		if err = rows.Scan(&id, &userID, &deviceID, &content, &sentByToken); err != nil {
+		if err = rows.Scan(&id, &userID, &deviceID, &content); err != nil {
 			return
 		}
+		if id > lastPos {
+			lastPos = id
+		}
 		event := types.SendToDeviceEvent{
 			ID:       id,
 			UserID:   userID,
 			DeviceID: deviceID,
 		}
 		if err = json.Unmarshal([]byte(content), &event.SendToDeviceEvent); err != nil {
-			return
-		}
-		if sentByToken != nil {
-			if token, err := types.NewStreamTokenFromString(*sentByToken); err == nil {
-				event.SentByToken = &token
-			}
+			continue
 		}
 		events = append(events, event)
-		if types.StreamPosition(id) > lastPos {
-			lastPos = types.StreamPosition(id)
-		}
 	}
-
+	if lastPos == 0 {
+		lastPos = to
+	}
 	return lastPos, events, rows.Err()
 }
 
-func (s *sendToDeviceStatements) UpdateSentSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, token string, nids []types.SendToDeviceNID,
+func (s *sendToDeviceStatements) DeleteSendToDeviceMessages(
+	ctx context.Context, txn *sql.Tx, userID, deviceID string, pos types.StreamPosition,
 ) (err error) {
-	_, err = sqlutil.TxStmt(txn, s.updateSentSendToDeviceMessagesStmt).ExecContext(ctx, token, pq.Array(nids))
+	_, err = sqlutil.TxStmt(txn, s.deleteSendToDeviceMessagesStmt).ExecContext(ctx, userID, deviceID, pos)
 	return
 }
 
-func (s *sendToDeviceStatements) DeleteSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, nids []types.SendToDeviceNID,
-) (err error) {
-	_, err = sqlutil.TxStmt(txn, s.deleteSendToDeviceMessagesStmt).ExecContext(ctx, pq.Array(nids))
+func (s *sendToDeviceStatements) SelectMaxSendToDeviceMessageID(
+	ctx context.Context, txn *sql.Tx,
+) (id int64, err error) {
+	var nullableID sql.NullInt64
+	stmt := sqlutil.TxStmt(txn, s.selectMaxSendToDeviceIDStmt)
+	err = stmt.QueryRowContext(ctx).Scan(&nullableID)
+	if nullableID.Valid {
+		id = nullableID.Int64
+	}
 	return
 }
diff --git a/syncapi/storage/postgres/syncserver.go b/syncapi/storage/postgres/syncserver.go
index 51840304..0fbf3c23 100644
--- a/syncapi/storage/postgres/syncserver.go
+++ b/syncapi/storage/postgres/syncserver.go
@@ -89,6 +89,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*SyncServerDatasource, e
 	}
 	m := sqlutil.NewMigrations()
 	deltas.LoadFixSequences(m)
+	deltas.LoadRemoveSendToDeviceSentColumn(m)
 	if err = m.RunDeltas(d.db, dbProperties); err != nil {
 		return nil, err
 	}
diff --git a/syncapi/storage/shared/syncserver.go b/syncapi/storage/shared/syncserver.go
index ebb99673..5b06aabc 100644
--- a/syncapi/storage/shared/syncserver.go
+++ b/syncapi/storage/shared/syncserver.go
@@ -29,6 +29,7 @@ import (
 	"github.com/matrix-org/dendrite/syncapi/storage/tables"
 	"github.com/matrix-org/dendrite/syncapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/sirupsen/logrus"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -85,6 +86,14 @@ func (d *Database) MaxStreamPositionForInvites(ctx context.Context) (types.Strea
 	return types.StreamPosition(id), nil
 }
 
+func (d *Database) MaxStreamPositionForSendToDeviceMessages(ctx context.Context) (types.StreamPosition, error) {
+	id, err := d.SendToDevice.SelectMaxSendToDeviceMessageID(ctx, nil)
+	if err != nil {
+		return 0, fmt.Errorf("d.SendToDevice.SelectMaxSendToDeviceMessageID: %w", err)
+	}
+	return types.StreamPosition(id), nil
+}
+
 func (d *Database) MaxStreamPositionForAccountData(ctx context.Context) (types.StreamPosition, error) {
 	id, err := d.AccountData.SelectMaxAccountDataID(ctx, nil)
 	if err != nil {
@@ -168,30 +177,6 @@ func (d *Database) GetEventsInStreamingRange(
 	return events, err
 }
 
-/*
-func (d *Database) AddTypingUser(
-	userID, roomID string, expireTime *time.Time,
-) types.StreamPosition {
-	return types.StreamPosition(d.EDUCache.AddTypingUser(userID, roomID, expireTime))
-}
-
-func (d *Database) RemoveTypingUser(
-	userID, roomID string,
-) types.StreamPosition {
-	return types.StreamPosition(d.EDUCache.RemoveUser(userID, roomID))
-}
-
-func (d *Database) SetTypingTimeoutCallback(fn cache.TimeoutCallbackFn) {
-	d.EDUCache.SetTimeoutCallback(fn)
-}
-*/
-
-/*
-func (d *Database) AddSendToDevice() types.StreamPosition {
-	return types.StreamPosition(d.EDUCache.AddSendToDeviceMessage())
-}
-*/
-
 func (d *Database) AllJoinedUsersInRooms(ctx context.Context) (map[string][]string, error) {
 	return d.CurrentRoomState.SelectJoinedUsers(ctx)
 }
@@ -891,16 +876,6 @@ func (d *Database) currentStateStreamEventsForRoom(
 	return s, nil
 }
 
-func (d *Database) SendToDeviceUpdatesWaiting(
-	ctx context.Context, userID, deviceID string,
-) (bool, error) {
-	count, err := d.SendToDevice.CountSendToDeviceMessages(ctx, nil, userID, deviceID)
-	if err != nil {
-		return false, err
-	}
-	return count > 0, nil
-}
-
 func (d *Database) StoreNewSendForDeviceMessage(
 	ctx context.Context, userID, deviceID string, event gomatrixserverlib.SendToDeviceEvent,
 ) (newPos types.StreamPosition, err error) {
@@ -919,77 +894,37 @@ func (d *Database) StoreNewSendForDeviceMessage(
 	if err != nil {
 		return 0, err
 	}
-	return 0, nil
+	return newPos, nil
 }
 
 func (d *Database) SendToDeviceUpdatesForSync(
 	ctx context.Context,
 	userID, deviceID string,
-	token types.StreamingToken,
-) (types.StreamPosition, []types.SendToDeviceEvent, []types.SendToDeviceNID, []types.SendToDeviceNID, error) {
+	from, to types.StreamPosition,
+) (types.StreamPosition, []types.SendToDeviceEvent, error) {
 	// First of all, get our send-to-device updates for this user.
-	lastPos, events, err := d.SendToDevice.SelectSendToDeviceMessages(ctx, nil, userID, deviceID)
+	lastPos, events, err := d.SendToDevice.SelectSendToDeviceMessages(ctx, nil, userID, deviceID, from, to)
 	if err != nil {
-		return 0, nil, nil, nil, fmt.Errorf("d.SendToDevice.SelectSendToDeviceMessages: %w", err)
+		return from, nil, fmt.Errorf("d.SendToDevice.SelectSendToDeviceMessages: %w", err)
 	}
-
 	// If there's nothing to do then stop here.
 	if len(events) == 0 {
-		return 0, nil, nil, nil, nil
-	}
-
-	// Work out whether we need to update any of the database entries.
-	toReturn := []types.SendToDeviceEvent{}
-	toUpdate := []types.SendToDeviceNID{}
-	toDelete := []types.SendToDeviceNID{}
-	for _, event := range events {
-		if event.SentByToken == nil {
-			// If the event has no sent-by token yet then we haven't attempted to send
-			// it. Record the current requested sync token in the database.
-			toUpdate = append(toUpdate, event.ID)
-			toReturn = append(toReturn, event)
-			event.SentByToken = &token
-		} else if token.IsAfter(*event.SentByToken) {
-			// The event had a sync token, therefore we've sent it before. The current
-			// sync token is now after the stored one so we can assume that the client
-			// successfully completed the previous sync (it would re-request it otherwise)
-			// so we can remove the entry from the database.
-			toDelete = append(toDelete, event.ID)
-		} else {
-			// It looks like the sync is being re-requested, maybe it timed out or
-			// failed. Re-send any that should have been acknowledged by now.
-			toReturn = append(toReturn, event)
-		}
+		return to, nil, nil
 	}
-
-	return lastPos, toReturn, toUpdate, toDelete, nil
+	return lastPos, events, nil
 }
 
 func (d *Database) CleanSendToDeviceUpdates(
 	ctx context.Context,
-	toUpdate, toDelete []types.SendToDeviceNID,
-	token types.StreamingToken,
+	userID, deviceID string, before types.StreamPosition,
 ) (err error) {
-	if len(toUpdate) == 0 && len(toDelete) == 0 {
-		return nil
+	if err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		return d.SendToDevice.DeleteSendToDeviceMessages(ctx, txn, userID, deviceID, before)
+	}); err != nil {
+		logrus.WithError(err).Errorf("Failed to clean up old send-to-device messages for user %q device %q", userID, deviceID)
+		return err
 	}
-	// If we need to write to the database then we'll ask the SendToDeviceWriter to
-	// do that for us. It'll guarantee that we don't lock the table for writes in
-	// more than one place.
-	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
-		// Delete any send-to-device messages marked for deletion.
-		if e := d.SendToDevice.DeleteSendToDeviceMessages(ctx, txn, toDelete); e != nil {
-			return fmt.Errorf("d.SendToDevice.DeleteSendToDeviceMessages: %w", e)
-		}
-
-		// Now update any outstanding send-to-device messages with the new sync token.
-		if e := d.SendToDevice.UpdateSentSendToDeviceMessages(ctx, txn, token.String(), toUpdate); e != nil {
-			return fmt.Errorf("d.SendToDevice.UpdateSentSendToDeviceMessages: %w", err)
-		}
-
-		return nil
-	})
-	return
+	return nil
 }
 
 // getMembershipFromEvent returns the value of content.membership iff the event is a state event
diff --git a/syncapi/storage/sqlite3/deltas/20201211125500_sequences.go b/syncapi/storage/sqlite3/deltas/20201211125500_sequences.go
index 64905013..8e7ebff8 100644
--- a/syncapi/storage/sqlite3/deltas/20201211125500_sequences.go
+++ b/syncapi/storage/sqlite3/deltas/20201211125500_sequences.go
@@ -24,6 +24,7 @@ import (
 
 func LoadFromGoose() {
 	goose.AddMigration(UpFixSequences, DownFixSequences)
+	goose.AddMigration(UpRemoveSendToDeviceSentColumn, DownRemoveSendToDeviceSentColumn)
 }
 
 func LoadFixSequences(m *sqlutil.Migrations) {
diff --git a/syncapi/storage/sqlite3/deltas/20210112130000_sendtodevice_sentcolumn.go b/syncapi/storage/sqlite3/deltas/20210112130000_sendtodevice_sentcolumn.go
new file mode 100644
index 00000000..e0c51410
--- /dev/null
+++ b/syncapi/storage/sqlite3/deltas/20210112130000_sendtodevice_sentcolumn.go
@@ -0,0 +1,67 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package deltas
+
+import (
+	"database/sql"
+	"fmt"
+
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+)
+
+func LoadRemoveSendToDeviceSentColumn(m *sqlutil.Migrations) {
+	m.AddMigration(UpRemoveSendToDeviceSentColumn, DownRemoveSendToDeviceSentColumn)
+}
+
+func UpRemoveSendToDeviceSentColumn(tx *sql.Tx) error {
+	_, err := tx.Exec(`
+		CREATE TEMPORARY TABLE syncapi_send_to_device_backup(id, user_id, device_id, content);
+		INSERT INTO syncapi_send_to_device_backup SELECT id, user_id, device_id, content FROM syncapi_send_to_device;
+		DROP TABLE syncapi_send_to_device;
+		CREATE TABLE syncapi_send_to_device(
+			id INTEGER PRIMARY KEY AUTOINCREMENT, 
+			user_id TEXT NOT NULL,
+			device_id TEXT NOT NULL,
+			content TEXT NOT NULL
+		);
+		INSERT INTO syncapi_send_to_device SELECT id, user_id, device_id, content FROM syncapi_send_to_device_backup;
+		DROP TABLE syncapi_send_to_device_backup;
+	`)
+	if err != nil {
+		return fmt.Errorf("failed to execute upgrade: %w", err)
+	}
+	return nil
+}
+
+func DownRemoveSendToDeviceSentColumn(tx *sql.Tx) error {
+	_, err := tx.Exec(`
+		CREATE TEMPORARY TABLE syncapi_send_to_device_backup(id, user_id, device_id, content);
+		INSERT INTO syncapi_send_to_device_backup SELECT id, user_id, device_id, content FROM syncapi_send_to_device;
+		DROP TABLE syncapi_send_to_device;
+		CREATE TABLE syncapi_send_to_device(
+			id INTEGER PRIMARY KEY AUTOINCREMENT, 
+			user_id TEXT NOT NULL, 
+			device_id TEXT NOT NULL, 
+			content TEXT NOT NULL,
+			sent_by_token TEXT
+		);
+		INSERT INTO syncapi_send_to_device SELECT id, user_id, device_id, content FROM syncapi_send_to_device_backup;
+		DROP TABLE syncapi_send_to_device_backup;
+	`)
+	if err != nil {
+		return fmt.Errorf("failed to execute upgrade: %w", err)
+	}
+	return nil
+}
diff --git a/syncapi/storage/sqlite3/send_to_device_table.go b/syncapi/storage/sqlite3/send_to_device_table.go
index e9dcb1bf..0b1d5bbf 100644
--- a/syncapi/storage/sqlite3/send_to_device_table.go
+++ b/syncapi/storage/sqlite3/send_to_device_table.go
@@ -18,12 +18,12 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
-	"strings"
 
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/syncapi/storage/tables"
 	"github.com/matrix-org/dendrite/syncapi/types"
+	"github.com/sirupsen/logrus"
 )
 
 const sendToDeviceSchema = `
@@ -36,11 +36,7 @@ CREATE TABLE IF NOT EXISTS syncapi_send_to_device (
 	-- The device ID to send the message to.
 	device_id TEXT NOT NULL,
 	-- The event content JSON.
-	content TEXT NOT NULL,
-	-- The token that was supplied to the /sync at the time that this
-	-- message was included in a sync response, or NULL if we haven't
-	-- included it in a /sync response yet.
-	sent_by_token TEXT
+	content TEXT NOT NULL
 );
 `
 
@@ -49,33 +45,27 @@ const insertSendToDeviceMessageSQL = `
 	  VALUES ($1, $2, $3)
 `
 
-const countSendToDeviceMessagesSQL = `
-	SELECT COUNT(*)
-	  FROM syncapi_send_to_device
-	  WHERE user_id = $1 AND device_id = $2
-`
-
 const selectSendToDeviceMessagesSQL = `
-	SELECT id, user_id, device_id, content, sent_by_token
+	SELECT id, user_id, device_id, content
 	  FROM syncapi_send_to_device
-	  WHERE user_id = $1 AND device_id = $2
+	  WHERE user_id = $1 AND device_id = $2 AND id > $3 AND id <= $4
 	  ORDER BY id DESC
 `
 
-const updateSentSendToDeviceMessagesSQL = `
-	UPDATE syncapi_send_to_device SET sent_by_token = $1
-	  WHERE id IN ($2)
-`
-
 const deleteSendToDeviceMessagesSQL = `
-	DELETE FROM syncapi_send_to_device WHERE id IN ($1)
+	DELETE FROM syncapi_send_to_device
+	  WHERE user_id = $1 AND device_id = $2 AND id < $3
 `
 
+const selectMaxSendToDeviceIDSQL = "" +
+	"SELECT MAX(id) FROM syncapi_send_to_device"
+
 type sendToDeviceStatements struct {
 	db                             *sql.DB
 	insertSendToDeviceMessageStmt  *sql.Stmt
 	selectSendToDeviceMessagesStmt *sql.Stmt
-	countSendToDeviceMessagesStmt  *sql.Stmt
+	deleteSendToDeviceMessagesStmt *sql.Stmt
+	selectMaxSendToDeviceIDStmt    *sql.Stmt
 }
 
 func NewSqliteSendToDeviceTable(db *sql.DB) (tables.SendToDevice, error) {
@@ -86,15 +76,18 @@ func NewSqliteSendToDeviceTable(db *sql.DB) (tables.SendToDevice, error) {
 	if err != nil {
 		return nil, err
 	}
-	if s.countSendToDeviceMessagesStmt, err = db.Prepare(countSendToDeviceMessagesSQL); err != nil {
-		return nil, err
-	}
 	if s.insertSendToDeviceMessageStmt, err = db.Prepare(insertSendToDeviceMessageSQL); err != nil {
 		return nil, err
 	}
 	if s.selectSendToDeviceMessagesStmt, err = db.Prepare(selectSendToDeviceMessagesSQL); err != nil {
 		return nil, err
 	}
+	if s.deleteSendToDeviceMessagesStmt, err = db.Prepare(deleteSendToDeviceMessagesSQL); err != nil {
+		return nil, err
+	}
+	if s.selectMaxSendToDeviceIDStmt, err = db.Prepare(selectMaxSendToDeviceIDSQL); err != nil {
+		return nil, err
+	}
 	return s, nil
 }
 
@@ -111,75 +104,57 @@ func (s *sendToDeviceStatements) InsertSendToDeviceMessage(
 	return
 }
 
-func (s *sendToDeviceStatements) CountSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, userID, deviceID string,
-) (count int, err error) {
-	row := sqlutil.TxStmt(txn, s.countSendToDeviceMessagesStmt).QueryRowContext(ctx, userID, deviceID)
-	if err = row.Scan(&count); err != nil {
-		return
-	}
-	return count, nil
-}
-
 func (s *sendToDeviceStatements) SelectSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, userID, deviceID string,
+	ctx context.Context, txn *sql.Tx, userID, deviceID string, from, to types.StreamPosition,
 ) (lastPos types.StreamPosition, events []types.SendToDeviceEvent, err error) {
-	rows, err := sqlutil.TxStmt(txn, s.selectSendToDeviceMessagesStmt).QueryContext(ctx, userID, deviceID)
+	rows, err := sqlutil.TxStmt(txn, s.selectSendToDeviceMessagesStmt).QueryContext(ctx, userID, deviceID, from, to)
 	if err != nil {
 		return
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "SelectSendToDeviceMessages: rows.close() failed")
 
 	for rows.Next() {
-		var id types.SendToDeviceNID
+		var id types.StreamPosition
 		var userID, deviceID, content string
-		var sentByToken *string
-		if err = rows.Scan(&id, &userID, &deviceID, &content, &sentByToken); err != nil {
+		if err = rows.Scan(&id, &userID, &deviceID, &content); err != nil {
+			logrus.WithError(err).Errorf("Failed to retrieve send-to-device message")
 			return
 		}
+		if id > lastPos {
+			lastPos = id
+		}
 		event := types.SendToDeviceEvent{
 			ID:       id,
 			UserID:   userID,
 			DeviceID: deviceID,
 		}
 		if err = json.Unmarshal([]byte(content), &event.SendToDeviceEvent); err != nil {
-			return
-		}
-		if sentByToken != nil {
-			if token, err := types.NewStreamTokenFromString(*sentByToken); err == nil {
-				event.SentByToken = &token
-			}
+			logrus.WithError(err).Errorf("Failed to unmarshal send-to-device message")
+			continue
 		}
 		events = append(events, event)
-		if types.StreamPosition(id) > lastPos {
-			lastPos = types.StreamPosition(id)
-		}
 	}
-
+	if lastPos == 0 {
+		lastPos = to
+	}
 	return lastPos, events, rows.Err()
 }
 
-func (s *sendToDeviceStatements) UpdateSentSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, token string, nids []types.SendToDeviceNID,
+func (s *sendToDeviceStatements) DeleteSendToDeviceMessages(
+	ctx context.Context, txn *sql.Tx, userID, deviceID string, pos types.StreamPosition,
 ) (err error) {
-	query := strings.Replace(updateSentSendToDeviceMessagesSQL, "($2)", sqlutil.QueryVariadic(1+len(nids)), 1)
-	params := make([]interface{}, 1+len(nids))
-	params[0] = token
-	for k, v := range nids {
-		params[k+1] = v
-	}
-	_, err = txn.ExecContext(ctx, query, params...)
+	_, err = sqlutil.TxStmt(txn, s.deleteSendToDeviceMessagesStmt).ExecContext(ctx, userID, deviceID, pos)
 	return
 }
 
-func (s *sendToDeviceStatements) DeleteSendToDeviceMessages(
-	ctx context.Context, txn *sql.Tx, nids []types.SendToDeviceNID,
-) (err error) {
-	query := strings.Replace(deleteSendToDeviceMessagesSQL, "($1)", sqlutil.QueryVariadic(len(nids)), 1)
-	params := make([]interface{}, 1+len(nids))
-	for k, v := range nids {
-		params[k] = v
+func (s *sendToDeviceStatements) SelectMaxSendToDeviceMessageID(
+	ctx context.Context, txn *sql.Tx,
+) (id int64, err error) {
+	var nullableID sql.NullInt64
+	stmt := sqlutil.TxStmt(txn, s.selectMaxSendToDeviceIDStmt)
+	err = stmt.QueryRowContext(ctx).Scan(&nullableID)
+	if nullableID.Valid {
+		id = nullableID.Int64
 	}
-	_, err = txn.ExecContext(ctx, query, params...)
 	return
 }
diff --git a/syncapi/storage/sqlite3/syncserver.go b/syncapi/storage/sqlite3/syncserver.go
index 7abe8dd0..fdb6ce4f 100644
--- a/syncapi/storage/sqlite3/syncserver.go
+++ b/syncapi/storage/sqlite3/syncserver.go
@@ -102,6 +102,7 @@ func (d *SyncServerDatasource) prepare(dbProperties *config.DatabaseOptions) (er
 	}
 	m := sqlutil.NewMigrations()
 	deltas.LoadFixSequences(m)
+	deltas.LoadRemoveSendToDeviceSentColumn(m)
 	if err = m.RunDeltas(d.db, dbProperties); err != nil {
 		return err
 	}
diff --git a/syncapi/storage/tables/interface.go b/syncapi/storage/tables/interface.go
index 1e654da5..fca88824 100644
--- a/syncapi/storage/tables/interface.go
+++ b/syncapi/storage/tables/interface.go
@@ -147,10 +147,9 @@ type BackwardsExtremities interface {
 // sync response, as the client is seemingly trying to repeat the same /sync.
 type SendToDevice interface {
 	InsertSendToDeviceMessage(ctx context.Context, txn *sql.Tx, userID, deviceID, content string) (pos types.StreamPosition, err error)
-	SelectSendToDeviceMessages(ctx context.Context, txn *sql.Tx, userID, deviceID string) (lastPos types.StreamPosition, events []types.SendToDeviceEvent, err error)
-	UpdateSentSendToDeviceMessages(ctx context.Context, txn *sql.Tx, token string, nids []types.SendToDeviceNID) (err error)
-	DeleteSendToDeviceMessages(ctx context.Context, txn *sql.Tx, nids []types.SendToDeviceNID) (err error)
-	CountSendToDeviceMessages(ctx context.Context, txn *sql.Tx, userID, deviceID string) (count int, err error)
+	SelectSendToDeviceMessages(ctx context.Context, txn *sql.Tx, userID, deviceID string, from, to types.StreamPosition) (lastPos types.StreamPosition, events []types.SendToDeviceEvent, err error)
+	DeleteSendToDeviceMessages(ctx context.Context, txn *sql.Tx, userID, deviceID string, from types.StreamPosition) (err error)
+	SelectMaxSendToDeviceMessageID(ctx context.Context, txn *sql.Tx) (id int64, err error)
 }
 
 type Filter interface {
diff --git a/syncapi/streams/stream_sendtodevice.go b/syncapi/streams/stream_sendtodevice.go
index 804f525d..a3aaf3d7 100644
--- a/syncapi/streams/stream_sendtodevice.go
+++ b/syncapi/streams/stream_sendtodevice.go
@@ -10,6 +10,16 @@ type SendToDeviceStreamProvider struct {
 	StreamProvider
 }
 
+func (p *SendToDeviceStreamProvider) Setup() {
+	p.StreamProvider.Setup()
+
+	id, err := p.DB.MaxStreamPositionForSendToDeviceMessages(context.Background())
+	if err != nil {
+		panic(err)
+	}
+	p.latest = id
+}
+
 func (p *SendToDeviceStreamProvider) CompleteSync(
 	ctx context.Context,
 	req *types.SyncRequest,
@@ -23,24 +33,19 @@ func (p *SendToDeviceStreamProvider) IncrementalSync(
 	from, to types.StreamPosition,
 ) types.StreamPosition {
 	// See if we have any new tasks to do for the send-to-device messaging.
-	lastPos, events, updates, deletions, err := p.DB.SendToDeviceUpdatesForSync(req.Context, req.Device.UserID, req.Device.ID, req.Since)
+	lastPos, events, err := p.DB.SendToDeviceUpdatesForSync(req.Context, req.Device.UserID, req.Device.ID, from, to)
 	if err != nil {
 		req.Log.WithError(err).Error("p.DB.SendToDeviceUpdatesForSync failed")
 		return from
 	}
 
-	// Before we return the sync response, make sure that we take action on
-	// any send-to-device database updates or deletions that we need to do.
-	// Then add the updates into the sync response.
-	if len(updates) > 0 || len(deletions) > 0 {
-		// Handle the updates and deletions in the database.
-		err = p.DB.CleanSendToDeviceUpdates(context.Background(), updates, deletions, req.Since)
-		if err != nil {
+	if len(events) > 0 {
+		// Clean up old send-to-device messages from before this stream position.
+		if err := p.DB.CleanSendToDeviceUpdates(req.Context, req.Device.UserID, req.Device.ID, from); err != nil {
 			req.Log.WithError(err).Error("p.DB.CleanSendToDeviceUpdates failed")
 			return from
 		}
-	}
-	if len(events) > 0 {
+
 		// Add the updates into the sync response.
 		for _, event := range events {
 			req.Response.ToDevice.Events = append(req.Response.ToDevice.Events, event.SendToDeviceEvent)
diff --git a/syncapi/types/types.go b/syncapi/types/types.go
index 90356f7d..4ccc8a48 100644
--- a/syncapi/types/types.go
+++ b/syncapi/types/types.go
@@ -492,14 +492,11 @@ func NewLeaveResponse() *LeaveResponse {
 	return &res
 }
 
-type SendToDeviceNID int
-
 type SendToDeviceEvent struct {
 	gomatrixserverlib.SendToDeviceEvent
-	ID          SendToDeviceNID
-	UserID      string
-	DeviceID    string
-	SentByToken *StreamingToken
+	ID       StreamPosition
+	UserID   string
+	DeviceID string
 }
 
 type PeekingDevice struct {