github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/task/kafka/consumer.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"github.com/johnnyeven/libtools/task/constants"
     7  	"github.com/segmentio/kafka-go"
     8  	"github.com/sirupsen/logrus"
     9  	"time"
    10  )
    11  
    12  type KafkaConsumer struct {
    13  	r               *kafka.Reader
    14  	config          kafka.ReaderConfig
    15  	workerProcessor constants.TaskProcessor
    16  }
    17  
    18  func NewKafkaConsumer(info constants.ConnectionInfo) *KafkaConsumer {
    19  	c := kafka.ReaderConfig{
    20  		Brokers:  []string{fmt.Sprintf("%s:%d", info.Host, info.Port)},
    21  		MinBytes: 10e3,
    22  		MaxBytes: 10e6,
    23  		CommitInterval: time.Second,
    24  	}
    25  
    26  	return &KafkaConsumer{
    27  		config: c,
    28  	}
    29  }
    30  
    31  func (c *KafkaConsumer) RegisterChannel(channel string, processor constants.TaskProcessor) error {
    32  	c.config.Topic = channel
    33  	c.config.GroupID = channel + "_GROUP"
    34  	c.workerProcessor = processor
    35  
    36  	r := kafka.NewReader(c.config)
    37  	c.r = r
    38  
    39  	return nil
    40  }
    41  
    42  func (c *KafkaConsumer) Work() {
    43  	logrus.Debug("KafkaConsumer.Working...")
    44  	ctx := context.Background()
    45  	for {
    46  		m, err := c.r.ReadMessage(ctx)
    47  		if err != nil {
    48  			logrus.Errorf("kafka.Reader.ReadMessage err: %v", err)
    49  			break
    50  		}
    51  		c.r.CommitMessages(ctx, m)
    52  
    53  		t := &constants.Task{}
    54  		err = constants.UnmarshalData(m.Value, t)
    55  		if err != nil {
    56  			logrus.Errorf("Work UnmarshalData err: %v", err)
    57  			continue
    58  		}
    59  		_, err = c.workerProcessor(t)
    60  		if err != nil {
    61  			continue
    62  		}
    63  	}
    64  }
    65  
    66  func (c *KafkaConsumer) Stop() {
    67  	c.r.Close()
    68  	logrus.Debug("KafkaConsumer.Stopped")
    69  }