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

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,36 +69,37 @@ 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) {
-	// Parse out the event JSON
-	userID := msg.Header.Get(jetstream.UserID)
-	var output eventutil.AccountData
-	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("client API server output log: message parse failure")
-		sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
-	}
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
+		// Parse out the event JSON
+		userID := msg.Header.Get(jetstream.UserID)
+		var output eventutil.AccountData
+		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("client API server output log: message parse failure")
+			sentry.CaptureException(err)
+			return true
+		}
 
-	log.WithFields(log.Fields{
-		"type":    output.Type,
-		"room_id": output.RoomID,
-	}).Info("received data from client API server")
-
-	streamPos, err := s.db.UpsertAccountData(
-		context.TODO(), userID, output.RoomID, output.Type,
-	)
-	if err != nil {
-		sentry.CaptureException(err)
 		log.WithFields(log.Fields{
-			"type":       output.Type,
-			"room_id":    output.RoomID,
-			log.ErrorKey: err,
-		}).Panicf("could not save account data")
-	}
+			"type":    output.Type,
+			"room_id": output.RoomID,
+		}).Info("received data from client API server")
 
-	s.stream.Advance(streamPos)
-	s.notifier.OnNewAccountData(userID, types.StreamingToken{AccountDataPosition: streamPos})
+		streamPos, err := s.db.UpsertAccountData(
+			s.ctx, userID, output.RoomID, output.Type,
+		)
+		if err != nil {
+			sentry.CaptureException(err)
+			log.WithFields(log.Fields{
+				"type":       output.Type,
+				"room_id":    output.RoomID,
+				log.ErrorKey: err,
+			}).Panicf("could not save account data")
+		}
 
-	_ = msg.Ack()
+		s.stream.Advance(streamPos)
+		s.notifier.OnNewAccountData(userID, types.StreamingToken{AccountDataPosition: streamPos})
+
+		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,30 +67,31 @@ func (s *OutputReceiptEventConsumer) Start() error {
 }
 
 func (s *OutputReceiptEventConsumer) onMessage(msg *nats.Msg) {
-	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
-	}
+	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)
+			return true
+		}
 
-	streamPos, err := s.db.StoreReceipt(
-		context.TODO(),
-		output.RoomID,
-		output.Type,
-		output.UserID,
-		output.EventID,
-		output.Timestamp,
-	)
-	if err != nil {
-		sentry.CaptureException(err)
-		return
-	}
+		streamPos, err := s.db.StoreReceipt(
+			s.ctx,
+			output.RoomID,
+			output.Type,
+			output.UserID,
+			output.EventID,
+			output.Timestamp,
+		)
+		if err != nil {
+			sentry.CaptureException(err)
+			return true
+		}
 
-	s.stream.Advance(streamPos)
-	s.notifier.OnNewReceipt(output.RoomID, types.StreamingToken{ReceiptPosition: streamPos})
+		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,48 +71,47 @@ func (s *OutputSendToDeviceEventConsumer) Start() error {
 }
 
 func (s *OutputSendToDeviceEventConsumer) onMessage(msg *nats.Msg) {
-	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
-	}
+	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)
+			return true
+		}
 
-	_, domain, err := gomatrixserverlib.SplitID('@', output.UserID)
-	if err != nil {
-		sentry.CaptureException(err)
-		_ = msg.Ack()
-		return
-	}
-	if domain != s.serverName {
-		_ = msg.Ack()
-		return
-	}
+		_, domain, err := gomatrixserverlib.SplitID('@', output.UserID)
+		if err != nil {
+			sentry.CaptureException(err)
+			return true
+		}
+		if domain != s.serverName {
+			return true
+		}
 
-	util.GetLogger(context.TODO()).WithFields(log.Fields{
-		"sender":     output.Sender,
-		"user_id":    output.UserID,
-		"device_id":  output.DeviceID,
-		"event_type": output.Type,
-	}).Info("sync API received send-to-device event from EDU server")
+		util.GetLogger(context.TODO()).WithFields(log.Fields{
+			"sender":     output.Sender,
+			"user_id":    output.UserID,
+			"device_id":  output.DeviceID,
+			"event_type": output.Type,
+		}).Info("sync API received send-to-device event from EDU server")
 
-	streamPos, err := s.db.StoreNewSendForDeviceMessage(
-		context.TODO(), output.UserID, output.DeviceID, output.SendToDeviceEvent,
-	)
-	if err != nil {
-		sentry.CaptureException(err)
-		log.WithError(err).Errorf("failed to store send-to-device message")
-		return
-	}
+		streamPos, err := s.db.StoreNewSendForDeviceMessage(
+			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 false
+		}
 
-	s.stream.Advance(streamPos)
-	s.notifier.OnNewSendToDevice(
-		output.UserID,
-		[]string{output.DeviceID},
-		types.StreamingToken{SendToDevicePosition: streamPos},
-	)
+		s.stream.Advance(streamPos)
+		s.notifier.OnNewSendToDevice(
+			output.UserID,
+			[]string{output.DeviceID},
+			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,35 +69,36 @@ func (s *OutputTypingEventConsumer) Start() error {
 }
 
 func (s *OutputTypingEventConsumer) onMessage(msg *nats.Msg) {
-	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
-	}
+	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)
+			return true
+		}
 
-	log.WithFields(log.Fields{
-		"room_id": output.Event.RoomID,
-		"user_id": output.Event.UserID,
-		"typing":  output.Event.Typing,
-	}).Debug("received data from EDU server")
+		log.WithFields(log.Fields{
+			"room_id": output.Event.RoomID,
+			"user_id": output.Event.UserID,
+			"typing":  output.Event.Typing,
+		}).Debug("received data from EDU server")
 
-	var typingPos types.StreamPosition
-	typingEvent := output.Event
-	if typingEvent.Typing {
-		typingPos = types.StreamPosition(
-			s.eduCache.AddTypingUser(typingEvent.UserID, typingEvent.RoomID, output.ExpireTime),
-		)
-	} else {
-		typingPos = types.StreamPosition(
-			s.eduCache.RemoveUser(typingEvent.UserID, typingEvent.RoomID),
-		)
-	}
+		var typingPos types.StreamPosition
+		typingEvent := output.Event
+		if typingEvent.Typing {
+			typingPos = types.StreamPosition(
+				s.eduCache.AddTypingUser(typingEvent.UserID, typingEvent.RoomID, output.ExpireTime),
+			)
+		} else {
+			typingPos = types.StreamPosition(
+				s.eduCache.RemoveUser(typingEvent.UserID, typingEvent.RoomID),
+			)
+		}
 
-	s.stream.Advance(typingPos)
-	s.notifier.OnNewTyping(output.Event.RoomID, types.StreamingToken{TypingPosition: typingPos})
+		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,55 +79,53 @@ 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) {
-	// 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
-	}
-
-	switch output.Type {
-	case api.OutputTypeNewRoomEvent:
-		// Ignore redaction events. We will add them to the database when they are
-		// validated (when we receive OutputTypeRedactedEvent)
-		event := output.NewRoomEvent.Event
-		if event.Type() == gomatrixserverlib.MRoomRedaction && event.StateKey() == nil {
-			// 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
-			}
+	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")
+			return true
 		}
-		err = s.onNewRoomEvent(context.TODO(), *output.NewRoomEvent)
-	case api.OutputTypeOldRoomEvent:
-		err = s.onOldRoomEvent(context.TODO(), *output.OldRoomEvent)
-	case api.OutputTypeNewInviteEvent:
-		s.onNewInviteEvent(context.TODO(), *output.NewInviteEvent)
-	case api.OutputTypeRetireInviteEvent:
-		s.onRetireInviteEvent(context.TODO(), *output.RetireInviteEvent)
-	case api.OutputTypeNewPeek:
-		s.onNewPeek(context.TODO(), *output.NewPeek)
-	case api.OutputTypeRetirePeek:
-		s.onRetirePeek(context.TODO(), *output.RetirePeek)
-	case api.OutputTypeRedactedEvent:
-		err = s.onRedactEvent(context.TODO(), *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
-	}
 
-	_ = msg.Ack()
+		switch output.Type {
+		case api.OutputTypeNewRoomEvent:
+			// Ignore redaction events. We will add them to the database when they are
+			// validated (when we receive OutputTypeRedactedEvent)
+			event := output.NewRoomEvent.Event
+			if event.Type() == gomatrixserverlib.MRoomRedaction && event.StateKey() == nil {
+				// 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() {
+					return true
+				}
+			}
+			err = s.onNewRoomEvent(s.ctx, *output.NewRoomEvent)
+		case api.OutputTypeOldRoomEvent:
+			err = s.onOldRoomEvent(s.ctx, *output.OldRoomEvent)
+		case api.OutputTypeNewInviteEvent:
+			s.onNewInviteEvent(s.ctx, *output.NewInviteEvent)
+		case api.OutputTypeRetireInviteEvent:
+			s.onRetireInviteEvent(s.ctx, *output.RetireInviteEvent)
+		case api.OutputTypeNewPeek:
+			s.onNewPeek(s.ctx, *output.NewPeek)
+		case api.OutputTypeRetirePeek:
+			s.onRetirePeek(s.ctx, *output.RetirePeek)
+		case api.OutputTypeRedactedEvent:
+			err = s.onRedactEvent(s.ctx, *output.RedactedEvent)
+		default:
+			log.WithField("type", output.Type).Debug(
+				"roomserver output log: ignoring unknown output type",
+			)
+		}
+		if err != nil {
+			log.WithError(err).Error("roomserver output log: failed to process event")
+			return false
+		}
+
+		return true
+	})
 }
 
 func (s *OutputRoomEventConsumer) onRedactEvent(