github.com/qxnw/lib4go@v0.0.0-20180426074627-c80c7e84b925/mq/mqtt/mqtt.consumer.go (about)

     1  package mqtt
     2  
     3  import (
     4  	"net"
     5  	"strings"
     6  	"sync"
     7  
     8  	"errors"
     9  
    10  	proto "github.com/huin/mqtt"
    11  	"github.com/jeffallen/mqtt"
    12  	"github.com/qxnw/lib4go/concurrent/cmap"
    13  	"github.com/qxnw/lib4go/logger"
    14  	"github.com/qxnw/lib4go/mq"
    15  	xnet "github.com/qxnw/lib4go/net"
    16  	"github.com/zkfy/stompngo"
    17  )
    18  
    19  type consumerChan struct {
    20  	msgChan     <-chan stompngo.MessageData
    21  	unconsumeCh chan struct{}
    22  }
    23  
    24  //Consumer Consumer
    25  type Consumer struct {
    26  	address    string
    27  	client     *mqtt.ClientConn
    28  	queues     cmap.ConcurrentMap
    29  	subChan    chan string
    30  	connecting bool
    31  	closeCh    chan struct{}
    32  	done       bool
    33  	lk         sync.Mutex
    34  	header     []string
    35  	once       sync.Once
    36  	*mq.OptionConf
    37  	conf *Conf
    38  }
    39  
    40  //NewConsumer 创建新的Consumer
    41  func NewConsumer(address string, opts ...mq.Option) (consumer *Consumer, err error) {
    42  	consumer = &Consumer{address: address}
    43  	consumer.OptionConf = &mq.OptionConf{Logger: logger.GetSession("mqtt", logger.CreateSession())}
    44  	consumer.closeCh = make(chan struct{})
    45  	consumer.queues = cmap.New(2)
    46  	consumer.subChan = make(chan string, 3)
    47  	for _, opt := range opts {
    48  		opt(consumer.OptionConf)
    49  	}
    50  	consumer.conf, err = NewConf(consumer.Raw)
    51  	return
    52  }
    53  
    54  //Connect  连接服务器
    55  func (consumer *Consumer) Connect() (err error) {
    56  	conn, err := net.Dial("tcp", consumer.conf.Address)
    57  	if err != nil {
    58  		return err
    59  	}
    60  	cc := mqtt.NewClientConn(conn)
    61  	cc.Dump = consumer.conf.DumpData
    62  	cc.ClientId = xnet.GetLocalIPAddress()
    63  	if err = cc.Connect(consumer.conf.UserName, consumer.conf.Password); err != nil {
    64  		return err
    65  	}
    66  	consumer.client = cc
    67  	go consumer.recvMessage()
    68  	return nil
    69  }
    70  
    71  //recvMessage 循环接收，并放入指定的队列
    72  func (consumer *Consumer) recvMessage() {
    73  
    74  START:
    75  	for {
    76  		select {
    77  		case <-consumer.closeCh:
    78  			break START
    79  		case q := <-consumer.subChan:
    80  			tq := make([]proto.TopicQos, 1)
    81  			tq[0].Topic = q
    82  			tq[0].Qos = proto.QosAtMostOnce
    83  			consumer.client.Subscribe(tq)
    84  		case msg := <-consumer.client.Incoming:
    85  			nmsg := NewMessage()
    86  			if err := msg.Payload.WritePayload(nmsg); err != nil {
    87  				consumer.Logger.Error(err)
    88  				continue
    89  			}
    90  			if nq, b := consumer.queues.Get(msg.TopicName); b {
    91  				nQ := nq.(chan *Message)
    92  				nQ <- nmsg
    93  			}
    94  		}
    95  	}
    96  }
    97  
    98  //Consume 注册消费信息
    99  func (consumer *Consumer) Consume(queue string, concurrency int, callback func(mq.IMessage)) (err error) {
   100  	if strings.EqualFold(queue, "") {
   101  		return errors.New("队列名字不能为空")
   102  	}
   103  	if callback == nil {
   104  		return errors.New("回调函数不能为nil")
   105  	}
   106  
   107  	_, _, err = consumer.queues.SetIfAbsentCb(queue, func(input ...interface{}) (c interface{}, err error) {
   108  		queue := input[0].(string)
   109  		if concurrency <= 0 {
   110  			concurrency = 10
   111  		}
   112  		msgChan := make(chan *Message, concurrency)
   113  		for i := 0; i < concurrency; i++ {
   114  			go func() {
   115  			START:
   116  				for {
   117  					select {
   118  					case message, ok := <-msgChan:
   119  						if !ok {
   120  							break START
   121  						}
   122  						go callback(message)
   123  					}
   124  				}
   125  			}()
   126  		}
   127  		consumer.subChan <- queue
   128  		return msgChan, nil
   129  	}, queue)
   130  	return
   131  }
   132  
   133  //UnConsume 取消注册消费
   134  func (consumer *Consumer) UnConsume(queue string) {
   135  
   136  }
   137  
   138  //Close 关闭当前连接
   139  func (consumer *Consumer) Close() {
   140  	consumer.once.Do(func() {
   141  		close(consumer.closeCh)
   142  	})
   143  
   144  	consumer.queues.RemoveIterCb(func(key string, value interface{}) bool {
   145  		ch := value.(chan *Message)
   146  		close(ch)
   147  		return true
   148  	})
   149  	if consumer.client == nil {
   150  		return
   151  	}
   152  	consumer.client.Disconnect()
   153  }
   154  
   155  type ConsumerResolver struct {
   156  }
   157  
   158  func (s *ConsumerResolver) Resolve(address string, opts ...mq.Option) (mq.MQConsumer, error) {
   159  	return NewConsumer(address, opts...)
   160  }
   161  func init() {
   162  	mq.RegisterCosnumer("mqtt", &ConsumerResolver{})
   163  }