Add jetstream.WithJetStreamMessage to make ack/nak-ing less messy, use process context in consumers

appservice/consumers/roomserver.go

@@ -32,6 +32,7 @@ import (
 
 // OutputRoomEventConsumer consumes events that originated in the room server.
 type OutputRoomEventConsumer struct {
+	ctx          context.Context
 	jetstream    nats.JetStreamContext
 	topic        string
 	asDB         storage.Database
@@ -51,6 +52,7 @@ func NewOutputRoomEventConsumer(
 	workerStates []types.ApplicationServiceWorkerState,
 ) *OutputRoomEventConsumer {
 	return &OutputRoomEventConsumer{
+		ctx:          process.Context(),
 		jetstream:    js,
 		topic:        cfg.Global.JetStream.TopicFor(jetstream.OutputRoomEvent),
 		asDB:         appserviceDB,
@@ -69,18 +71,17 @@ func (s *OutputRoomEventConsumer) Start() error {
 // onMessage is called when the appservice component receives a new event from
 // the room server output log.
 func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		// Parse out the event JSON
 		var output api.OutputEvent
 		if err := json.Unmarshal(msg.Data, &output); err != nil {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("roomserver output log: message parse failure")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		if output.Type != api.OutputTypeNewRoomEvent {
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		events := []*gomatrixserverlib.HeaderedEvent{output.NewRoomEvent.Event}
@@ -89,10 +90,11 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 		// Send event to any relevant application services
 		if err := s.filterRoomserverEvents(context.TODO(), events); err != nil {
 			log.WithError(err).Errorf("roomserver output log: filter error")
-		return
+			return true
 		}
 
-	_ = msg.Ack()
+		return true
+	})
 }
 
 // filterRoomserverEvents takes in events and decides whether any of them need

federationapi/consumers/eduserver.go

@@ -32,6 +32,7 @@ import (
 
 // OutputEDUConsumer consumes events that originate in EDU server.
 type OutputEDUConsumer struct {
+	ctx               context.Context
 	jetstream         nats.JetStreamContext
 	db                storage.Database
 	queues            *queue.OutgoingQueues
@@ -50,6 +51,7 @@ func NewOutputEDUConsumer(
 	store storage.Database,
 ) *OutputEDUConsumer {
 	return &OutputEDUConsumer{
+		ctx:               process.Context(),
 		jetstream:         js,
 		queues:            queues,
 		db:                store,
@@ -78,31 +80,28 @@ func (t *OutputEDUConsumer) Start() error {
 // send-to-device events topic from the EDU server.
 func (t *OutputEDUConsumer) onSendToDeviceEvent(msg *nats.Msg) {
 	// Extract the send-to-device event from msg.
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		var ote api.OutputSendToDeviceEvent
 		if err := json.Unmarshal(msg.Data, &ote); err != nil {
 			log.WithError(err).Errorf("eduserver output log: message parse failed (expected send-to-device)")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		// only send send-to-device events which originated from us
 		_, originServerName, err := gomatrixserverlib.SplitID('@', ote.Sender)
 		if err != nil {
 			log.WithError(err).WithField("user_id", ote.Sender).Error("Failed to extract domain from send-to-device sender")
-		_ = msg.Ack()
-		return
+			return true
 		}
 		if originServerName != t.ServerName {
 			log.WithField("other_server", originServerName).Info("Suppressing send-to-device: originated elsewhere")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		_, destServerName, err := gomatrixserverlib.SplitID('@', ote.UserID)
 		if err != nil {
 			log.WithError(err).WithField("user_id", ote.UserID).Error("Failed to extract domain from send-to-device destination")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		// Pack the EDU and marshal it
@@ -122,28 +121,30 @@ func (t *OutputEDUConsumer) onSendToDeviceEvent(msg *nats.Msg) {
 		}
 		if edu.Content, err = json.Marshal(tdm); err != nil {
 			log.WithError(err).Error("failed to marshal EDU JSON")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		log.Infof("Sending send-to-device message into %q destination queue", destServerName)
 		if err := t.queues.SendEDU(edu, t.ServerName, []gomatrixserverlib.ServerName{destServerName}); err != nil {
 			log.WithError(err).Error("failed to send EDU")
+			return false
 		}
 
-	_ = msg.Ack()
+		return true
+	})
 }
 
 // onTypingEvent is called in response to a message received on the typing
 // events topic from the EDU server.
 func (t *OutputEDUConsumer) onTypingEvent(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		// Extract the typing event from msg.
 		var ote api.OutputTypingEvent
 		if err := json.Unmarshal(msg.Data, &ote); err != nil {
 			// Skip this msg but continue processing messages.
 			log.WithError(err).Errorf("eduserver output log: message parse failed (expected typing)")
 			_ = msg.Ack()
-		return
+			return true
 		}
 
 		// only send typing events which originated from us
@@ -151,16 +152,16 @@ func (t *OutputEDUConsumer) onTypingEvent(msg *nats.Msg) {
 		if err != nil {
 			log.WithError(err).WithField("user_id", ote.Event.UserID).Error("Failed to extract domain from typing sender")
 			_ = msg.Ack()
-		return
+			return true
 		}
 		if typingServerName != t.ServerName {
-		return
+			return true
 		}
 
-	joined, err := t.db.GetJoinedHosts(context.TODO(), ote.Event.RoomID)
+		joined, err := t.db.GetJoinedHosts(t.ctx, ote.Event.RoomID)
 		if err != nil {
 			log.WithError(err).WithField("room_id", ote.Event.RoomID).Error("failed to get joined hosts for room")
-		return
+			return false
 		}
 
 		names := make([]gomatrixserverlib.ServerName, len(joined))
@@ -175,45 +176,44 @@ func (t *OutputEDUConsumer) onTypingEvent(msg *nats.Msg) {
 			"typing":  ote.Event.Typing,
 		}); err != nil {
 			log.WithError(err).Error("failed to marshal EDU JSON")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		if err := t.queues.SendEDU(edu, t.ServerName, names); err != nil {
 			log.WithError(err).Error("failed to send EDU")
+			return false
 		}
 
-	_ = msg.Ack()
+		return true
+	})
 }
 
 // onReceiptEvent is called in response to a message received on the receipt
 // events topic from the EDU server.
 func (t *OutputEDUConsumer) onReceiptEvent(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		// Extract the typing event from msg.
 		var receipt api.OutputReceiptEvent
 		if err := json.Unmarshal(msg.Data, &receipt); err != nil {
 			// Skip this msg but continue processing messages.
 			log.WithError(err).Errorf("eduserver output log: message parse failed (expected receipt)")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		// only send receipt events which originated from us
 		_, receiptServerName, err := gomatrixserverlib.SplitID('@', receipt.UserID)
 		if err != nil {
 			log.WithError(err).WithField("user_id", receipt.UserID).Error("failed to extract domain from receipt sender")
-		_ = msg.Ack()
-		return
+			return true
 		}
 		if receiptServerName != t.ServerName {
-		_ = msg.Ack()
-		return // don't log, very spammy as it logs for each remote receipt
+			return true
 		}
 
-	joined, err := t.db.GetJoinedHosts(context.TODO(), receipt.RoomID)
+		joined, err := t.db.GetJoinedHosts(t.ctx, receipt.RoomID)
 		if err != nil {
 			log.WithError(err).WithField("room_id", receipt.RoomID).Error("failed to get joined hosts for room")
-		return
+			return false
 		}
 
 		names := make([]gomatrixserverlib.ServerName, len(joined))
@@ -239,13 +239,14 @@ func (t *OutputEDUConsumer) onReceiptEvent(msg *nats.Msg) {
 		}
 		if edu.Content, err = json.Marshal(content); err != nil {
 			log.WithError(err).Error("failed to marshal EDU JSON")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		if err := t.queues.SendEDU(edu, t.ServerName, names); err != nil {
 			log.WithError(err).Error("failed to send EDU")
+			return false
 		}
 
-	_ = msg.Ack()
+		return true
+	})
 }

federationapi/consumers/keychange.go

@@ -35,6 +35,7 @@ import (
 
 // KeyChangeConsumer consumes events that originate in key server.
 type KeyChangeConsumer struct {
+	ctx        context.Context
 	consumer   *internal.ContinualConsumer
 	db         storage.Database
 	queues     *queue.OutgoingQueues
@@ -52,6 +53,7 @@ func NewKeyChangeConsumer(
 	rsAPI roomserverAPI.RoomserverInternalAPI,
 ) *KeyChangeConsumer {
 	c := &KeyChangeConsumer{
+		ctx: process.Context(),
 		consumer: &internal.ContinualConsumer{
 			Process:        process,
 			ComponentName:  "federationsender/keychange",
@@ -117,7 +119,7 @@ func (t *KeyChangeConsumer) onDeviceKeyMessage(m api.DeviceMessage) error {
 	}
 
 	var queryRes roomserverAPI.QueryRoomsForUserResponse
-	err = t.rsAPI.QueryRoomsForUser(context.Background(), &roomserverAPI.QueryRoomsForUserRequest{
+	err = t.rsAPI.QueryRoomsForUser(t.ctx, &roomserverAPI.QueryRoomsForUserRequest{
 		UserID:         m.UserID,
 		WantMembership: "join",
 	}, &queryRes)
@@ -126,7 +128,7 @@ func (t *KeyChangeConsumer) onDeviceKeyMessage(m api.DeviceMessage) error {
 		return nil
 	}
 	// send this key change to all servers who share rooms with this user.
-	destinations, err := t.db.GetJoinedHostsForRooms(context.Background(), queryRes.RoomIDs)
+	destinations, err := t.db.GetJoinedHostsForRooms(t.ctx, queryRes.RoomIDs)
 	if err != nil {
 		logger.WithError(err).Error("failed to calculate joined hosts for rooms user is in")
 		return nil
@@ -169,7 +171,7 @@ func (t *KeyChangeConsumer) onCrossSigningMessage(m api.DeviceMessage) error {
 	logger := logrus.WithField("user_id", output.UserID)
 
 	var queryRes roomserverAPI.QueryRoomsForUserResponse
-	err = t.rsAPI.QueryRoomsForUser(context.Background(), &roomserverAPI.QueryRoomsForUserRequest{
+	err = t.rsAPI.QueryRoomsForUser(t.ctx, &roomserverAPI.QueryRoomsForUserRequest{
 		UserID:         output.UserID,
 		WantMembership: "join",
 	}, &queryRes)
@@ -178,7 +180,7 @@ func (t *KeyChangeConsumer) onCrossSigningMessage(m api.DeviceMessage) error {
 		return nil
 	}
 	// send this key change to all servers who share rooms with this user.
-	destinations, err := t.db.GetJoinedHostsForRooms(context.Background(), queryRes.RoomIDs)
+	destinations, err := t.db.GetJoinedHostsForRooms(t.ctx, queryRes.RoomIDs)
 	if err != nil {
 		logger.WithError(err).Error("fedsender key change consumer: failed to calculate joined hosts for rooms user is in")
 		return nil

federationapi/consumers/roomserver.go

@@ -33,6 +33,7 @@ import (
 
 // OutputRoomEventConsumer consumes events that originated in the room server.
 type OutputRoomEventConsumer struct {
+	ctx       context.Context
 	cfg       *config.FederationAPI
 	rsAPI     api.RoomserverInternalAPI
 	jetstream nats.JetStreamContext
@@ -51,6 +52,7 @@ func NewOutputRoomEventConsumer(
 	rsAPI api.RoomserverInternalAPI,
 ) *OutputRoomEventConsumer {
 	return &OutputRoomEventConsumer{
+		ctx:       process.Context(),
 		cfg:       cfg,
 		jetstream: js,
 		db:        store,
@@ -71,13 +73,13 @@ func (s *OutputRoomEventConsumer) Start() error {
 // because updates it will likely fail with a types.EventIDMismatchError when it
 // realises that it cannot update the room state using the deltas.
 func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		// Parse out the event JSON
 		var output api.OutputEvent
 		if err := json.Unmarshal(msg.Data, &output); err != nil {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("roomserver output log: message parse failure")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		switch output.Type {
@@ -85,9 +87,9 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 			ev := output.NewRoomEvent.Event
 
 			if output.NewRoomEvent.RewritesState {
-			if err := s.db.PurgeRoomState(context.TODO(), ev.RoomID()); err != nil {
+				if err := s.db.PurgeRoomState(s.ctx, ev.RoomID()); err != nil {
 					log.WithError(err).Errorf("roomserver output log: purge room state failure")
-				return
+					return false
 				}
 			}
 
@@ -97,7 +99,6 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 					log.WithField("error", output.Type).Info(
 						err.Error(),
 					)
-				_ = msg.Ack()
 				default:
 					// panic rather than continue with an inconsistent database
 					log.WithFields(log.Fields{
@@ -108,28 +109,25 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 						log.ErrorKey: err,
 					}).Panicf("roomserver output log: write room event failure")
 				}
-			return
 			}
 
-		_ = msg.Ack()
-
 		case api.OutputTypeNewInboundPeek:
 			if err := s.processInboundPeek(*output.NewInboundPeek); err != nil {
 				log.WithFields(log.Fields{
 					"event":      output.NewInboundPeek,
 					log.ErrorKey: err,
 				}).Panicf("roomserver output log: remote peek event failure")
-			return
+				return false
 			}
-		_ = msg.Ack()
 
 		default:
 			log.WithField("type", output.Type).Debug(
 				"roomserver output log: ignoring unknown output type",
 			)
-		_ = msg.Ack()
-		return
 		}
+
+		return true
+	})
 }
 
 // processInboundPeek starts tracking a new federated inbound peek (replacing the existing one if any)
@@ -146,7 +144,7 @@ func (s *OutputRoomEventConsumer) processInboundPeek(orp api.OutputNewInboundPee
 	//
 	// This is making the tests flakey.
 
-	return s.db.AddInboundPeek(context.TODO(), orp.ServerName, orp.RoomID, orp.PeekID, orp.RenewalInterval)
+	return s.db.AddInboundPeek(s.ctx, orp.ServerName, orp.RoomID, orp.PeekID, orp.RenewalInterval)
 }
 
 // processMessage updates the list of currently joined hosts in the room
@@ -162,7 +160,7 @@ func (s *OutputRoomEventConsumer) processMessage(ore api.OutputNewRoomEvent) err
 	// TODO(#290): handle EventIDMismatchError and recover the current state by
 	// talking to the roomserver
 	oldJoinedHosts, err := s.db.UpdateRoom(
-		context.TODO(),
+		s.ctx,
 		ore.Event.RoomID(),
 		ore.LastSentEventID,
 		ore.Event.EventID(),
@@ -255,7 +253,7 @@ func (s *OutputRoomEventConsumer) joinedHostsAtEvent(
 	}
 
 	// handle peeking hosts
-	inboundPeeks, err := s.db.GetInboundPeeks(context.TODO(), ore.Event.Event.RoomID())
+	inboundPeeks, err := s.db.GetInboundPeeks(s.ctx, ore.Event.Event.RoomID())
 	if err != nil {
 		return nil, err
 	}
@@ -373,7 +371,7 @@ func (s *OutputRoomEventConsumer) lookupStateEvents(
 	// from the roomserver using the query API.
 	eventReq := api.QueryEventsByIDRequest{EventIDs: missing}
 	var eventResp api.QueryEventsByIDResponse
-	if err := s.rsAPI.QueryEventsByID(context.TODO(), &eventReq, &eventResp); err != nil {
+	if err := s.rsAPI.QueryEventsByID(s.ctx, &eventReq, &eventResp); err != nil {
 		return nil, err
 	}
 

setup/jetstream/helpers.go

@@ -0,0 +1,11 @@
+package jetstream
+
+import "github.com/nats-io/nats.go"
+
+func WithJetStreamMessage(msg *nats.Msg, f func(msg *nats.Msg) bool) {
+	if f(msg) {
+		_ = msg.Ack()
+	} else {
+		_ = msg.Nak()
+	}
+}

syncapi/consumers/clientapi.go

@@ -32,6 +32,7 @@ import (
 
 // OutputClientDataConsumer consumes events that originated in the client API server.
 type OutputClientDataConsumer struct {
+	ctx       context.Context
 	jetstream nats.JetStreamContext
 	topic     string
 	db        storage.Database
@@ -49,6 +50,7 @@ func NewOutputClientDataConsumer(
 	stream types.StreamProvider,
 ) *OutputClientDataConsumer {
 	return &OutputClientDataConsumer{
+		ctx:       process.Context(),
 		jetstream: js,
 		topic:     cfg.Matrix.JetStream.TopicFor(jetstream.OutputClientData),
 		db:        store,
@@ -67,6 +69,7 @@ func (s *OutputClientDataConsumer) Start() error {
 // It is not safe for this function to be called from multiple goroutines, or else the
 // sync stream position may race and be incorrectly calculated.
 func (s *OutputClientDataConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		// Parse out the event JSON
 		userID := msg.Header.Get(jetstream.UserID)
 		var output eventutil.AccountData
@@ -74,8 +77,7 @@ func (s *OutputClientDataConsumer) onMessage(msg *nats.Msg) {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("client API server output log: message parse failure")
 			sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		log.WithFields(log.Fields{
@@ -84,7 +86,7 @@ func (s *OutputClientDataConsumer) onMessage(msg *nats.Msg) {
 		}).Info("received data from client API server")
 
 		streamPos, err := s.db.UpsertAccountData(
-		context.TODO(), userID, output.RoomID, output.Type,
+			s.ctx, userID, output.RoomID, output.Type,
 		)
 		if err != nil {
 			sentry.CaptureException(err)
@@ -98,5 +100,6 @@ func (s *OutputClientDataConsumer) onMessage(msg *nats.Msg) {
 		s.stream.Advance(streamPos)
 		s.notifier.OnNewAccountData(userID, types.StreamingToken{AccountDataPosition: streamPos})
 
-	_ = msg.Ack()
+		return true
+	})
 }

syncapi/consumers/eduserver_receipts.go

@@ -32,6 +32,7 @@ import (
 
 // OutputReceiptEventConsumer consumes events that originated in the EDU server.
 type OutputReceiptEventConsumer struct {
+	ctx       context.Context
 	jetstream nats.JetStreamContext
 	topic     string
 	db        storage.Database
@@ -50,6 +51,7 @@ func NewOutputReceiptEventConsumer(
 	stream types.StreamProvider,
 ) *OutputReceiptEventConsumer {
 	return &OutputReceiptEventConsumer{
+		ctx:       process.Context(),
 		jetstream: js,
 		topic:     cfg.Matrix.JetStream.TopicFor(jetstream.OutputReceiptEvent),
 		db:        store,
@@ -65,17 +67,17 @@ func (s *OutputReceiptEventConsumer) Start() error {
 }
 
 func (s *OutputReceiptEventConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		var output api.OutputReceiptEvent
 		if err := json.Unmarshal(msg.Data, &output); err != nil {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("EDU server output log: message parse failure")
 			sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		streamPos, err := s.db.StoreReceipt(
-		context.TODO(),
+			s.ctx,
 			output.RoomID,
 			output.Type,
 			output.UserID,
@@ -84,11 +86,12 @@ func (s *OutputReceiptEventConsumer) onMessage(msg *nats.Msg) {
 		)
 		if err != nil {
 			sentry.CaptureException(err)
-		return
+			return true
 		}
 
 		s.stream.Advance(streamPos)
 		s.notifier.OnNewReceipt(output.RoomID, types.StreamingToken{ReceiptPosition: streamPos})
 
-	_ = msg.Ack()
+		return true
+	})
 }

syncapi/consumers/eduserver_sendtodevice.go

@@ -34,6 +34,7 @@ import (
 
 // OutputSendToDeviceEventConsumer consumes events that originated in the EDU server.
 type OutputSendToDeviceEventConsumer struct {
+	ctx        context.Context
 	jetstream  nats.JetStreamContext
 	topic      string
 	db         storage.Database
@@ -53,6 +54,7 @@ func NewOutputSendToDeviceEventConsumer(
 	stream types.StreamProvider,
 ) *OutputSendToDeviceEventConsumer {
 	return &OutputSendToDeviceEventConsumer{
+		ctx:        process.Context(),
 		jetstream:  js,
 		topic:      cfg.Matrix.JetStream.TopicFor(jetstream.OutputSendToDeviceEvent),
 		db:         store,
@@ -69,24 +71,22 @@ func (s *OutputSendToDeviceEventConsumer) Start() error {
 }
 
 func (s *OutputSendToDeviceEventConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		var output api.OutputSendToDeviceEvent
 		if err := json.Unmarshal(msg.Data, &output); err != nil {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("EDU server output log: message parse failure")
 			sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		_, domain, err := gomatrixserverlib.SplitID('@', output.UserID)
 		if err != nil {
 			sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
+			return true
 		}
 		if domain != s.serverName {
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		util.GetLogger(context.TODO()).WithFields(log.Fields{
@@ -97,12 +97,12 @@ func (s *OutputSendToDeviceEventConsumer) onMessage(msg *nats.Msg) {
 		}).Info("sync API received send-to-device event from EDU server")
 
 		streamPos, err := s.db.StoreNewSendForDeviceMessage(
-		context.TODO(), output.UserID, output.DeviceID, output.SendToDeviceEvent,
+			s.ctx, output.UserID, output.DeviceID, output.SendToDeviceEvent,
 		)
 		if err != nil {
 			sentry.CaptureException(err)
 			log.WithError(err).Errorf("failed to store send-to-device message")
-		return
+			return false
 		}
 
 		s.stream.Advance(streamPos)
@@ -112,5 +112,6 @@ func (s *OutputSendToDeviceEventConsumer) onMessage(msg *nats.Msg) {
 			types.StreamingToken{SendToDevicePosition: streamPos},
 		)
 
-	_ = msg.Ack()
+		return true
+	})
 }

syncapi/consumers/eduserver_typing.go

@@ -15,6 +15,7 @@
 package consumers
 
 import (
+	"context"
 	"encoding/json"
 
 	"github.com/getsentry/sentry-go"
@@ -32,6 +33,7 @@ import (
 
 // OutputTypingEventConsumer consumes events that originated in the EDU server.
 type OutputTypingEventConsumer struct {
+	ctx       context.Context
 	jetstream nats.JetStreamContext
 	topic     string
 	eduCache  *cache.EDUCache
@@ -51,6 +53,7 @@ func NewOutputTypingEventConsumer(
 	stream types.StreamProvider,
 ) *OutputTypingEventConsumer {
 	return &OutputTypingEventConsumer{
+		ctx:       process.Context(),
 		jetstream: js,
 		topic:     cfg.Matrix.JetStream.TopicFor(jetstream.OutputTypingEvent),
 		eduCache:  eduCache,
@@ -66,13 +69,13 @@ func (s *OutputTypingEventConsumer) Start() error {
 }
 
 func (s *OutputTypingEventConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		var output api.OutputTypingEvent
 		if err := json.Unmarshal(msg.Data, &output); err != nil {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("EDU server output log: message parse failure")
 			sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		log.WithFields(log.Fields{
@@ -96,5 +99,6 @@ func (s *OutputTypingEventConsumer) onMessage(msg *nats.Msg) {
 		s.stream.Advance(typingPos)
 		s.notifier.OnNewTyping(output.Event.RoomID, types.StreamingToken{TypingPosition: typingPos})
 
-	_ = msg.Ack()
+		return true
+	})
 }

syncapi/consumers/keychange.go

@@ -34,6 +34,7 @@ import (
 
 // OutputKeyChangeEventConsumer consumes events that originated in the key server.
 type OutputKeyChangeEventConsumer struct {
+	ctx                 context.Context
 	keyChangeConsumer   *internal.ContinualConsumer
 	db                  storage.Database
 	notifier            *notifier.Notifier
@@ -68,6 +69,7 @@ func NewOutputKeyChangeEventConsumer(
 	}
 
 	s := &OutputKeyChangeEventConsumer{
+		ctx:                 process.Context(),
 		keyChangeConsumer:   &consumer,
 		db:                  store,
 		serverName:          serverName,
@@ -131,7 +133,7 @@ func (s *OutputKeyChangeEventConsumer) onDeviceKeyMessage(m api.DeviceMessage, o
 	output := m.DeviceKeys
 	// work out who we need to notify about the new key
 	var queryRes roomserverAPI.QuerySharedUsersResponse
-	err := s.rsAPI.QuerySharedUsers(context.Background(), &roomserverAPI.QuerySharedUsersRequest{
+	err := s.rsAPI.QuerySharedUsers(s.ctx, &roomserverAPI.QuerySharedUsersRequest{
 		UserID: output.UserID,
 	}, &queryRes)
 	if err != nil {
@@ -158,7 +160,7 @@ func (s *OutputKeyChangeEventConsumer) onCrossSigningMessage(m api.DeviceMessage
 	output := m.CrossSigningKeyUpdate
 	// work out who we need to notify about the new key
 	var queryRes roomserverAPI.QuerySharedUsersResponse
-	err := s.rsAPI.QuerySharedUsers(context.Background(), &roomserverAPI.QuerySharedUsersRequest{
+	err := s.rsAPI.QuerySharedUsers(s.ctx, &roomserverAPI.QuerySharedUsersRequest{
 		UserID: output.UserID,
 	}, &queryRes)
 	if err != nil {

syncapi/consumers/roomserver.go

@@ -34,6 +34,7 @@ import (
 
 // OutputRoomEventConsumer consumes events that originated in the room server.
 type OutputRoomEventConsumer struct {
+	ctx          context.Context
 	cfg          *config.SyncAPI
 	rsAPI        api.RoomserverInternalAPI
 	jetstream    nats.JetStreamContext
@@ -56,6 +57,7 @@ func NewOutputRoomEventConsumer(
 	rsAPI api.RoomserverInternalAPI,
 ) *OutputRoomEventConsumer {
 	return &OutputRoomEventConsumer{
+		ctx:          process.Context(),
 		cfg:          cfg,
 		jetstream:    js,
 		topic:        cfg.Matrix.JetStream.TopicFor(jetstream.OutputRoomEvent),
@@ -77,14 +79,14 @@ func (s *OutputRoomEventConsumer) Start() error {
 // It is not safe for this function to be called from multiple goroutines, or else the
 // sync stream position may race and be incorrectly calculated.
 func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 		// Parse out the event JSON
 		var err error
 		var output api.OutputEvent
 		if err = json.Unmarshal(msg.Data, &output); err != nil {
 			// If the message was invalid, log it and move on to the next message in the stream
 			log.WithError(err).Errorf("roomserver output log: message parse failure")
-		_ = msg.Ack()
-		return
+			return true
 		}
 
 		switch output.Type {
@@ -96,36 +98,34 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 				// in the special case where the event redacts itself, just pass the message through because
 				// we will never see the other part of the pair
 				if event.Redacts() != event.EventID() {
-				_ = msg.Ack()
-				return
+					return true
 				}
 			}
-		err = s.onNewRoomEvent(context.TODO(), *output.NewRoomEvent)
+			err = s.onNewRoomEvent(s.ctx, *output.NewRoomEvent)
 		case api.OutputTypeOldRoomEvent:
-		err = s.onOldRoomEvent(context.TODO(), *output.OldRoomEvent)
+			err = s.onOldRoomEvent(s.ctx, *output.OldRoomEvent)
 		case api.OutputTypeNewInviteEvent:
-		s.onNewInviteEvent(context.TODO(), *output.NewInviteEvent)
+			s.onNewInviteEvent(s.ctx, *output.NewInviteEvent)
 		case api.OutputTypeRetireInviteEvent:
-		s.onRetireInviteEvent(context.TODO(), *output.RetireInviteEvent)
+			s.onRetireInviteEvent(s.ctx, *output.RetireInviteEvent)
 		case api.OutputTypeNewPeek:
-		s.onNewPeek(context.TODO(), *output.NewPeek)
+			s.onNewPeek(s.ctx, *output.NewPeek)
 		case api.OutputTypeRetirePeek:
-		s.onRetirePeek(context.TODO(), *output.RetirePeek)
+			s.onRetirePeek(s.ctx, *output.RetirePeek)
 		case api.OutputTypeRedactedEvent:
-		err = s.onRedactEvent(context.TODO(), *output.RedactedEvent)
+			err = s.onRedactEvent(s.ctx, *output.RedactedEvent)
 		default:
 			log.WithField("type", output.Type).Debug(
 				"roomserver output log: ignoring unknown output type",
 			)
-		_ = msg.Ack()
 		}
 		if err != nil {
 			log.WithError(err).Error("roomserver output log: failed to process event")
-		_ = msg.Nak()
-		return
+			return false
 		}
 
-	_ = msg.Ack()
+		return true
+	})
 }
 
 func (s *OutputRoomEventConsumer) onRedactEvent(