github.com/angenalZZZ/gofunc@v0.0.0-20210507121333-48ff1be3917b/data/queue/consumer/nsq.go (about)

     1  package consumer
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/angenalZZZ/gofunc/data/queue/message"
     6  	"github.com/angenalZZZ/gofunc/log"
     7  	"github.com/nsqio/go-nsq"
     8  	"github.com/rs/zerolog"
     9  	"runtime"
    10  )
    11  
    12  // NsqConsumer NSQ messages consumer.
    13  type NsqConsumer struct {
    14  	C        chan struct{}
    15  	Log      *log.Logger
    16  	Config   *nsq.Config
    17  	handlers map[message.TMessage]*message.Queue
    18  }
    19  
    20  func NewNsqConsumer() *NsqConsumer {
    21  	var (
    22  		l  *log.Logger
    23  		ll zerolog.Level
    24  	)
    25  
    26  	if log.Log != nil {
    27  		l = log.Log
    28  	} else {
    29  		l = log.InitConsole("2006-01-02 15:04:05.000", false)
    30  		ll = zerolog.ErrorLevel
    31  		l.Level(ll)
    32  	}
    33  
    34  	return &NsqConsumer{
    35  		C:        make(chan struct{}),
    36  		Log:      l,
    37  		handlers: make(map[message.TMessage]*message.Queue),
    38  	}
    39  }
    40  
    41  // Register create topic/channel handler for messages,
    42  // This function creates a new nsq Reader.
    43  func (c *NsqConsumer) Register(topic, channel string, maxInFlight int, handler message.Handler) error {
    44  	tch := message.TMessage{Topic: topic, Channel: channel}
    45  
    46  	var config *nsq.Config
    47  	if c.Config == nil {
    48  		config = nsq.NewConfig()
    49  		_ = config.Set("verbose", false)
    50  		_ = config.Set("max_in_flight", maxInFlight)
    51  	} else {
    52  		config = c.Config
    53  	}
    54  
    55  	r, err := nsq.NewConsumer(topic, channel, config)
    56  	if err != nil {
    57  		return err
    58  	}
    59  
    60  	r.SetLogger(c, nsq.LogLevel(int(c.Log.GetLevel())))
    61  
    62  	q := &message.Queue{H: handler, Consumer: r}
    63  	r.AddConcurrentHandlers(q, maxInFlight)
    64  	c.handlers[tch] = q
    65  	return nil
    66  }
    67  
    68  // Connect - Connects all readers to NSQ.
    69  func (c *NsqConsumer) Connect(addr ...string) error {
    70  	for _, q := range c.handlers {
    71  		for _, add := range addr {
    72  			if err := q.ConnectToNSQD(add); err != nil {
    73  				return err
    74  			}
    75  		}
    76  	}
    77  	return nil
    78  }
    79  
    80  // ConnectLookupD Connects all readers to NSQ LookupD.
    81  func (c *NsqConsumer) ConnectLookupD(addr ...string) error {
    82  	for _, q := range c.handlers {
    83  		for _, add := range addr {
    84  			if err := q.ConnectToNSQLookupd(add); err != nil {
    85  				return err
    86  			}
    87  		}
    88  	}
    89  	return nil
    90  }
    91  
    92  // Start Just waits.
    93  func (c *NsqConsumer) Start() error {
    94  	if c.C == nil {
    95  		c.C = make(chan struct{})
    96  	}
    97  	<-c.C
    98  	return nil
    99  }
   100  
   101  // Stop Gracefully closes all consumers.
   102  func (c *NsqConsumer) Stop() {
   103  	for _, h := range c.handlers {
   104  		h.Stop()
   105  	}
   106  	if c.C != nil {
   107  		close(c.C)
   108  	}
   109  }
   110  
   111  // Output log.
   112  func (c *NsqConsumer) Output(calldepth int, s string) error {
   113  	if c.Log.GetLevel() == zerolog.DebugLevel {
   114  		_, file, line, ok := runtime.Caller(calldepth)
   115  		if !ok {
   116  			file = "???"
   117  			line = 0
   118  		}
   119  		s = fmt.Sprintf("%s %04d: %s", file, line, s)
   120  	}
   121  	c.Log.Print(s)
   122  	return nil
   123  }