Pull consumers

appservice/consumers/roomserver.go

@@ -66,14 +66,15 @@ func NewOutputRoomEventConsumer(
 
 // Start consuming from room servers
 func (s *OutputRoomEventConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(s.topic, s.onMessage, s.durable)
+	return jetstream.JetStreamConsumer(
-	return err
+		s.ctx, s.jetstream, s.topic, s.onMessage,
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
+	)
 }
 
 // onMessage is called when the appservice component receives a new event from
 // the room server output log.
-func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
+func (s *OutputRoomEventConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	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 {
@@ -96,7 +97,6 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 	}
 
 	return true
-	})
 }
 
 // filterRoomserverEvents takes in events and decides whether any of them need

federationapi/consumers/eduserver.go

@@ -66,13 +66,22 @@ func NewOutputEDUConsumer(
 
 // Start consuming from EDU servers
 func (t *OutputEDUConsumer) Start() error {
-	if _, err := t.jetstream.Subscribe(t.typingTopic, t.onTypingEvent, t.durable); err != nil {
+	if err := jetstream.JetStreamConsumer(
+		t.ctx, t.jetstream, t.typingTopic, t.onTypingEvent,
+		t.durable, nats.DeliverAll(), nats.ManualAck(),
+	); err != nil {
 		return err
 	}
-	if _, err := t.jetstream.Subscribe(t.sendToDeviceTopic, t.onSendToDeviceEvent, t.durable); err != nil {
+	if err := jetstream.JetStreamConsumer(
+		t.ctx, t.jetstream, t.sendToDeviceTopic, t.onSendToDeviceEvent,
+		t.durable, nats.DeliverAll(), nats.ManualAck(),
+	); err != nil {
 		return err
 	}
-	if _, err := t.jetstream.Subscribe(t.receiptTopic, t.onReceiptEvent, t.durable); err != nil {
+	if err := jetstream.JetStreamConsumer(
+		t.ctx, t.jetstream, t.receiptTopic, t.onReceiptEvent,
+		t.durable, nats.DeliverAll(), nats.ManualAck(),
+	); err != nil {
 		return err
 	}
 	return nil
@@ -80,9 +89,8 @@ func (t *OutputEDUConsumer) Start() error {
 
 // onSendToDeviceEvent is called in response to a message received on the
 // send-to-device events topic from the EDU server.
-func (t *OutputEDUConsumer) onSendToDeviceEvent(msg *nats.Msg) {
+func (t *OutputEDUConsumer) onSendToDeviceEvent(ctx context.Context, msg *nats.Msg) bool {
 	// 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)")
@@ -133,13 +141,11 @@ func (t *OutputEDUConsumer) onSendToDeviceEvent(msg *nats.Msg) {
 	}
 
 	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) {
+func (t *OutputEDUConsumer) onTypingEvent(ctx context.Context, msg *nats.Msg) bool {
-	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 {
@@ -160,7 +166,7 @@ func (t *OutputEDUConsumer) onTypingEvent(msg *nats.Msg) {
 		return true
 	}
 
-		joined, err := t.db.GetJoinedHosts(t.ctx, ote.Event.RoomID)
+	joined, err := t.db.GetJoinedHosts(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 false
@@ -187,13 +193,11 @@ func (t *OutputEDUConsumer) onTypingEvent(msg *nats.Msg) {
 	}
 
 	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) {
+func (t *OutputEDUConsumer) onReceiptEvent(ctx context.Context, msg *nats.Msg) bool {
-	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 {
@@ -212,7 +216,7 @@ func (t *OutputEDUConsumer) onReceiptEvent(msg *nats.Msg) {
 		return true
 	}
 
-		joined, err := t.db.GetJoinedHosts(t.ctx, receipt.RoomID)
+	joined, err := t.db.GetJoinedHosts(ctx, receipt.RoomID)
 	if err != nil {
 		log.WithError(err).WithField("room_id", receipt.RoomID).Error("failed to get joined hosts for room")
 		return false
@@ -250,5 +254,4 @@ func (t *OutputEDUConsumer) onReceiptEvent(msg *nats.Msg) {
 	}
 
 	return true
-	})
 }

federationapi/consumers/roomserver.go

@@ -66,20 +66,17 @@ func NewOutputRoomEventConsumer(
 
 // Start consuming from room servers
 func (s *OutputRoomEventConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(
+	return jetstream.JetStreamConsumer(
-		s.topic, s.onMessage, s.durable,
+		s.ctx, s.jetstream, s.topic, s.onMessage,
-		nats.DeliverAll(),
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
-		nats.ManualAck(),
 	)
-	return err
 }
 
 // onMessage is called when the federation server receives a new event from the room server output log.
 // It is unsafe to call this with messages for the same room in multiple gorountines
 // 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) {
+func (s *OutputRoomEventConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	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 {
@@ -133,7 +130,6 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 	}
 
 	return true
-	})
 }
 
 // processInboundPeek starts tracking a new federated inbound peek (replacing the existing one if any)

setup/jetstream/helpers.go

@@ -1,12 +1,48 @@
 package jetstream
 
-import "github.com/nats-io/nats.go"
+import (
+	"context"
+	"fmt"
 
-func WithJetStreamMessage(msg *nats.Msg, f func(msg *nats.Msg) bool) {
+	"github.com/nats-io/nats.go"
-	_ = msg.InProgress()
+	"github.com/sirupsen/logrus"
-	if f(msg) {
+)
-		_ = msg.Ack()
+
+func JetStreamConsumer(
+	ctx context.Context, nats nats.JetStreamContext, subj string,
+	f func(ctx context.Context, msg *nats.Msg) bool,
+	opts ...nats.SubOpt,
+) error {
+	sub, err := nats.SubscribeSync(subj, opts...)
+	if err != nil {
+		return fmt.Errorf("nats.SubscribeSync: %w", err)
+	}
+	go func() {
+		handle := func(err error) {
+			if err == context.Canceled || err == context.DeadlineExceeded {
+				logrus.WithContext(ctx).WithField("subject", subj).Warn(err)
+				return
+			}
+			logrus.WithContext(ctx).WithField("subject", subj).Fatal(err)
+		}
+		for {
+			msg, err := sub.NextMsgWithContext(ctx)
+			if err != nil {
+				handle(fmt.Errorf("sub.NextMsgWithContext: %w", err))
+			}
+			if err = msg.InProgress(); err != nil {
+				handle(fmt.Errorf("msg.InProgress: %w", err))
+			}
+			if f(ctx, msg) {
+				if err = msg.Ack(); err != nil {
+					handle(fmt.Errorf("msg.Ack: %w", err))
+				}
 			} else {
-		_ = msg.Nak()
+				if err = msg.Nak(); err != nil {
+					handle(fmt.Errorf("msg.Nak: %w", err))
 				}
 			}
+		}
+	}()
+	return nil
+}

syncapi/consumers/clientapi.go

@@ -63,15 +63,16 @@ func NewOutputClientDataConsumer(
 
 // Start consuming from room servers
 func (s *OutputClientDataConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(s.topic, s.onMessage, s.durable)
+	return jetstream.JetStreamConsumer(
-	return err
+		s.ctx, s.jetstream, s.topic, s.onMessage,
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
+	)
 }
 
 // onMessage is called when the sync server receives a new event from the client API server output log.
 // 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) {
+func (s *OutputClientDataConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 	// Parse out the event JSON
 	userID := msg.Header.Get(jetstream.UserID)
 	var output eventutil.AccountData
@@ -103,5 +104,4 @@ func (s *OutputClientDataConsumer) onMessage(msg *nats.Msg) {
 	s.notifier.OnNewAccountData(userID, types.StreamingToken{AccountDataPosition: streamPos})
 
 	return true
-	})
 }

syncapi/consumers/eduserver_receipts.go

@@ -64,12 +64,13 @@ func NewOutputReceiptEventConsumer(
 
 // Start consuming from EDU api
 func (s *OutputReceiptEventConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(s.topic, s.onMessage, s.durable)
+	return jetstream.JetStreamConsumer(
-	return err
+		s.ctx, s.jetstream, s.topic, s.onMessage,
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
+	)
 }
 
-func (s *OutputReceiptEventConsumer) onMessage(msg *nats.Msg) {
+func (s *OutputReceiptEventConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	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
@@ -95,5 +96,4 @@ func (s *OutputReceiptEventConsumer) onMessage(msg *nats.Msg) {
 	s.notifier.OnNewReceipt(output.RoomID, types.StreamingToken{ReceiptPosition: streamPos})
 
 	return true
-	})
 }

syncapi/consumers/eduserver_sendtodevice.go

@@ -68,12 +68,13 @@ func NewOutputSendToDeviceEventConsumer(
 
 // Start consuming from EDU api
 func (s *OutputSendToDeviceEventConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(s.topic, s.onMessage, s.durable)
+	return jetstream.JetStreamConsumer(
-	return err
+		s.ctx, s.jetstream, s.topic, s.onMessage,
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
+	)
 }
 
-func (s *OutputSendToDeviceEventConsumer) onMessage(msg *nats.Msg) {
+func (s *OutputSendToDeviceEventConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	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
@@ -115,5 +116,4 @@ func (s *OutputSendToDeviceEventConsumer) onMessage(msg *nats.Msg) {
 	)
 
 	return true
-	})
 }

syncapi/consumers/eduserver_typing.go

@@ -66,12 +66,13 @@ func NewOutputTypingEventConsumer(
 
 // Start consuming from EDU api
 func (s *OutputTypingEventConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(s.topic, s.onMessage, s.durable)
+	return jetstream.JetStreamConsumer(
-	return err
+		s.ctx, s.jetstream, s.topic, s.onMessage,
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
+	)
 }
 
-func (s *OutputTypingEventConsumer) onMessage(msg *nats.Msg) {
+func (s *OutputTypingEventConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	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
@@ -102,5 +103,4 @@ func (s *OutputTypingEventConsumer) onMessage(msg *nats.Msg) {
 	s.notifier.OnNewTyping(output.Event.RoomID, types.StreamingToken{TypingPosition: typingPos})
 
 	return true
-	})
 }

syncapi/consumers/roomserver.go

@@ -73,19 +73,16 @@ func NewOutputRoomEventConsumer(
 
 // Start consuming from room servers
 func (s *OutputRoomEventConsumer) Start() error {
-	_, err := s.jetstream.Subscribe(
+	return jetstream.JetStreamConsumer(
-		s.topic, s.onMessage, s.durable,
+		s.ctx, s.jetstream, s.topic, s.onMessage,
-		nats.DeliverAll(),
+		s.durable, nats.DeliverAll(), nats.ManualAck(),
-		nats.ManualAck(),
 	)
-	return err
 }
 
 // onMessage is called when the sync server receives a new event from the room server output log.
 // 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) {
+func (s *OutputRoomEventConsumer) onMessage(ctx context.Context, msg *nats.Msg) bool {
-	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
 	// Parse out the event JSON
 	var err error
 	var output api.OutputEvent
@@ -131,7 +128,6 @@ func (s *OutputRoomEventConsumer) onMessage(msg *nats.Msg) {
 	}
 
 	return true
-	})
 }
 
 func (s *OutputRoomEventConsumer) onRedactEvent(