github.com/glide-im/glide@v1.6.0/pkg/store/kafka_store_consumer.go (about)

     1  package store
     2  
     3  import (
     4  	"github.com/Shopify/sarama"
     5  	"github.com/glide-im/glide/pkg/logger"
     6  	"github.com/glide-im/glide/pkg/messages"
     7  )
     8  
     9  type KafkaConsumer struct {
    10  	consumer  sarama.Consumer
    11  	cf        func(m *messages.ChatMessage)
    12  	offlineCf func(m *messages.ChatMessage)
    13  	channelCf func(m *messages.ChatMessage)
    14  }
    15  
    16  func NewKafkaConsumer(address []string) (*KafkaConsumer, error) {
    17  
    18  	consumer, err := sarama.NewConsumer(address, sarama.NewConfig())
    19  	if err != nil {
    20  		return nil, err
    21  	}
    22  	c := &KafkaConsumer{
    23  		consumer: consumer,
    24  	}
    25  	if err = c.run(); err != nil {
    26  		return nil, err
    27  	}
    28  
    29  	return c, nil
    30  }
    31  
    32  func (c *KafkaConsumer) run() error {
    33  
    34  	partitions, err2 := c.consumer.Partitions(KafkaChatMessageTopic)
    35  
    36  	if err2 != nil {
    37  		return err2
    38  	}
    39  
    40  	for _, partition := range partitions {
    41  
    42  		consumer, err := c.consumer.ConsumePartition(KafkaChatMessageTopic, partition, sarama.OffsetNewest)
    43  		if err != nil {
    44  			return err
    45  		}
    46  
    47  		go func(pc sarama.PartitionConsumer) {
    48  			for m := range pc.Messages() {
    49  				var cm = messages.ChatMessage{}
    50  				err2 := messages.JsonCodec.Decode(m.Value, &cm)
    51  				if err2 != nil {
    52  					logger.E("message decode error %v", err2)
    53  					continue
    54  				}
    55  				if c.cf != nil {
    56  					c.cf(&cm)
    57  				}
    58  			}
    59  		}(consumer)
    60  
    61  		consumer2, err := c.consumer.ConsumePartition(KafkaChannelMessageTopic, partition, sarama.OffsetNewest)
    62  		if err != nil {
    63  			return err
    64  		}
    65  
    66  		go func(pc sarama.PartitionConsumer) {
    67  			for m := range pc.Messages() {
    68  				var cm = messages.ChatMessage{}
    69  				err2 := messages.JsonCodec.Decode(m.Value, &cm)
    70  				if err2 != nil {
    71  					logger.E("message decode error %v", err2)
    72  					continue
    73  				}
    74  				if c.channelCf != nil {
    75  					c.channelCf(&cm)
    76  				}
    77  			}
    78  		}(consumer2)
    79  
    80  		consumer3, err := c.consumer.ConsumePartition(KafkaChatOfflineMessageTopic, partition, sarama.OffsetNewest)
    81  		if err != nil {
    82  			return err
    83  		}
    84  
    85  		go func(pc sarama.PartitionConsumer) {
    86  			for m := range pc.Messages() {
    87  				var cm = messages.ChatMessage{}
    88  				err2 := messages.JsonCodec.Decode(m.Value, &cm)
    89  				if err2 != nil {
    90  					logger.E("message decode error %v", err2)
    91  					continue
    92  				}
    93  				if c.offlineCf != nil {
    94  					c.offlineCf(&cm)
    95  				}
    96  			}
    97  		}(consumer3)
    98  	}
    99  
   100  	return nil
   101  }
   102  
   103  func (c *KafkaConsumer) Close() error {
   104  	return c.consumer.Close()
   105  }
   106  
   107  func (c *KafkaConsumer) ConsumeChatMessage(cf func(m *messages.ChatMessage)) {
   108  	c.cf = cf
   109  }
   110  
   111  func (c *KafkaConsumer) ConsumeChannelMessage(cf func(m *messages.ChatMessage)) {
   112  	c.channelCf = cf
   113  }
   114  
   115  func (c *KafkaConsumer) ConsumeOfflineMessage(cf func(m *messages.ChatMessage)) {
   116  	c.offlineCf = cf
   117  }