feat(consumer): add batch send deadline

Author: ggml1

Diff for: README.md

@@ -29,6 +29,7 @@ To create new injectors for your topics, you should create a new kubernetes depl
 - `ELASTICSEARCH_DISABLE_SNIFFING` if set to "true", the client will not sniff Elasticsearch nodes during the node discovery process. Defaults to false. **OPTIONAL**
 - `KAFKA_CONSUMER_CONCURRENCY` Number of parallel goroutines working as a consumer. Default value is 1 **OPTIONAL**
 - `KAFKA_CONSUMER_BATCH_SIZE` Number of records to accumulate before sending them to Elasticsearch (for each goroutine). Default value is 100 **OPTIONAL**
+- `KAFKA_CONSUMER_BATCH_DEADLINE` If no new records are added to the batch after this time duration, the batch will be sent to Elasticsearch. Default value is 1m **OPTIONAL**
 - `ES_INDEX_COLUMN` Record field to append to index name. Ex: to create one ES index per campaign, use "campaign_id" here **OPTIONAL**
 - `ES_BLACKLISTED_COLUMNS` Comma separated list of record fields to filter before sending to Elasticsearch. Defaults to empty string. **OPTIONAL**
 - `ES_DOC_ID_COLUMN` Record field to be the document ID of Elasticsearch. Defaults to "kafkaRecordPartition:kafkaRecordOffset". **OPTIONAL**

Diff for: cmd/injector.go

@@ -42,6 +42,7 @@ func main() {
 		ConsumerGroup:         os.Getenv("KAFKA_CONSUMER_GROUP"),
 		Concurrency:           os.Getenv("KAFKA_CONSUMER_CONCURRENCY"),
 		BatchSize:             os.Getenv("KAFKA_CONSUMER_BATCH_SIZE"),
+		BatchDeadline:         os.Getenv("KAFKA_CONSUMER_BATCH_DEADLINE"),
 		BufferSize:            os.Getenv("KAFKA_CONSUMER_BUFFER_SIZE"),
 		MetricsUpdateInterval: os.Getenv("KAFKA_CONSUMER_METRICS_UPDATE_INTERVAL"),
 		RecordType:            os.Getenv("KAFKA_CONSUMER_RECORD_TYPE"),

Diff for: src/injector/injector.go

@@ -22,6 +22,11 @@ func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *s
 		level.Warn(logger).Log("err", err, "message", "failed to get consumer batch size")
 		batchSize = 100
 	}
+	batchDeadline, err := time.ParseDuration(kafkaConfig.BatchDeadline)
+	if err != nil {
+		level.Warn(logger).Log("err", err, "message", "failed to get consumer batch deadline")
+		batchDeadline = time.Minute
+	}
 	metricsUpdateInterval, err := time.ParseDuration(kafkaConfig.MetricsUpdateInterval)
 	if err != nil {
 		level.Warn(logger).Log("err", err, "message", "failed to get consumer metrics update interval")
@@ -54,6 +59,7 @@ func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *s
 		Logger:                logger,
 		Concurrency:           concurrency,
 		BatchSize:             batchSize,
+		BatchDeadline:         batchDeadline,
 		MetricsUpdateInterval: metricsUpdateInterval,
 		BufferSize:            bufferSize,
 		IncludeKey:            includeKey,

Diff for: src/kafka/config.go

@@ -10,6 +10,7 @@ type Config struct {
 	ConsumerGroup         string
 	Concurrency           string
 	BatchSize             string
+	BatchDeadline         string
 	MetricsUpdateInterval string
 	BufferSize            string
 	RecordType            string

Diff for: src/kafka/consumer.go

@@ -2,8 +2,8 @@ package kafka
 
 import (
 	"context"
-	"os"
 	"errors"
+	"os"
 
 	"time"
 
@@ -12,9 +12,9 @@ import (
 	"github.com/go-kit/kit/endpoint"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
+	e "github.com/inloco/kafka-elasticsearch-injector/src/errors"
 	"github.com/inloco/kafka-elasticsearch-injector/src/metrics"
 	"github.com/inloco/kafka-elasticsearch-injector/src/models"
-	e "github.com/inloco/kafka-elasticsearch-injector/src/errors"
 )
 
 type Notification int32
@@ -41,6 +41,7 @@ type Consumer struct {
 	Logger                log.Logger
 	Concurrency           int
 	BatchSize             int
+	BatchDeadline         time.Duration
 	MetricsUpdateInterval time.Duration
 	BufferSize            int
 	IncludeKey            bool
@@ -80,8 +81,9 @@ func (k *kafka) Start(signals chan os.Signal, notifications chan<- Notification)
 	defer consumer.Close()
 
 	buffSize := k.consumer.BatchSize
+	batchDeadline := k.consumer.BatchDeadline
 	for i := 0; i < concurrency; i++ {
-		go k.worker(consumer, buffSize, notifications)
+		go k.worker(consumer, buffSize, batchDeadline, notifications)
 	}
 	go func() {
 		for {
@@ -134,45 +136,65 @@ func (k *kafka) Start(signals chan os.Signal, notifications chan<- Notification)
 	}
 }
 
-func (k *kafka) worker(consumer *cluster.Consumer, buffSize int, notifications chan<- Notification) {
+func (k *kafka) decodeMessages(buf []*sarama.ConsumerMessage, bufIdx int) []*models.Record {
+	decoded := make([]*models.Record, 0)
+	for i := 0; i < bufIdx; i++ {
+		req, err := k.consumer.Decoder(nil, buf[i], k.consumer.IncludeKey)
+		if err != nil {
+			if errors.Is(err, e.ErrNilMessage) {
+				continue
+			}
+
+			level.Error(k.consumer.Logger).Log(
+				"message", "Error decoding message",
+				"err", err.Error(),
+			)
+			continue
+		}
+		decoded = append(decoded, req)
+	}
+
+	return decoded
+}
+
+func (k *kafka) flushMessages(buf []*sarama.ConsumerMessage, bufIdx int, consumer *cluster.Consumer, notifications chan<- Notification) {
+	records := k.decodeMessages(buf, bufIdx)
+	for {
+		if res, err := k.consumer.Endpoint(context.Background(), records); err != nil {
+			level.Error(k.consumer.Logger).Log("message", "error on endpoint call", "err", err.Error())
+			var _ = res // ignore res (for now)
+			continue
+		}
+		break
+	}
+
+	notifications <- Inserted
+	k.metricsPublisher.IncrementRecordsConsumed(len(records))
+	for i := 0; i < bufIdx; i++ {
+		k.offsetCh <- &topicPartitionOffset{buf[i].Topic, buf[i].Partition, buf[i].Offset}
+		consumer.MarkOffset(buf[i], "") // mark message as processed
+	}
+}
+
+func (k *kafka) worker(consumer *cluster.Consumer, buffSize int, batchDeadline time.Duration, notifications chan<- Notification) {
 	buf := make([]*sarama.ConsumerMessage, buffSize)
-	var decoded []*models.Record
+	var lastReceivedMsg time.Time
 	idx := 0
 	for {
-		kafkaMsg := <-k.consumerCh
-		buf[idx] = kafkaMsg
-		idx++
-		for idx == buffSize {
-			if decoded == nil {
-				for _, msg := range buf {
-					req, err := k.consumer.Decoder(nil, msg, k.consumer.IncludeKey)
-					if err != nil {
-						if errors.Is(err, e.ErrNilMessage) {
-							continue
-						}
-
-						level.Error(k.consumer.Logger).Log(
-							"message", "Error decoding message",
-							"err", err.Error(),
-						)
-						continue
-					}
-					decoded = append(decoded, req)
-				}
-			}
-			if res, err := k.consumer.Endpoint(context.Background(), decoded); err != nil {
-				level.Error(k.consumer.Logger).Log("message", "error on endpoint call", "err", err.Error())
-				var _ = res // ignore res (for now)
-				continue
+		select {
+		case kafkaMsg := <-k.consumerCh:
+			lastReceivedMsg = time.Now()
+			buf[idx] = kafkaMsg
+			idx++
+			if idx == buffSize {
+				k.flushMessages(buf, idx, consumer, notifications)
+				idx = 0
 			}
-			notifications <- Inserted
-			k.metricsPublisher.IncrementRecordsConsumed(buffSize)
-			for _, msg := range buf {
-				k.offsetCh <- &topicPartitionOffset{msg.Topic, msg.Partition, msg.Offset}
-				consumer.MarkOffset(msg, "") // mark message as processed
+		default:
+			if idx > 0 && time.Since(lastReceivedMsg) > batchDeadline {
+				k.flushMessages(buf, idx, consumer, notifications)
+				idx = 0
 			}
-			decoded = nil
-			idx = 0
 		}
 	}
 }