Repository: beatlabs/patron
Pull request #413 (2021-09-13 18:11:33)

New option to notify a Kafka consumer has reached the end of the topic
Signed-off-by: teivah <t.harsanyi@thebeat.co>

## Which problem is this PR solving?

https://github.com/beatlabs/patron/issues/412

## Short description of the changes

* Add a new `WithNotificationOnceReachingLatestOffset` option (please let me know if you find a better option name 😄)
* In each partition method (`partitionsFromOffset` and `partitionsSinceDuration`), if the option is enabled, we retrieve the latest offsets per partition
* In the main `Consume` method, once all the consumers have reached their latest offset, we close the input channel.

teivah: CI is now fixed, my bad 🙏 
teivah: A quick update as I didn't have the opportunity to test in k8s. Everything seems to be working fine now.

Patch:
diff --git a/component/async/kafka/kafka.go b/component/async/kafka/kafka.go
index a7a0303e6..3434b3def 100644
--- a/component/async/kafka/kafka.go
+++ b/component/async/kafka/kafka.go
@@ -88,13 +88,14 @@ func init() {
 
 // ConsumerConfig is the common configuration of patron kafka consumers.
 type ConsumerConfig struct {
-	Brokers               []string
-	Buffer                int
-	DecoderFunc           encoding.DecodeRawFunc
-	DurationBasedConsumer bool
-	DurationOffset        time.Duration
-	TimeExtractor         func(*sarama.ConsumerMessage) (time.Time, error)
-	SaramaConfig          *sarama.Config
+	Brokers                 []string
+	Buffer                  int
+	DecoderFunc             encoding.DecodeRawFunc
+	DurationBasedConsumer   bool
+	DurationOffset          time.Duration
+	TimeExtractor           func(*sarama.ConsumerMessage) (time.Time, error)
+	SaramaConfig            *sarama.Config
+	LatestOffsetReachedChan chan<- struct{}
 }
 
 type message struct {
diff --git a/component/async/kafka/simple/simple.go b/component/async/kafka/simple/simple.go
index 4611acf8b..749c22215 100644
--- a/component/async/kafka/simple/simple.go
+++ b/component/async/kafka/simple/simple.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/Shopify/sarama"
@@ -34,6 +35,18 @@ func WithDurationOffset(since time.Duration, timeExtractor TimeExtractor) kafka.
 	}
 }
 
+// WithNotificationOnceReachingLatestOffset closes the input channel once all the partition consumers have reached the
+// latest offset.
+func WithNotificationOnceReachingLatestOffset(ch chan<- struct{}) kafka.OptionFunc {
+	return func(c *kafka.ConsumerConfig) error {
+		if ch == nil {
+			return errors.New("nil channel")
+		}
+		c.LatestOffsetReachedChan = ch
+		return nil
+	}
+}
+
 // Factory definition of a consumer factory.
 type Factory struct {
 	name    string
@@ -93,16 +106,22 @@ func (f *Factory) Create() (async.Consumer, error) {
 		c.partitions = c.partitionsSinceDuration
 	}
 
+	if c.config.LatestOffsetReachedChan != nil {
+		c.latestOffsetReachedChan = c.config.LatestOffsetReachedChan
+	}
+
 	return c, nil
 }
 
 // consumer members can be injected or overwritten with the usage of OptionFunc arguments.
 type consumer struct {
-	topic      string
-	cnl        context.CancelFunc
-	ms         sarama.Consumer
-	config     kafka.ConsumerConfig
-	partitions func(context.Context) ([]sarama.PartitionConsumer, error)
+	topic                   string
+	cnl                     context.CancelFunc
+	ms                      sarama.Consumer
+	config                  kafka.ConsumerConfig
+	partitions              func(context.Context) ([]sarama.PartitionConsumer, error)
+	latestOffsetReachedChan chan<- struct{}
+	latestOffsets           []int64
 }
 
 // Close handles closing consumer.
@@ -135,8 +154,23 @@ func (c *consumer) Consume(ctx context.Context) (<-chan async.Message, <-chan er
 		return nil, nil, errors.New("got 0 partitions")
 	}
 
-	for _, pc := range pcs {
-		go func(consumer sarama.PartitionConsumer) {
+	var wg sync.WaitGroup
+	if c.latestOffsetReachedChan != nil {
+		wg.Add(len(pcs))
+		go func() {
+			// Wait for all the partition consumers to have reached the latest offset before closing the input channel.
+			wg.Wait()
+			close(c.latestOffsetReachedChan)
+		}()
+	}
+
+	for i, pc := range pcs {
+		var latestOffset int64
+		if c.latestOffsetReachedChan != nil {
+			latestOffset = c.latestOffsets[i]
+		}
+		go func(consumer sarama.PartitionConsumer, latestOffset int64) {
+			latestOffsetReached := false
 			for {
 				select {
 				case <-ctx.Done():
@@ -148,6 +182,11 @@ func (c *consumer) Consume(ctx context.Context) (<-chan async.Message, <-chan er
 					chErr <- consumerError
 					return
 				case m := <-consumer.Messages():
+					if c.latestOffsetReachedChan != nil && !latestOffsetReached && m.Offset >= latestOffset {
+						latestOffsetReached = true
+						wg.Done()
+					}
+
 					kafka.TopicPartitionOffsetDiffGaugeSet("", m.Topic, m.Partition, consumer.HighWaterMarkOffset(), m.Offset)
 					kafka.MessageStatusCountInc(kafka.MessageReceived, "", m.Topic)
 
@@ -161,14 +200,19 @@ func (c *consumer) Consume(ctx context.Context) (<-chan async.Message, <-chan er
 					chMsg <- msg
 				}
 			}
-		}(pc)
+		}(pc, latestOffset)
 	}
 
 	return chMsg, chErr, nil
 }
 
 func (c *consumer) partitionsFromOffset(_ context.Context) ([]sarama.PartitionConsumer, error) {
-	ms, err := sarama.NewConsumer(c.config.Brokers, c.config.SaramaConfig)
+	client, err := sarama.NewClient(c.config.Brokers, c.config.SaramaConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create client: %w", err)
+	}
+
+	ms, err := sarama.NewConsumerFromClient(client)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create simple consumer: %w", err)
 	}
@@ -182,7 +226,6 @@ func (c *consumer) partitionsFromOffset(_ context.Context) ([]sarama.PartitionCo
 	pcs := make([]sarama.PartitionConsumer, len(partitions))
 
 	for i, partition := range partitions {
-
 		pc, err := c.ms.ConsumePartition(c.topic, partition, c.config.SaramaConfig.Consumer.Offsets.Initial)
 		if nil != err {
 			return nil, fmt.Errorf("failed to get partition consumer: %w", err)
@@ -190,6 +233,13 @@ func (c *consumer) partitionsFromOffset(_ context.Context) ([]sarama.PartitionCo
 		pcs[i] = pc
 	}
 
+	if c.latestOffsetReachedChan != nil {
+		err := c.setLatestOffsets(client, partitions)
+		if err != nil {
+			return nil, fmt.Errorf("failed to set latest offsets: %w", err)
+		}
+	}
+
 	return pcs, nil
 }
 
@@ -240,9 +290,31 @@ func (c *consumer) partitionsSinceDuration(ctx context.Context) ([]sarama.Partit
 		pcs[i] = pc
 	}
 
+	if c.latestOffsetReachedChan != nil {
+		err := c.setLatestOffsets(client, partitions)
+		if err != nil {
+			return nil, fmt.Errorf("failed to set latest offsets: %w", err)
+		}
+	}
+
 	return pcs, nil
 }
 
+func (c *consumer) setLatestOffsets(client sarama.Client, partitions []int32) error {
+	offsets := make([]int64, len(partitions))
+	for i, partitionID := range partitions {
+		offset, err := client.GetOffset(c.topic, partitionID, sarama.OffsetNewest)
+		if err != nil {
+			return err
+		}
+		// At this stage, offset is the offset of the next message in the partition
+		offsets[i] = offset - 1
+	}
+
+	c.latestOffsets = offsets
+	return nil
+}
+
 func closePartitionConsumer(cns sarama.PartitionConsumer) {
 	if cns == nil {
 		return


Test patch:
diff --git a/component/async/kafka/simple/simple_test.go b/component/async/kafka/simple/simple_test.go
index c1f0a5125..2d77c0e34 100644
--- a/component/async/kafka/simple/simple_test.go
+++ b/component/async/kafka/simple/simple_test.go
@@ -154,3 +154,37 @@ func TestWithDurationOffset(t *testing.T) {
 		})
 	}
 }
+
+func TestWithNotificationOnceReachingLatestOffset(t *testing.T) {
+	type args struct {
+		ch chan<- struct{}
+	}
+	testCases := map[string]struct {
+		args        args
+		expectedErr error
+	}{
+		"success": {
+			args: args{
+				ch: make(chan struct{}),
+			},
+		},
+		"error - nil channel": {
+			args: args{
+				ch: nil,
+			},
+			expectedErr: errors.New("nil channel"),
+		},
+	}
+	for name, tt := range testCases {
+		t.Run(name, func(t *testing.T) {
+			c := kafka.ConsumerConfig{}
+			err := WithNotificationOnceReachingLatestOffset(tt.args.ch)(&c)
+			if tt.expectedErr != nil {
+				assert.EqualError(t, err, tt.expectedErr.Error())
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.args.ch, c.LatestOffsetReachedChan)
+			}
+		})
+	}
+}
diff --git a/test/docker/kafka/consumer_simple_integration_test.go b/test/docker/kafka/consumer_simple_integration_test.go
index 73b9a3cf8..ce41f6339 100644
--- a/test/docker/kafka/consumer_simple_integration_test.go
+++ b/test/docker/kafka/consumer_simple_integration_test.go
@@ -167,6 +167,64 @@ func TestSimpleConsume_WithDurationOffset(t *testing.T) {
 	assert.Equal(t, sent[2:], received)
 }
 
+func TestSimpleConsume_WithNotificationOnceReachingLatestOffset(t *testing.T) {
+	messages := make([]*sarama.ProducerMessage, 0)
+	numberOfMessages := 10
+	for i := 0; i < numberOfMessages; i++ {
+		messages = append(messages, getProducerMessage(simpleTopic4, "foo"))
+	}
+
+	err := sendMessages(messages...)
+	require.NoError(t, err)
+
+	chMessages := make(chan []string)
+	chErr := make(chan error)
+	chNotif := make(chan struct{})
+	go func() {
+		factory, err := simple.New("test4", simpleTopic4, Brokers(), kafka.DecoderJSON(), kafka.Version(sarama.V2_1_0_0.String()),
+			kafka.StartFromOldest(), simple.WithNotificationOnceReachingLatestOffset(chNotif))
+		if err != nil {
+			chErr <- err
+			return
+		}
+
+		consumer, err := factory.Create()
+		if err != nil {
+			chErr <- err
+			return
+		}
+		defer func() {
+			_ = consumer.Close()
+		}()
+
+		received, err := consumeMessages(consumer, numberOfMessages)
+		if err != nil {
+			chErr <- err
+			return
+		}
+
+		chMessages <- received
+	}()
+
+	time.Sleep(5 * time.Second)
+
+	select {
+	case <-chMessages:
+		break
+	case err = <-chErr:
+		require.NoError(t, err)
+	}
+
+	// At this stage, we have received all the expected messages.
+	// We should also check that the notification channel is also eventually closed.
+	select {
+	case <-time.After(time.Second):
+		assert.FailNow(t, "notification channel not closed")
+	case _, open := <-chNotif:
+		assert.False(t, open)
+	}
+}
+
 func createTimestampPayload(timestamps ...time.Time) []string {
 	payloads := make([]string, len(timestamps))
 	for i, timestamp := range timestamps {
diff --git a/test/docker/kafka/kafka_test.go b/test/docker/kafka/kafka_test.go
index 1e4f249e7..bf08a68e8 100644
--- a/test/docker/kafka/kafka_test.go
+++ b/test/docker/kafka/kafka_test.go
@@ -25,6 +25,7 @@ const (
 	simpleTopic1         = "simpleTopic1"
 	simpleTopic2         = "simpleTopic2"
 	simpleTopic3         = "simpleTopic3"
+	simpleTopic4         = "simpleTopic4"
 	groupTopic1          = "groupTopic1"
 	groupTopic2          = "groupTopic2"
 	successTopic1        = "successTopic1"
@@ -39,6 +40,8 @@ func TestMain(m *testing.M) {
 	topics := []string{
 		getTopic(simpleTopic1),
 		getTopic(simpleTopic2),
+		getTopic(simpleTopic3),
+		getTopic(simpleTopic4),
 		getTopic(groupTopic1),
 		getTopic(groupTopic2),
 		getTopic(failAllRetriesTopic1),