github.com/sereiner/library@v0.0.0-20200518095232-1fa3e640cc5f/mq/mqtt/mqtt.consumer.go (about) 1 package mqtt 2 3 import ( 4 "crypto/tls" 5 "crypto/x509" 6 "errors" 7 "fmt" 8 "io/ioutil" 9 xnet "net" 10 "strings" 11 "sync" 12 "time" 13 14 "github.com/sereiner/library/concurrent/cmap" 15 logger "github.com/sereiner/library/log" 16 "github.com/sereiner/library/mq" 17 "github.com/sereiner/library/net" 18 "github.com/sereiner/library/utility" 19 "github.com/yosssi/gmq/mqtt" 20 "github.com/yosssi/gmq/mqtt/client" 21 "github.com/zkfy/stompngo" 22 ) 23 24 type consumerChan struct { 25 msgChan <-chan stompngo.MessageData 26 unconsumeCh chan struct{} 27 } 28 29 //Consumer Consumer 30 type Consumer struct { 31 address string 32 client *client.Client 33 queues cmap.ConcurrentMap 34 subChan chan string 35 connecting bool 36 closeCh chan struct{} 37 connCh chan int 38 done bool 39 lk sync.Mutex 40 header []string 41 once sync.Once 42 *mq.OptionConf 43 conf *Conf 44 } 45 46 //NewConsumer 创建新的Consumer 47 func NewConsumer(address string, opts ...mq.Option) (consumer *Consumer, err error) { 48 consumer = &Consumer{address: address} 49 consumer.OptionConf = &mq.OptionConf{QueueCount: 250, Logger: logger.GetSession("mqtt.consumer", logger.CreateSession())} 50 consumer.closeCh = make(chan struct{}) 51 consumer.connCh = make(chan int, 1) 52 consumer.queues = cmap.New(2) 53 for _, opt := range opts { 54 opt(consumer.OptionConf) 55 } 56 consumer.subChan = make(chan string, consumer.OptionConf.QueueCount) 57 consumer.conf, err = NewConf(consumer.Raw) 58 if err != nil { 59 return nil, err 60 } 61 return consumer, nil 62 } 63 64 //Connect 连接服务器 65 func (consumer *Consumer) Connect() (err error) { 66 cc, _, err := consumer.connect() 67 if err != nil { 68 return err 69 } 70 consumer.client = cc 71 go consumer.reconnect() 72 go consumer.subscribe() 73 return nil 74 } 75 func (consumer *Consumer) reconnect() { 76 for { 77 select { 78 case <-time.After(time.Second * 3): //延迟重连 79 select { 80 case <-consumer.connCh: 81 consumer.Logger.Debug("consumer与服务器断开连接,准备重连") 82 func() { 83 defer recover() 84 consumer.client.Disconnect() 85 consumer.client.Terminate() 86 }() 87 client, b, err := consumer.connect() 88 if err != nil { 89 consumer.Logger.Error("连接失败:", err) 90 } 91 if b { 92 consumer.Logger.Info("consumer成功连接到服务器") 93 consumer.client = client 94 consumer.queues.IterCb(func(k string, v interface{}) bool { 95 consumer.subChan <- k 96 return true 97 }) 98 } 99 default: 100 101 } 102 } 103 } 104 } 105 106 func (consumer *Consumer) connect() (*client.Client, bool, error) { 107 consumer.lk.Lock() 108 defer consumer.lk.Unlock() 109 cert, err := consumer.getCert(consumer.conf) 110 if err != nil { 111 return nil, false, err 112 } 113 cc := client.New(&client.Options{ 114 ErrorHandler: func(err error) { 115 select { 116 case consumer.connCh <- 1: //发送重连消息 117 default: 118 } 119 }, 120 }) 121 host, port, err := xnet.SplitHostPort(consumer.conf.Address) 122 if err != nil { 123 return nil, false, err 124 } 125 addrs, err := xnet.LookupHost(host) 126 if err != nil { 127 return nil, false, err 128 } 129 if err != nil { 130 return nil, false, err 131 } 132 for _, addr := range addrs { 133 if err := cc.Connect(&client.ConnectOptions{ 134 Network: "tcp", 135 Address: addr + ":" + port, 136 UserName: []byte(consumer.conf.UserName), 137 Password: []byte(consumer.conf.Password), 138 ClientID: []byte(fmt.Sprintf("%s-%s", net.GetLocalIPAddress(), utility.GetGUID()[0:6])), 139 TLSConfig: cert, 140 KeepAlive: 3, 141 }); err == nil { 142 return cc, true, nil 143 } 144 } 145 return nil, false, fmt.Errorf("连接失败:%v[%v](%s-%s/%s)", err, consumer.conf.Address, addrs, consumer.conf.UserName, consumer.conf.Password) 146 } 147 148 func (consumer *Consumer) getCert(conf *Conf) (*tls.Config, error) { 149 if conf.CertPath == "" { 150 return nil, nil 151 } 152 b, err := ioutil.ReadFile(conf.CertPath) 153 if err != nil { 154 return nil, fmt.Errorf("读取证书失败:%s(%v)", conf.CertPath, err) 155 } 156 roots := x509.NewCertPool() 157 if ok := roots.AppendCertsFromPEM(b); !ok { 158 return nil, fmt.Errorf("failed to parse root certificate") 159 } 160 return &tls.Config{ 161 RootCAs: roots, 162 }, nil 163 } 164 165 //subscribe 循环接收,并放入指定的队列 166 // QoS0,最多一次送达。也就是发出去就fire掉,没有后面的事情了。 167 // QoS1,至少一次送达。发出去之后必须等待ack,没有ack,就要找时机重发 168 // QoS2,准确一次送达。消息id将拥有一个简单的生命周期。 169 func (consumer *Consumer) subscribe() { 170 171 START: 172 for { 173 select { 174 case <-consumer.closeCh: 175 break START 176 case q := <-consumer.subChan: 177 err := consumer.client.Subscribe(&client.SubscribeOptions{ 178 SubReqs: []*client.SubReq{ 179 &client.SubReq{ 180 TopicFilter: []byte(q), 181 QoS: mqtt.QoS0, 182 Handler: func(topicName, message []byte) { 183 nmsg := NewMessage() 184 _, err := nmsg.Write(message) 185 if err != nil { 186 consumer.Logger.Error("写入消息失败:", string(message)) 187 return 188 } 189 if nq, b := consumer.queues.Get(string(topicName)); b { 190 nQ := nq.(chan *Message) 191 nQ <- nmsg 192 } 193 }, 194 }, 195 }, 196 }) 197 if err != nil { 198 consumer.Logger.Error("消息订阅出错", err) 199 } 200 } 201 } 202 } 203 204 //Consume 注册消费信息 205 func (consumer *Consumer) Consume(queue string, concurrency int, callback func(mq.IMessage)) (err error) { 206 if strings.EqualFold(queue, "") { 207 return errors.New("队列名字不能为空") 208 } 209 if callback == nil { 210 return errors.New("回调函数不能为nil") 211 } 212 213 _, _, err = consumer.queues.SetIfAbsentCb(queue, func(input ...interface{}) (c interface{}, err error) { 214 queue := input[0].(string) 215 if concurrency <= 0 { 216 concurrency = 10 217 } 218 msgChan := make(chan *Message, concurrency) 219 for i := 0; i < concurrency; i++ { 220 go func() { 221 START: 222 for { 223 select { 224 case message, ok := <-msgChan: 225 if !ok { 226 break START 227 } 228 go callback(message) 229 } 230 } 231 }() 232 } 233 consumer.subChan <- queue 234 return msgChan, nil 235 }, queue) 236 return 237 } 238 239 //UnConsume 取消注册消费 240 func (consumer *Consumer) UnConsume(queue string) { 241 err := consumer.client.Unsubscribe(&client.UnsubscribeOptions{ 242 TopicFilters: [][]byte{ 243 []byte(queue), 244 }, 245 }) 246 if err != nil { 247 consumer.Logger.Errorf("取消订单消息出错(queue:%s)err:%v", queue, err) 248 } 249 } 250 251 //Close 关闭当前连接 252 func (consumer *Consumer) Close() { 253 consumer.once.Do(func() { 254 close(consumer.closeCh) 255 }) 256 257 consumer.queues.RemoveIterCb(func(key string, value interface{}) bool { 258 ch := value.(chan *Message) 259 close(ch) 260 return true 261 }) 262 if consumer.client == nil { 263 return 264 } 265 consumer.client.Disconnect() 266 consumer.client.Terminate() 267 } 268 269 type ConsumerResolver struct { 270 } 271 272 func (s *ConsumerResolver) Resolve(address string, opts ...mq.Option) (mq.MQConsumer, error) { 273 return NewConsumer(address, opts...) 274 } 275 func init() { 276 mq.RegisterCosnumer("mqtt", &ConsumerResolver{}) 277 }