git.zd.zone/hrpc/hrpc@v0.0.12/mq/kafka/consumer.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  
     6  	"git.zd.zone/hrpc/hrpc/log"
     7  	"github.com/Shopify/sarama"
     8  )
     9  
    10  type consumerGroupHandler struct {
    11  	ctx context.Context
    12  	h   Handler
    13  }
    14  
    15  func (c consumerGroupHandler) Setup(_ sarama.ConsumerGroupSession) error {
    16  	return nil
    17  }
    18  func (c consumerGroupHandler) Cleanup(_ sarama.ConsumerGroupSession) error {
    19  	return nil
    20  }
    21  func (c consumerGroupHandler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
    22  	for msg := range claim.Messages() {
    23  		if err := c.h.Handle(c.ctx, msg.Topic, msg.Key, msg.Value, msg.Partition, msg.Offset); err != nil {
    24  			return err
    25  		}
    26  		sess.MarkMessage(msg, "")
    27  		sess.Commit()
    28  	}
    29  	return nil
    30  }
    31  
    32  func RegisterGroupConsumer(ctx context.Context, h Handler, groupID string, topics ...string) {
    33  	g, err := sarama.NewConsumerGroupFromClient(groupID, k.client)
    34  	if err != nil {
    35  		log.WithFields(ctx).Error(err)
    36  		return
    37  	}
    38  
    39  	go func(ctx context.Context, h Handler, g sarama.ConsumerGroup, topics ...string) {
    40  		defer g.Close()
    41  		for {
    42  			if err := g.Consume(ctx, topics, consumerGroupHandler{
    43  				ctx: ctx,
    44  				h:   h,
    45  			}); err != nil {
    46  				log.WithFields(ctx).Error(err)
    47  			}
    48  		}
    49  	}(ctx, h, g, topics...)
    50  }
    51  
    52  type Handler interface {
    53  	Handle(ctx context.Context, topic string, key, value []byte, partition int32, offset int64) error
    54  }
    55  
    56  func Consume(topic string, partition int, f func(m *Message)) error {
    57  	c, err := sarama.NewConsumerFromClient(k.client)
    58  	if err != nil {
    59  		return err
    60  	}
    61  	defer c.Close()
    62  	p, err := c.ConsumePartition(topic, int32(partition), k.opt.OffsetInitial)
    63  	if err != nil {
    64  		return err
    65  	}
    66  	defer p.Close()
    67  
    68  	for {
    69  		msg := <-p.Messages()
    70  		f(&Message{
    71  			Topic:     msg.Topic,
    72  			Value:     msg.Value,
    73  			Offset:    msg.Offset,
    74  			Partition: msg.Partition,
    75  			Timestamp: msg.Timestamp,
    76  			Key:       msg.Key,
    77  		})
    78  	}
    79  }