github.com/sereiner/library@v0.0.0-20200518095232-1fa3e640cc5f/mq/redis/redis.consumer.go (about)

     1  package redis
     2  
     3  import (
     4  	"strings"
     5  	"sync"
     6  	"time"
     7  
     8  	"errors"
     9  
    10  	"github.com/sereiner/library/concurrent/cmap"
    11  	logger "github.com/sereiner/library/log"
    12  	"github.com/sereiner/library/mq"
    13  	"github.com/sereiner/library/redis"
    14  	"github.com/zkfy/stompngo"
    15  )
    16  
    17  type consumerChan struct {
    18  	msgChan     <-chan stompngo.MessageData
    19  	unconsumeCh chan struct{}
    20  }
    21  
    22  //RedisConsumer Consumer
    23  type RedisConsumer struct {
    24  	address    string
    25  	client     *redis.Client
    26  	queues     cmap.ConcurrentMap
    27  	connecting bool
    28  	closeCh    chan struct{}
    29  	done       bool
    30  	lk         sync.Mutex
    31  	header     []string
    32  	once       sync.Once
    33  	*mq.OptionConf
    34  }
    35  
    36  //NewRedisConsumer 创建新的Consumer
    37  func NewRedisConsumer(address string, opts ...mq.Option) (consumer *RedisConsumer, err error) {
    38  	consumer = &RedisConsumer{address: address}
    39  	consumer.OptionConf = &mq.OptionConf{Logger: logger.GetSession("mq.redis", logger.CreateSession())}
    40  	consumer.closeCh = make(chan struct{})
    41  	consumer.queues = cmap.New(2)
    42  	for _, opt := range opts {
    43  		opt(consumer.OptionConf)
    44  	}
    45  	return
    46  }
    47  
    48  //Connect  连接服务器
    49  func (consumer *RedisConsumer) Connect() (err error) {
    50  	consumer.client, err = redis.NewClientByJSON(consumer.Raw)
    51  	return
    52  }
    53  
    54  //Consume 注册消费信息
    55  func (consumer *RedisConsumer) Consume(queue string, concurrency int, callback func(mq.IMessage)) (err error) {
    56  	if strings.EqualFold(queue, "") {
    57  		return errors.New("队列名字不能为空")
    58  	}
    59  	if callback == nil {
    60  		return errors.New("回调函数不能为nil")
    61  	}
    62  
    63  	_, _, err = consumer.queues.SetIfAbsentCb(queue, func(input ...interface{}) (c interface{}, err error) {
    64  		queue := input[0].(string)
    65  		unconsumeCh := make(chan struct{}, 1)
    66  		nconcurrency := concurrency
    67  		if concurrency <= 0 {
    68  			nconcurrency = 10
    69  		}
    70  		msgChan := make(chan *RedisMessage, nconcurrency)
    71  		for i := 0; i < nconcurrency; i++ {
    72  			go func() {
    73  			START:
    74  				for {
    75  					select {
    76  					case message, ok := <-msgChan:
    77  						if !ok {
    78  							break START
    79  						}
    80  						if concurrency == 0 {
    81  							go callback(message)
    82  						} else {
    83  							callback(message)
    84  						}
    85  					}
    86  				}
    87  			}()
    88  		}
    89  
    90  		go func() {
    91  		START:
    92  			for {
    93  				select {
    94  				case <-consumer.closeCh:
    95  					break START
    96  				case <-unconsumeCh:
    97  					break START
    98  				case <-time.After(time.Millisecond * 50):
    99  					if consumer.client != nil && !consumer.done {
   100  						cmd := consumer.client.BLPop(time.Second, queue)
   101  						message := NewRedisMessage(cmd)
   102  						if message.Has() {
   103  							msgChan <- message
   104  
   105  						}
   106  					}
   107  
   108  				}
   109  			}
   110  			close(msgChan)
   111  		}()
   112  		return unconsumeCh, nil
   113  	}, queue)
   114  	return
   115  }
   116  
   117  //UnConsume 取消注册消费
   118  func (consumer *RedisConsumer) UnConsume(queue string) {
   119  	if consumer.client == nil {
   120  		return
   121  	}
   122  	if c, ok := consumer.queues.Get(queue); ok {
   123  		close(c.(chan struct{}))
   124  	}
   125  	consumer.queues.Remove(queue)
   126  }
   127  
   128  //Close 关闭当前连接
   129  func (consumer *RedisConsumer) Close() {
   130  
   131  	consumer.once.Do(func() {
   132  		close(consumer.closeCh)
   133  	})
   134  
   135  	consumer.queues.RemoveIterCb(func(key string, value interface{}) bool {
   136  		ch := value.(chan struct{})
   137  		close(ch)
   138  		return true
   139  	})
   140  	if consumer.client == nil {
   141  		return
   142  	}
   143  	consumer.client.Close()
   144  }
   145  
   146  type redisConsumerResolver struct {
   147  }
   148  
   149  func (s *redisConsumerResolver) Resolve(address string, opts ...mq.Option) (mq.MQConsumer, error) {
   150  	return NewRedisConsumer(address, opts...)
   151  }
   152  func init() {
   153  	mq.RegisterCosnumer("redis", &redisConsumerResolver{})
   154  }