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  }