github.com/aaabigfish/gopkg@v1.1.0/mq/kafka/reader.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	kf "github.com/segmentio/kafka-go"
     8  )
     9  
    10  // 注意消费者的kafka版本必须0.10以上的版本
    11  type Reader interface {
    12  	ReadMessage(ctx context.Context) (Message, error)
    13  	FetchMessage(ctx context.Context) (Message, error)
    14  	CommitMessages(ctx context.Context, msgs ...Message) error
    15  	GetReader() *kf.Reader
    16  	AddHook(...ReaderFunc)
    17  	Do(m Message) error
    18  	Close() error
    19  }
    20  
    21  type reader struct {
    22  	kfReader *kf.Reader
    23  	hooks    []ReaderFunc
    24  }
    25  
    26  func NewReader(brokers []string, topic string, groupId string) Reader {
    27  	dialer := &kf.Dialer{
    28  		Timeout:   10 * time.Second,
    29  		DualStack: true,
    30  	}
    31  
    32  	// make a new reader that consumes from topic-A
    33  	r := kf.NewReader(kf.ReaderConfig{
    34  		Brokers:  brokers,
    35  		GroupID:  groupId,
    36  		Dialer:   dialer,
    37  		Topic:    topic,
    38  		MinBytes: 10e3, // 10KB
    39  		MaxBytes: 10e6, // 10MB
    40  		// CommitInterval: 10 * time.Second, // flushes commits to Kafka every second
    41  	})
    42  
    43  	return &reader{
    44  		kfReader: r,
    45  		hooks:    make([]ReaderFunc, 0),
    46  	}
    47  }
    48  
    49  func (r *reader) GetReader() *kf.Reader {
    50  	return r.kfReader
    51  }
    52  
    53  func (r *reader) ReadMessage(ctx context.Context) (Message, error) {
    54  	return r.kfReader.ReadMessage(ctx)
    55  }
    56  
    57  func (r *reader) FetchMessage(ctx context.Context) (Message, error) {
    58  	return r.kfReader.FetchMessage(ctx)
    59  }
    60  
    61  func (r *reader) CommitMessages(ctx context.Context, msgs ...Message) error {
    62  	return r.kfReader.CommitMessages(ctx, msgs...)
    63  }
    64  
    65  func (r *reader) AddHook(hook ...ReaderFunc) {
    66  	r.hooks = append(r.hooks, hook...)
    67  }
    68  
    69  func (r *reader) Do(m Message) error {
    70  	for _, hook := range r.hooks {
    71  		if err := hook(m.Partition, m.Offset, m.Key, m.Value); err != nil {
    72  			return err
    73  		}
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  func (r *reader) Close() error {
    80  	return r.kfReader.Close()
    81  }