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 }