github.com/sereiner/library@v0.0.0-20200518095232-1fa3e640cc5f/mq/nsq/nsq.consumer.go (about)

     1  package nsq
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/signal"
     7  	"strings"
     8  	"time"
     9  
    10  	"github.com/nsqio/go-nsq"
    11  
    12  	"github.com/sereiner/library/concurrent/cmap"
    13  	"github.com/sereiner/library/mq"
    14  )
    15  
    16  type NsqConsumer struct {
    17  	address   string
    18  	consumers cmap.ConcurrentMap
    19  	quitChan  chan struct{}
    20  	*mq.OptionConf
    21  }
    22  
    23  type nsqConsumer struct {
    24  	consumer *nsq.Consumer
    25  	msgQueue chan *nsq.Message
    26  	closeCh  chan struct{}
    27  }
    28  
    29  func (n *nsqConsumer) HandleMessage(message *nsq.Message) error {
    30  	time.Sleep(time.Nanosecond * 100)
    31  	n.msgQueue <- message
    32  	return nil
    33  }
    34  
    35  func (n *nsqConsumer) Messages() chan *nsq.Message {
    36  	return n.msgQueue
    37  }
    38  
    39  // NewNsqConsumer 初始化nsq Consumer
    40  func NewNsqConsumer(address string, opts ...mq.Option) (nsq *NsqConsumer, err error) {
    41  	nsq = &NsqConsumer{address: address, quitChan: make(chan struct{}, 0)}
    42  	nsq.OptionConf = &mq.OptionConf{}
    43  	nsq.consumers = cmap.New(2)
    44  	for _, opt := range opts {
    45  		opt(nsq.OptionConf)
    46  	}
    47  	return
    48  }
    49  
    50  //Connect 连接到服务器
    51  func (n *NsqConsumer) Connect() error {
    52  	return nil
    53  }
    54  
    55  //Consume 订阅消息
    56  func (n *NsqConsumer) Consume(queue string, concurrency int, call func(mq.IMessage)) (err error) {
    57  	_, cname, _ := n.consumers.SetIfAbsentCb(queue, func(i ...interface{}) (interface{}, error) {
    58  
    59  		c := &nsqConsumer{}
    60  		queueArr := strings.Split(queue, "#")
    61  		if len(queueArr) != 2 {
    62  			panic("nsq 消息队列格式错误 topic#channel")
    63  		}
    64  
    65  		config := nsq.NewConfig()
    66  		c.consumer, err = nsq.NewConsumer(queueArr[0], queueArr[1], config)
    67  		c.msgQueue = make(chan *nsq.Message, 10000)
    68  		c.closeCh = make(chan struct{}, 1)
    69  		c.consumer.AddHandler(c)
    70  		err := c.consumer.ConnectToNSQD(n.address)
    71  		if err != nil {
    72  			fmt.Println(err)
    73  			return nil, err
    74  		}
    75  		return c, nil
    76  	})
    77  
    78  	consumer := cname.(*nsqConsumer)
    79  
    80  	var (
    81  		chanmsg = make(chan *nsq.Message, 10000)
    82  		closing = make(chan struct{})
    83  	)
    84  
    85  	go func() {
    86  		signals := make(chan os.Signal, 1)
    87  		signal.Notify(signals, os.Kill, os.Interrupt)
    88  		<-signals
    89  		fmt.Println("Initiating shutdown of consumer...")
    90  		close(closing)
    91  	}()
    92  
    93  	go func(consumer *nsqConsumer) {
    94  		for {
    95  			select {
    96  			case <-consumer.closeCh:
    97  				break
    98  			case message, ok := <-consumer.Messages():
    99  				if ok {
   100  					chanmsg <- message
   101  				}
   102  			}
   103  		}
   104  	}(consumer)
   105  
   106  	go func() {
   107  	LOOP:
   108  		for {
   109  			select {
   110  			case <-consumer.closeCh:
   111  				break LOOP
   112  			case msg, ok := <-chanmsg:
   113  				if ok {
   114  					go call(NewNsqMessage(msg))
   115  				} else {
   116  					break LOOP
   117  				}
   118  			}
   119  		}
   120  		close(chanmsg)
   121  	}()
   122  
   123  	//if err := consumer.consumer.Close(); err != nil {
   124  	//	fmt.Println("Failed to close consumer: ", err)
   125  	//}
   126  	return nil
   127  }
   128  
   129  //UnConsume 取消注册消费
   130  func (n *NsqConsumer) UnConsume(queue string) {
   131  	if c, ok := n.consumers.Get(queue); ok {
   132  		consumer := c.(*nsqConsumer)
   133  		close(consumer.msgQueue)
   134  	}
   135  }
   136  
   137  //Close 关闭当前连接
   138  func (n *NsqConsumer) Close() {
   139  	close(n.quitChan)
   140  	n.consumers.IterCb(func(key string, value interface{}) bool {
   141  		c := value.(*nsqConsumer)
   142  		c.consumer.Stop()
   143  		time.Sleep(time.Second)
   144  		c.closeCh <- struct{}{}
   145  		close(c.msgQueue)
   146  		return true
   147  	})
   148  }
   149  
   150  type nsqConsumerResolver struct {
   151  }
   152  
   153  func (s *nsqConsumerResolver) Resolve(address string, opts ...mq.Option) (mq.MQConsumer, error) {
   154  	return NewNsqConsumer(address, opts...)
   155  }
   156  func init() {
   157  	mq.RegisterCosnumer("nsq", &nsqConsumerResolver{})
   158  }