github.com/aaabigfish/gopkg@v1.1.0/mq/nsq/sub.go (about)

     1  package nsq
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/nsqio/go-nsq"
     7  )
     8  
     9  type HandlerFunc func(*NsqMessage) error
    10  
    11  type Handler interface {
    12  	HandleMessage(message *NsqMessage) error
    13  }
    14  
    15  // HandleMessage implements the Handler interface
    16  func (h HandlerFunc) HandleMessage(m *NsqMessage) error {
    17  	return h(m)
    18  }
    19  
    20  type Reader interface {
    21  	GetConsumer() *nsq.Consumer
    22  	AddHandler(HandlerFunc)
    23  	AddHandlers(int, HandlerFunc)
    24  	Run() error
    25  	Close()
    26  }
    27  
    28  type reader struct {
    29  	topic    string
    30  	channel  string
    31  	addrs    []string
    32  	c        *NsqConfig
    33  	consumer *nsq.Consumer
    34  }
    35  
    36  func NewConsumer(topic, channel string, addrs []string, c ...*NsqConfig) Reader {
    37  	cfg := &NsqConfig{}
    38  	if len(c) > 0 {
    39  		cfg = c[0]
    40  	} else {
    41  		cfg = NewNsqConfig()
    42  	}
    43  
    44  	consumer, err := nsq.NewConsumer(topic, channel, cfg)
    45  	if err != nil {
    46  		panic(fmt.Sprintf("new nsq consumer failed err(%v)", err))
    47  	}
    48  
    49  	return &reader{
    50  		topic:    topic,
    51  		channel:  channel,
    52  		addrs:    addrs,
    53  		c:        cfg,
    54  		consumer: consumer,
    55  	}
    56  }
    57  
    58  func (r *reader) GetConsumer() *nsq.Consumer {
    59  	return r.consumer
    60  }
    61  
    62  func (r *reader) AddHandler(hookFunc HandlerFunc) {
    63  	r.consumer.AddHandler(HandlerFunc(hookFunc))
    64  }
    65  
    66  // 并发处理消息，n是并发数量
    67  func (r *reader) AddHandlers(n int, hookFunc HandlerFunc) {
    68  	r.consumer.AddConcurrentHandlers(HandlerFunc(hookFunc), n)
    69  }
    70  
    71  func (r *reader) Run() error {
    72  	return r.consumer.ConnectToNSQLookupds(r.addrs)
    73  }
    74  
    75  func (r *reader) Close() {
    76  	if r.consumer != nil {
    77  		r.consumer.Stop()
    78  		return
    79  	}
    80  }