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 }