Make roomserver Consumer use the factored out Consumer

src/github.com/matrix-org/dendrite/cmd/dendrite-roomserver/main.go

@@ -2,16 +2,18 @@ package main
 
 import (
 	"fmt"
-	"github.com/matrix-org/dendrite/roomserver/input"
-	"github.com/matrix-org/dendrite/roomserver/query"
-	"github.com/matrix-org/dendrite/roomserver/storage"
-	"github.com/prometheus/client_golang/prometheus"
-	sarama "gopkg.in/Shopify/sarama.v1"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
 	"strconv"
 	"strings"
+
+	"github.com/matrix-org/dendrite/common"
+	"github.com/matrix-org/dendrite/roomserver/input"
+	"github.com/matrix-org/dendrite/roomserver/query"
+	"github.com/matrix-org/dendrite/roomserver/storage"
+	"github.com/prometheus/client_golang/prometheus"
+	sarama "gopkg.in/Shopify/sarama.v1"
 )
 
 var (
@@ -43,10 +45,13 @@ func main() {
 	}
 
 	consumer := input.Consumer{
+		ContinualConsumer: common.ContinualConsumer{
+			Topic:          inputRoomEventTopic,
 			Consumer:       kafkaConsumer,
+			PartitionStore: db,
+		},
 		DB:                   db,
 		Producer:             kafkaProducer,
-		InputRoomEventTopic:  inputRoomEventTopic,
 		OutputRoomEventTopic: outputRoomEventTopic,
 	}
 

src/github.com/matrix-org/dendrite/roomserver/input/consumer.go

@@ -14,10 +14,7 @@ import (
 // A ConsumerDatabase has the storage APIs needed by the consumer.
 type ConsumerDatabase interface {
 	RoomEventDatabase
-	// PartitionOffsets returns the offsets the consumer has reached for each partition.
+	common.PartitionStorer
-	PartitionOffsets(topic string) ([]common.PartitionOffset, error)
-	// SetPartitionOffset records where the consumer has reached for a partition.
-	SetPartitionOffset(topic string, partition int32, offset int64) error
 }
 
 // An ErrorLogger handles the errors encountered by the consumer.
@@ -32,16 +29,10 @@ type ErrorLogger interface {
 // The events needed to construct the state at the event should already be stored on the roomserver.
 // If the event is not valid then it will be discarded and an error will be logged.
 type Consumer struct {
-	// A kafkaesque stream consumer providing the APIs for talking to the event source.
+	ContinualConsumer common.ContinualConsumer
-	// The interface is taken from a client library for Apache Kafka.
-	// But any equivalent event streaming protocol could be made to implement the same interface.
-	Consumer sarama.Consumer
 	// The database used to store the room events.
 	DB       ConsumerDatabase
 	Producer sarama.SyncProducer
-	// The kafkaesque topic to consume room events from.
-	// This is the name used in kafka to identify the stream to consume events from.
-	InputRoomEventTopic string
 	// The kafkaesque topic to output new room events to.
 	// This is the name used in kafka to identify the stream to write events to.
 	OutputRoomEventTopic string
@@ -76,48 +67,12 @@ func (c *Consumer) WriteOutputRoomEvent(output api.OutputRoomEvent) error {
 // Returns nil once all the goroutines are started.
 // Returns an error if it can't start consuming for any of the partitions.
 func (c *Consumer) Start() error {
-	offsets := map[int32]int64{}
+	c.ContinualConsumer.ProcessMessage = c.processMessage
-
+	c.ContinualConsumer.ShutdownCallback = c.shutdown
-	partitions, err := c.Consumer.Partitions(c.InputRoomEventTopic)
+	return c.ContinualConsumer.Start()
-	if err != nil {
-		return err
-	}
-	for _, partition := range partitions {
-		// Default all the offsets to the beginning of the stream.
-		offsets[partition] = sarama.OffsetOldest
 }
 
-	storedOffsets, err := c.DB.PartitionOffsets(c.InputRoomEventTopic)
+func (c *Consumer) processMessage(message *sarama.ConsumerMessage) error {
-	if err != nil {
-		return err
-	}
-	for _, offset := range storedOffsets {
-		// We've already processed events from this partition so advance the offset to where we got to.
-		offsets[offset.Partition] = offset.Offset
-	}
-
-	var partitionConsumers []sarama.PartitionConsumer
-	for partition, offset := range offsets {
-		pc, err := c.Consumer.ConsumePartition(c.InputRoomEventTopic, partition, offset)
-		if err != nil {
-			for _, p := range partitionConsumers {
-				p.Close()
-			}
-			return err
-		}
-		partitionConsumers = append(partitionConsumers, pc)
-	}
-	for _, pc := range partitionConsumers {
-		go c.consumePartition(pc)
-	}
-
-	return nil
-}
-
-// consumePartition consumes the room events for a single partition of the kafkaesque stream.
-func (c *Consumer) consumePartition(pc sarama.PartitionConsumer) {
-	defer pc.Close()
-	for message := range pc.Messages() {
 	var input api.InputRoomEvent
 	if err := json.Unmarshal(message.Value, &input); err != nil {
 		// If the message is invalid then log it and move onto the next message in the stream.
@@ -132,10 +87,6 @@ func (c *Consumer) consumePartition(pc sarama.PartitionConsumer) {
 			c.logError(message, err)
 		}
 	}
-		// Advance our position in the stream so that we will start at the right position after a restart.
-		if err := c.DB.SetPartitionOffset(c.InputRoomEventTopic, message.Partition, message.Offset); err != nil {
-			c.logError(message, err)
-		}
 	// Update the number of processed messages using atomic addition because it is accessed from multiple goroutines.
 	processed := atomic.AddInt64(&c.processed, 1)
 	// Check if we should stop processing.
@@ -144,12 +95,15 @@ func (c *Consumer) consumePartition(pc sarama.PartitionConsumer) {
 	// between M and (N + M) messages because the N goroutines could all try to process what they think will be the
 	// last message. We could be more careful here but this is good enough for getting rough benchmarks.
 	if c.StopProcessingAfter != nil && processed >= int64(*c.StopProcessingAfter) {
+		return common.ErrShutdown
+	}
+	return nil
+}
+
+func (c *Consumer) shutdown() {
 	if c.ShutdownCallback != nil {
 		c.ShutdownCallback(fmt.Sprintf("Stopping processing after %d messages", c.processed))
 	}
-			return
-		}
-	}
 }
 
 // logError is a convenience method for logging errors.