gitee.com/woood2/luca@v1.0.4/cmd/consumer/internal/subscriber/retry_consumer.go

package subscriber

import (
	"context"
	"gitee.com/woood2/luca/cmd/consumer/internal/alarm"
	"github.com/Shopify/sarama"
	"go.uber.org/zap"
	"runtime/debug"
	"time"
)

// RetryConsumer represents a Sarama consumer group consumer
type RetryConsumer struct {
	GroupID    string
	MsgHandler MsgHandler
	Logger     *zap.Logger
}

// Setup is run at the beginning of a new session, before ConsumeClaim
func (consumer *RetryConsumer) Setup(sarama.ConsumerGroupSession) error {
	return nil
}

// Cleanup is run at the end of a session, once all ConsumeClaim goroutines have exited
func (consumer *RetryConsumer) Cleanup(sarama.ConsumerGroupSession) error {
	return nil
}

// ConsumeClaim must start a consumer loop of ConsumerGroupClaim's Messages().
func (consumer *RetryConsumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	for message := range claim.Messages() {
		retry := 1
		delay := 2 * time.Second
		maxDelay := 5 * time.Minute
		ch := make(chan bool)
	retryLoop:
		for {
			go func() {
				rst := consumer.handle(message, retry) //长任务注意幂等性
				select {
				case ch <- rst:
					return
				case <-session.Context().Done():
					select {
					case ch <- rst:
					default: //防止goroutine泄漏
					}
					return
				}
			}()
			select {
			case <-session.Context().Done():
				select {
				case <-time.NewTimer(5 * time.Second).C:
					return nil
				case b := <-ch:
					if b {
						session.MarkMessage(message, "")
					} else {
						alarm.Yell(retry, consumer.GroupID, message)
					}
					return nil
				}
			case b := <-ch:
				if b {
					session.MarkMessage(message, "")
					select {
					case <-session.Context().Done():
						return nil
					default:
						break retryLoop
					}
				} else {
					alarm.Yell(retry, consumer.GroupID, message)
					select {
					case <-session.Context().Done():
						return nil
					case <-time.NewTimer(delay).C:
					}
					retry++
					delay = min(delay*2, maxDelay)
				}
			}
		}
	}
	return nil
}

func (consumer *RetryConsumer) handle(msg *sarama.ConsumerMessage, retry int) bool {
	defer func() {
		if r := recover(); r != nil {
			consumer.Logger.Error("panic recover",
				zap.Any("err", r),
				zap.String("topic", msg.Topic),
				zap.String("groupID", consumer.GroupID),
				zap.Int32("partition", msg.Partition),
				zap.Int64("offset", msg.Offset),
				zap.String("stack", string(debug.Stack())),
			)
		}
	}()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rst := consumer.MsgHandler(ctx, msg, retry)
	return rst
}

func min(a, b time.Duration) time.Duration {
	if a >= b {
		return b
	} else {
		return a
	}
}