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

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,174 +80,173 @@ 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.
-	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
-	}
+	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)")
+			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
-	}
-	if originServerName != t.ServerName {
-		log.WithField("other_server", originServerName).Info("Suppressing send-to-device: originated elsewhere")
-		_ = msg.Ack()
-		return
-	}
+		// 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")
+			return true
+		}
+		if originServerName != t.ServerName {
+			log.WithField("other_server", originServerName).Info("Suppressing send-to-device: originated elsewhere")
+			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
-	}
+		_, 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")
+			return true
+		}
 
-	// Pack the EDU and marshal it
-	edu := &gomatrixserverlib.EDU{
-		Type:   gomatrixserverlib.MDirectToDevice,
-		Origin: string(t.ServerName),
-	}
-	tdm := gomatrixserverlib.ToDeviceMessage{
-		Sender:    ote.Sender,
-		Type:      ote.Type,
-		MessageID: util.RandomString(32),
-		Messages: map[string]map[string]json.RawMessage{
-			ote.UserID: {
-				ote.DeviceID: ote.Content,
+		// Pack the EDU and marshal it
+		edu := &gomatrixserverlib.EDU{
+			Type:   gomatrixserverlib.MDirectToDevice,
+			Origin: string(t.ServerName),
+		}
+		tdm := gomatrixserverlib.ToDeviceMessage{
+			Sender:    ote.Sender,
+			Type:      ote.Type,
+			MessageID: util.RandomString(32),
+			Messages: map[string]map[string]json.RawMessage{
+				ote.UserID: {
+					ote.DeviceID: ote.Content,
+				},
 			},
-		},
-	}
-	if edu.Content, err = json.Marshal(tdm); err != nil {
-		log.WithError(err).Error("failed to marshal EDU JSON")
-		_ = msg.Ack()
-		return
-	}
+		}
+		if edu.Content, err = json.Marshal(tdm); err != nil {
+			log.WithError(err).Error("failed to marshal EDU JSON")
+			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")
-	}
+		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) {
-	// 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
-	}
+	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 true
+		}
 
-	// only send typing events which originated from us
-	_, typingServerName, err := gomatrixserverlib.SplitID('@', ote.Event.UserID)
-	if err != nil {
-		log.WithError(err).WithField("user_id", ote.Event.UserID).Error("Failed to extract domain from typing sender")
-		_ = msg.Ack()
-		return
-	}
-	if typingServerName != t.ServerName {
-		return
-	}
+		// only send typing events which originated from us
+		_, typingServerName, err := gomatrixserverlib.SplitID('@', ote.Event.UserID)
+		if err != nil {
+			log.WithError(err).WithField("user_id", ote.Event.UserID).Error("Failed to extract domain from typing sender")
+			_ = msg.Ack()
+			return true
+		}
+		if typingServerName != t.ServerName {
+			return true
+		}
 
-	joined, err := t.db.GetJoinedHosts(context.TODO(), ote.Event.RoomID)
-	if err != nil {
-		log.WithError(err).WithField("room_id", ote.Event.RoomID).Error("failed to get joined hosts for room")
-		return
-	}
+		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 false
+		}
 
-	names := make([]gomatrixserverlib.ServerName, len(joined))
-	for i := range joined {
-		names[i] = joined[i].ServerName
-	}
+		names := make([]gomatrixserverlib.ServerName, len(joined))
+		for i := range joined {
+			names[i] = joined[i].ServerName
+		}
 
-	edu := &gomatrixserverlib.EDU{Type: ote.Event.Type}
-	if edu.Content, err = json.Marshal(map[string]interface{}{
-		"room_id": ote.Event.RoomID,
-		"user_id": ote.Event.UserID,
-		"typing":  ote.Event.Typing,
-	}); err != nil {
-		log.WithError(err).Error("failed to marshal EDU JSON")
-		_ = msg.Ack()
-		return
-	}
+		edu := &gomatrixserverlib.EDU{Type: ote.Event.Type}
+		if edu.Content, err = json.Marshal(map[string]interface{}{
+			"room_id": ote.Event.RoomID,
+			"user_id": ote.Event.UserID,
+			"typing":  ote.Event.Typing,
+		}); err != nil {
+			log.WithError(err).Error("failed to marshal EDU JSON")
+			return true
+		}
 
-	if err := t.queues.SendEDU(edu, t.ServerName, names); err != nil {
-		log.WithError(err).Error("failed to send EDU")
-	}
+		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) {
-	// 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
-	}
+	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)")
+			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
-	}
-	if receiptServerName != t.ServerName {
-		_ = msg.Ack()
-		return // don't log, very spammy as it logs for each remote receipt
-	}
+		// 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")
+			return true
+		}
+		if receiptServerName != t.ServerName {
+			return true
+		}
 
-	joined, err := t.db.GetJoinedHosts(context.TODO(), receipt.RoomID)
-	if err != nil {
-		log.WithError(err).WithField("room_id", receipt.RoomID).Error("failed to get joined hosts for room")
-		return
-	}
+		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 false
+		}
 
-	names := make([]gomatrixserverlib.ServerName, len(joined))
-	for i := range joined {
-		names[i] = joined[i].ServerName
-	}
+		names := make([]gomatrixserverlib.ServerName, len(joined))
+		for i := range joined {
+			names[i] = joined[i].ServerName
+		}
 
-	content := map[string]api.FederationReceiptMRead{}
-	content[receipt.RoomID] = api.FederationReceiptMRead{
-		User: map[string]api.FederationReceiptData{
-			receipt.UserID: {
-				Data: api.ReceiptTS{
-					TS: receipt.Timestamp,
+		content := map[string]api.FederationReceiptMRead{}
+		content[receipt.RoomID] = api.FederationReceiptMRead{
+			User: map[string]api.FederationReceiptData{
+				receipt.UserID: {
+					Data: api.ReceiptTS{
+						TS: receipt.Timestamp,
+					},
+					EventIDs: []string{receipt.EventID},
 				},
-				EventIDs: []string{receipt.EventID},
 			},
-		},
-	}
+		}
 
-	edu := &gomatrixserverlib.EDU{
-		Type:   gomatrixserverlib.MReceipt,
-		Origin: string(t.ServerName),
-	}
-	if edu.Content, err = json.Marshal(content); err != nil {
-		log.WithError(err).Error("failed to marshal EDU JSON")
-		_ = msg.Ack()
-		return
-	}
+		edu := &gomatrixserverlib.EDU{
+			Type:   gomatrixserverlib.MReceipt,
+			Origin: string(t.ServerName),
+		}
+		if edu.Content, err = json.Marshal(content); err != nil {
+			log.WithError(err).Error("failed to marshal EDU JSON")
+			return true
+		}
 
-	if err := t.queues.SendEDU(edu, t.ServerName, names); err != nil {
-		log.WithError(err).Error("failed to send EDU")
-	}
+		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,65 +73,61 @@ 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) {
-	// 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
-	}
-
-	switch output.Type {
-	case api.OutputTypeNewRoomEvent:
-		ev := output.NewRoomEvent.Event
-
-		if output.NewRoomEvent.RewritesState {
-			if err := s.db.PurgeRoomState(context.TODO(), ev.RoomID()); err != nil {
-				log.WithError(err).Errorf("roomserver output log: purge room state failure")
-				return
-			}
+	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")
+			return true
 		}
 
-		if err := s.processMessage(*output.NewRoomEvent); err != nil {
-			switch err.(type) {
-			case *queue.ErrorFederationDisabled:
-				log.WithField("error", output.Type).Info(
-					err.Error(),
-				)
-				_ = msg.Ack()
-			default:
-				// panic rather than continue with an inconsistent database
+		switch output.Type {
+		case api.OutputTypeNewRoomEvent:
+			ev := output.NewRoomEvent.Event
+
+			if output.NewRoomEvent.RewritesState {
+				if err := s.db.PurgeRoomState(s.ctx, ev.RoomID()); err != nil {
+					log.WithError(err).Errorf("roomserver output log: purge room state failure")
+					return false
+				}
+			}
+
+			if err := s.processMessage(*output.NewRoomEvent); err != nil {
+				switch err.(type) {
+				case *queue.ErrorFederationDisabled:
+					log.WithField("error", output.Type).Info(
+						err.Error(),
+					)
+				default:
+					// panic rather than continue with an inconsistent database
+					log.WithFields(log.Fields{
+						"event_id":   ev.EventID(),
+						"event":      string(ev.JSON()),
+						"add":        output.NewRoomEvent.AddsStateEventIDs,
+						"del":        output.NewRoomEvent.RemovesStateEventIDs,
+						log.ErrorKey: err,
+					}).Panicf("roomserver output log: write room event failure")
+				}
+			}
+
+		case api.OutputTypeNewInboundPeek:
+			if err := s.processInboundPeek(*output.NewInboundPeek); err != nil {
 				log.WithFields(log.Fields{
-					"event_id":   ev.EventID(),
-					"event":      string(ev.JSON()),
-					"add":        output.NewRoomEvent.AddsStateEventIDs,
-					"del":        output.NewRoomEvent.RemovesStateEventIDs,
+					"event":      output.NewInboundPeek,
 					log.ErrorKey: err,
-				}).Panicf("roomserver output log: write room event failure")
+				}).Panicf("roomserver output log: remote peek event failure")
+				return false
 			}
-			return
+
+		default:
+			log.WithField("type", output.Type).Debug(
+				"roomserver output log: ignoring unknown output type",
+			)
 		}
 
-		_ = 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
-		}
-		_ = 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
 	}