feat: add consume loop function

Author: zhufuyi

pkg/kafka/README.md

@@ -31,7 +31,7 @@ func main() {
 	defer p.Close()
 
 	// Case 1: send sarama.ProducerMessage type message
-	msg := testData[0].(*sarama.ProducerMessage) // testData is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/producer_test.go#L18
+	msg := testData[0].(*sarama.ProducerMessage) // testData is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/producer_test.go#L19
 	partition, offset, err := p.SendMessage(msg)
 	if err != nil {
 		fmt.Println(err)
@@ -73,7 +73,7 @@ func main() {
 		kafka.AsyncProducerWithVersion(sarama.V3_6_0_0),
 		kafka.AsyncProducerWithRequiredAcks(sarama.WaitForLocal),
 		kafka.AsyncProducerWithFlushMessages(50),
-		kafka.AsyncProducerWithFlushFrequency(time.milliseconds*500),
+		kafka.AsyncProducerWithFlushFrequency(time.Millisecond*500),
 	)
 	if err != nil {
 		fmt.Println(err)
@@ -82,7 +82,7 @@ func main() {
 	defer p.Close()
 
 	// Case 1: send sarama.ProducerMessage type message, supports multiple messages
-	msg := testData[0].(*sarama.ProducerMessage) // testData is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/producer_test.go#L18
+	msg := testData[0].(*sarama.ProducerMessage) // testData is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/producer_test.go#L19
 	err = p.SendMessage(msg, msg)
 	if err != nil {
 		fmt.Println(err)
@@ -112,6 +112,7 @@ package main
 import (
 	"fmt"
 	"time"
+	"context"
 	"github.com/IBM/sarama"
 	"github.com/go-dev-frame/sponge/pkg/kafka"
 )
@@ -130,11 +131,11 @@ func main() {
 	defer cg.Close()
 
 	// Case 1: consume default handle message
-	go cg.Consume(context.Background(), []string{testTopic}, handleMsgFn) // handleMsgFn is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/consumer_test.go#L19
+	go cg.ConsumeLoop(context.Background(), []string{testTopic}, handleMsgFn) // handleMsgFn is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/consumer_test.go#L19
 
 	// Case 2: consume custom handle message
-	go cg.ConsumeCustom(context.Background(), []string{testTopic}, &myConsumerGroupHandler{ // myConsumerGroupHandler is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/consumer_test.go#L26
-		autoCommitEnable: cg.autoCommitEnable,
+	go cg.ConsumeCustomLoop(context.Background(), []string{testTopic}, &myConsumerGroupHandler{ // myConsumerGroupHandler is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/consumer_test.go#L26
+		autoCommitEnable: true,
 	})
 
 	<-time.After(time.Minute) // wait exit
@@ -150,9 +151,10 @@ package main
 
 import (
 	"fmt"
+	"time"
+	"context"
 	"github.com/IBM/sarama"
 	"github.com/go-dev-frame/sponge/pkg/kafka"
-	"time"
 )
 
 func main() {
@@ -167,7 +169,7 @@ func main() {
 	defer c.Close()
 
 	// Case 1: consume one partition
-	go c.ConsumePartition(context.Background(), testTopic, 0, sarama.OffsetNewest, handleMsgFn) // // handleMsgFn is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/consumer_test.go#L19
+	go c.ConsumePartition(context.Background(), testTopic, 0, sarama.OffsetNewest, handleMsgFn) // handleMsgFn is https://github.com/go-dev-frame/sponge/blob/main/pkg/kafka/consumer_test.go#L19
 
 	// Case 2: consume all partition
 	c.ConsumeAllPartition(context.Background(), testTopic, sarama.OffsetNewest, handleMsgFn)

pkg/kafka/consumer.go

@@ -2,6 +2,7 @@ package kafka
 
 import (
 	"context"
+	"time"
 
 	"github.com/IBM/sarama"
 	"go.uber.org/zap"
@@ -51,6 +52,40 @@ func InitConsumerGroup(addrs []string, groupID string, opts ...ConsumerOption) (
 	}, nil
 }
 
+// ConsumeLoop consume messages in a loop, with rebalanced handling
+func (c *ConsumerGroup) ConsumeLoop(ctx context.Context, topics []string, handleMessageFn HandleMessageFn) {
+	backoff := time.Second
+	maxBackoff := 30 * time.Second
+
+	for {
+		err := c.Consume(ctx, topics, handleMessageFn)
+		if err != nil {
+			select {
+			case <-time.After(backoff):
+			case <-ctx.Done():
+				return
+			}
+
+			if backoff < maxBackoff {
+				backoff *= 2
+				if backoff > maxBackoff {
+					backoff = maxBackoff
+				}
+			}
+			continue
+		}
+
+		// Normal exit (mostly due to rebalanced), should reset the backoff and wait for the next time to join the consumption group
+		backoff = time.Second
+
+		if ctx.Err() != nil {
+			return
+		}
+
+		c.zapLogger.Info("rebalanced, starting new session", zap.String("group_id", c.groupID), zap.Strings("topics", topics))
+	}
+}
+
 // Consume consume messages
 func (c *ConsumerGroup) Consume(ctx context.Context, topics []string, handleMessageFn HandleMessageFn) error {
 	handler := &defaultConsumerHandler{
@@ -68,11 +103,45 @@ func (c *ConsumerGroup) Consume(ctx context.Context, topics []string, handleMess
 	return nil
 }
 
+// ConsumeCustomLoop consume messages for custom handler in a loop, with rebalanced handling
+func (c *ConsumerGroup) ConsumeCustomLoop(ctx context.Context, topics []string, handler sarama.ConsumerGroupHandler) {
+	backoff := time.Second
+	maxBackoff := 30 * time.Second
+
+	for {
+		err := c.ConsumeCustom(ctx, topics, handler)
+		if err != nil {
+			select {
+			case <-time.After(backoff):
+			case <-ctx.Done():
+				return
+			}
+
+			if backoff < maxBackoff {
+				backoff *= 2
+				if backoff > maxBackoff {
+					backoff = maxBackoff
+				}
+			}
+			continue
+		}
+
+		// Normal exit (mostly due to rebalanced), should reset the backoff and wait for the next time to join the consumption group
+		backoff = time.Second
+
+		if ctx.Err() != nil {
+			return
+		}
+
+		c.zapLogger.Info("rebalanced, starting new session", zap.String("group_id", c.groupID), zap.Strings("topics", topics))
+	}
+}
+
 // ConsumeCustom consume messages for custom handler, you need to implement the sarama.ConsumerGroupHandler interface
 func (c *ConsumerGroup) ConsumeCustom(ctx context.Context, topics []string, handler sarama.ConsumerGroupHandler) error {
 	err := c.Group.Consume(ctx, topics, handler)
 	if err != nil {
-		c.zapLogger.Error("failed to consume messages", zap.String("group_id", c.groupID), zap.Strings("topics", topics), zap.Error(err))
+		c.zapLogger.Error("failed to consume custom messages", zap.String("group_id", c.groupID), zap.Strings("topics", topics), zap.Error(err))
 		return err
 	}
 	return nil
@@ -192,11 +261,11 @@ func (c *Consumer) ConsumePartition(ctx context.Context, topic string, partition
 	for {
 		select {
 		case msg := <-pc.Messages():
-			err := handleFn(msg)
+			err = handleFn(msg)
 			if err != nil {
 				c.zapLogger.Warn("failed to handle message", zap.Error(err), zap.String("topic", topic), zap.Int32("partition", partition), zap.Int64("offset", msg.Offset))
 			}
-		case err := <-pc.Errors():
+		case err = <-pc.Errors():
 			c.zapLogger.Error("partition consumer error", zap.Error(err))
 		case <-ctx.Done():
 			return