Remove internal.ContinualConsumer from syncapi

syncapi/consumers/keychange.go

@@ -18,29 +18,30 @@ import (
 	"context"
 	"encoding/json"
 
-	"github.com/Shopify/sarama"
 	"github.com/getsentry/sentry-go"
-	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/keyserver/api"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
+	"github.com/matrix-org/dendrite/setup/jetstream"
 	"github.com/matrix-org/dendrite/setup/process"
 	"github.com/matrix-org/dendrite/syncapi/notifier"
 	"github.com/matrix-org/dendrite/syncapi/storage"
 	"github.com/matrix-org/dendrite/syncapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/nats-io/nats.go"
 	"github.com/sirupsen/logrus"
 )
 
 // OutputKeyChangeEventConsumer consumes events that originated in the key server.
 type OutputKeyChangeEventConsumer struct {
-	ctx               context.Context
+	ctx        context.Context
-	keyChangeConsumer *internal.ContinualConsumer
+	jetstream  nats.JetStreamContext
-	db                storage.Database
+	topic      string
-	notifier          *notifier.Notifier
+	db         storage.Database
-	stream            types.StreamProvider
+	notifier   *notifier.Notifier
-	serverName        gomatrixserverlib.ServerName // our server name
+	stream     types.StreamProvider
-	rsAPI             roomserverAPI.RoomserverInternalAPI
+	serverName gomatrixserverlib.ServerName // our server name
-	keyAPI            api.KeyInternalAPI
+	rsAPI      roomserverAPI.RoomserverInternalAPI
+	keyAPI     api.KeyInternalAPI
 }
 
 // NewOutputKeyChangeEventConsumer creates a new OutputKeyChangeEventConsumer.
@@ -49,67 +50,63 @@ func NewOutputKeyChangeEventConsumer(
 	process *process.ProcessContext,
 	serverName gomatrixserverlib.ServerName,
 	topic string,
-	kafkaConsumer sarama.Consumer,
+	js nats.JetStreamContext,
 	keyAPI api.KeyInternalAPI,
 	rsAPI roomserverAPI.RoomserverInternalAPI,
 	store storage.Database,
 	notifier *notifier.Notifier,
 	stream types.StreamProvider,
 ) *OutputKeyChangeEventConsumer {
-
-	consumer := internal.ContinualConsumer{
-		Process:        process,
-		ComponentName:  "syncapi/keychange",
-		Topic:          topic,
-		Consumer:       kafkaConsumer,
-		PartitionStore: store,
-	}
-
 	s := &OutputKeyChangeEventConsumer{
-		ctx:               process.Context(),
+		ctx:        process.Context(),
-		keyChangeConsumer: &consumer,
+		jetstream:  js,
-		db:                store,
+		topic:      topic,
-		serverName:        serverName,
+		db:         store,
-		keyAPI:            keyAPI,
+		serverName: serverName,
-		rsAPI:             rsAPI,
+		keyAPI:     keyAPI,
-		notifier:          notifier,
+		rsAPI:      rsAPI,
-		stream:            stream,
+		notifier:   notifier,
+		stream:     stream,
 	}
 
-	consumer.ProcessMessage = s.onMessage
-
 	return s
 }
 
 // Start consuming from the key server
 func (s *OutputKeyChangeEventConsumer) Start() error {
-	return s.keyChangeConsumer.Start()
+	_, err := s.jetstream.Subscribe(
+		s.topic, s.onMessage,
+	)
+	return err
 }
 
-func (s *OutputKeyChangeEventConsumer) onMessage(msg *sarama.ConsumerMessage) error {
+func (s *OutputKeyChangeEventConsumer) onMessage(msg *nats.Msg) {
-	var m api.DeviceMessage
+	jetstream.WithJetStreamMessage(msg, func(msg *nats.Msg) bool {
-	if err := json.Unmarshal(msg.Value, &m); err != nil {
+		var m api.DeviceMessage
-		logrus.WithError(err).Errorf("failed to read device message from key change topic")
+		if err := json.Unmarshal(msg.Data, &m); err != nil {
-		return nil
+			logrus.WithError(err).Errorf("failed to read device message from key change topic")
-	}
+			return true
-	if m.DeviceKeys == nil && m.OutputCrossSigningKeyUpdate == nil {
+		}
-		// This probably shouldn't happen but stops us from panicking if we come
+		if m.DeviceKeys == nil && m.OutputCrossSigningKeyUpdate == nil {
-		// across an update that doesn't satisfy either types.
+			// This probably shouldn't happen but stops us from panicking if we come
-		return nil
+			// across an update that doesn't satisfy either types.
-	}
+			return true
-	switch m.Type {
+		}
-	case api.TypeCrossSigningUpdate:
+		switch m.Type {
-		return s.onCrossSigningMessage(m, m.DeviceChangeID)
+		case api.TypeCrossSigningUpdate:
-	case api.TypeDeviceKeyUpdate:
+			return s.onCrossSigningMessage(m, m.DeviceChangeID)
-		fallthrough
+		case api.TypeDeviceKeyUpdate:
-	default:
+			fallthrough
-		return s.onDeviceKeyMessage(m, m.DeviceChangeID)
+		default:
-	}
+			return s.onDeviceKeyMessage(m, m.DeviceChangeID)
+		}
+	})
+
 }
 
-func (s *OutputKeyChangeEventConsumer) onDeviceKeyMessage(m api.DeviceMessage, deviceChangeID int64) error {
+func (s *OutputKeyChangeEventConsumer) onDeviceKeyMessage(m api.DeviceMessage, deviceChangeID int64) bool {
 	if m.DeviceKeys == nil {
-		return nil
+		return true
 	}
 	output := m.DeviceKeys
 	// work out who we need to notify about the new key
@@ -120,7 +117,7 @@ func (s *OutputKeyChangeEventConsumer) onDeviceKeyMessage(m api.DeviceMessage, d
 	if err != nil {
 		logrus.WithError(err).Error("syncapi: failed to QuerySharedUsers for key change event from key server")
 		sentry.CaptureException(err)
-		return err
+		return true
 	}
 	// make sure we get our own key updates too!
 	queryRes.UserIDsToCount[output.UserID] = 1
@@ -131,10 +128,10 @@ func (s *OutputKeyChangeEventConsumer) onDeviceKeyMessage(m api.DeviceMessage, d
 		s.notifier.OnNewKeyChange(types.StreamingToken{DeviceListPosition: posUpdate}, userID, output.UserID)
 	}
 
-	return nil
+	return true
 }
 
-func (s *OutputKeyChangeEventConsumer) onCrossSigningMessage(m api.DeviceMessage, deviceChangeID int64) error {
+func (s *OutputKeyChangeEventConsumer) onCrossSigningMessage(m api.DeviceMessage, deviceChangeID int64) bool {
 	output := m.CrossSigningKeyUpdate
 	// work out who we need to notify about the new key
 	var queryRes roomserverAPI.QuerySharedUsersResponse
@@ -144,7 +141,7 @@ func (s *OutputKeyChangeEventConsumer) onCrossSigningMessage(m api.DeviceMessage
 	if err != nil {
 		logrus.WithError(err).Error("syncapi: failed to QuerySharedUsers for key change event from key server")
 		sentry.CaptureException(err)
-		return err
+		return true
 	}
 	// make sure we get our own key updates too!
 	queryRes.UserIDsToCount[output.UserID] = 1
@@ -155,5 +152,5 @@ func (s *OutputKeyChangeEventConsumer) onCrossSigningMessage(m api.DeviceMessage
 		s.notifier.OnNewKeyChange(types.StreamingToken{DeviceListPosition: posUpdate}, userID, output.UserID)
 	}
 
-	return nil
+	return true
 }

syncapi/internal/keychange.go

@@ -18,9 +18,9 @@ import (
 	"context"
 	"strings"
 
-	"github.com/Shopify/sarama"
 	keyapi "github.com/matrix-org/dendrite/keyserver/api"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
+	"github.com/matrix-org/dendrite/setup/jetstream"
 	"github.com/matrix-org/dendrite/syncapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
@@ -64,8 +64,8 @@ func DeviceListCatchup(
 	}
 
 	// now also track users who we already share rooms with but who have updated their devices between the two tokens
-	offset := sarama.OffsetOldest
+	offset := jetstream.OffsetOldest
-	toOffset := sarama.OffsetNewest
+	toOffset := jetstream.OffsetNewest
 	if to > 0 && to > from {
 		toOffset = int64(to)
 	}

syncapi/storage/interface.go

@@ -19,7 +19,6 @@ import (
 
 	eduAPI "github.com/matrix-org/dendrite/eduserver/api"
 
-	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/dendrite/syncapi/types"
 	userapi "github.com/matrix-org/dendrite/userapi/api"
@@ -27,8 +26,6 @@ import (
 )
 
 type Database interface {
-	internal.PartitionStorer
-
 	MaxStreamPositionForPDUs(ctx context.Context) (types.StreamPosition, error)
 	MaxStreamPositionForReceipts(ctx context.Context) (types.StreamPosition, error)
 	MaxStreamPositionForInvites(ctx context.Context) (types.StreamPosition, error)

syncapi/syncapi.go

@@ -48,7 +48,7 @@ func AddPublicRoutes(
 	federation *gomatrixserverlib.FederationClient,
 	cfg *config.SyncAPI,
 ) {
-	js, consumer, _ := jetstream.Prepare(&cfg.Matrix.JetStream)
+	js := jetstream.Prepare(&cfg.Matrix.JetStream)
 
 	syncDB, err := storage.NewSyncServerDatasource(&cfg.Database)
 	if err != nil {
@@ -66,7 +66,7 @@ func AddPublicRoutes(
 
 	keyChangeConsumer := consumers.NewOutputKeyChangeEventConsumer(
 		process, cfg.Matrix.ServerName, cfg.Matrix.JetStream.TopicFor(jetstream.OutputKeyChangeEvent),
-		consumer, keyAPI, rsAPI, syncDB, notifier,
+		js, keyAPI, rsAPI, syncDB, notifier,
 		streams.DeviceListStreamProvider,
 	)
 	if err = keyChangeConsumer.Start(); err != nil {