github.com/qxnw/lib4go@v0.0.0-20180426074627-c80c7e84b925/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/qxnw/lib4go/concurrent/cmap"
    11  	"github.com/qxnw/lib4go/logger"
    12  	"github.com/qxnw/lib4go/mq"
    13  	"github.com/qxnw/lib4go/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  		if concurrency <= 0 {
    67  			concurrency = 10
    68  		}
    69  		msgChan := make(chan *RedisMessage, concurrency)
    70  		for i := 0; i < concurrency; i++ {
    71  			go func() {
    72  			START:
    73  				for {
    74  					select {
    75  					case message, ok := <-msgChan:
    76  						if !ok {
    77  							break START
    78  						}
    79  						go callback(message)
    80  					}
    81  				}
    82  			}()
    83  		}
    84  		go func() {
    85  		START:
    86  			for {
    87  				select {
    88  				case <-consumer.closeCh:
    89  					break START
    90  				case <-unconsumeCh:
    91  					break START
    92  				case <-time.After(time.Millisecond * 50):
    93  					if consumer.client != nil && !consumer.done {
    94  						cmd := consumer.client.BLPop(time.Second, queue)
    95  						message := NewRedisMessage(cmd)
    96  						if message.Has() {
    97  							msgChan <- message
    98  
    99  						}
   100  					}
   101  
   102  				}
   103  			}
   104  			close(msgChan)
   105  		}()
   106  		return unconsumeCh, nil
   107  	}, queue)
   108  	return
   109  }
   110  
   111  //UnConsume 取消注册消费
   112  func (consumer *RedisConsumer) UnConsume(queue string) {
   113  	if consumer.client == nil {
   114  		return
   115  	}
   116  	if c, ok := consumer.queues.Get(queue); ok {
   117  		close(c.(chan struct{}))
   118  	}
   119  	consumer.queues.Remove(queue)
   120  }
   121  
   122  //Close 关闭当前连接
   123  func (consumer *RedisConsumer) Close() {
   124  
   125  	consumer.once.Do(func() {
   126  		close(consumer.closeCh)
   127  	})
   128  
   129  	consumer.queues.RemoveIterCb(func(key string, value interface{}) bool {
   130  		ch := value.(chan struct{})
   131  		close(ch)
   132  		return true
   133  	})
   134  	if consumer.client == nil {
   135  		return
   136  	}
   137  	consumer.client.Close()
   138  }
   139  
   140  type redisConsumerResolver struct {
   141  }
   142  
   143  func (s *redisConsumerResolver) Resolve(address string, opts ...mq.Option) (mq.MQConsumer, error) {
   144  	return NewRedisConsumer(address, opts...)
   145  }
   146  func init() {
   147  	mq.RegisterCosnumer("redis", &redisConsumerResolver{})
   148  }