feature: added nsq as broker for events (#6740)

vendor/github.com/nsqio/go-nsq/producer.go

@@ -0,0 +1,393 @@
+package nsq
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type producerConn interface {
+	String() string
+	SetLogger(logger, LogLevel, string)
+	Connect() (*IdentifyResponse, error)
+	Close() error
+	WriteCommand(*Command) error
+}
+
+// Producer is a high-level type to publish to NSQ.
+//
+// A Producer instance is 1:1 with a destination `nsqd`
+// and will lazily connect to that instance (and re-connect)
+// when Publish commands are executed.
+type Producer struct {
+	id     int64
+	addr   string
+	conn   producerConn
+	config Config
+
+	logger   logger
+	logLvl   LogLevel
+	logGuard sync.RWMutex
+
+	responseChan chan []byte
+	errorChan    chan []byte
+	closeChan    chan int
+
+	transactionChan chan *ProducerTransaction
+	transactions    []*ProducerTransaction
+	state           int32
+
+	concurrentProducers int32
+	stopFlag            int32
+	exitChan            chan int
+	wg                  sync.WaitGroup
+	guard               sync.Mutex
+}
+
+// ProducerTransaction is returned by the async publish methods
+// to retrieve metadata about the command after the
+// response is received.
+type ProducerTransaction struct {
+	cmd      *Command
+	doneChan chan *ProducerTransaction
+	Error    error         // the error (or nil) of the publish command
+	Args     []interface{} // the slice of variadic arguments passed to PublishAsync or MultiPublishAsync
+}
+
+func (t *ProducerTransaction) finish() {
+	if t.doneChan != nil {
+		t.doneChan <- t
+	}
+}
+
+// NewProducer returns an instance of Producer for the specified address
+//
+// The only valid way to create a Config is via NewConfig, using a struct literal will panic.
+// After Config is passed into NewProducer the values are no longer mutable (they are copied).
+func NewProducer(addr string, config *Config) (*Producer, error) {
+	config.assertInitialized()
+	err := config.Validate()
+	if err != nil {
+		return nil, err
+	}
+
+	p := &Producer{
+		id: atomic.AddInt64(&instCount, 1),
+
+		addr:   addr,
+		config: *config,
+
+		logger: log.New(os.Stderr, "", log.Flags()),
+		logLvl: LogLevelInfo,
+
+		transactionChan: make(chan *ProducerTransaction),
+		exitChan:        make(chan int),
+		responseChan:    make(chan []byte),
+		errorChan:       make(chan []byte),
+	}
+	return p, nil
+}
+
+// Ping causes the Producer to connect to it's configured nsqd (if not already
+// connected) and send a `Nop` command, returning any error that might occur.
+//
+// This method can be used to verify that a newly-created Producer instance is
+// configured correctly, rather than relying on the lazy "connect on Publish"
+// behavior of a Producer.
+func (w *Producer) Ping() error {
+	if atomic.LoadInt32(&w.state) != StateConnected {
+		err := w.connect()
+		if err != nil {
+			return err
+		}
+	}
+
+	return w.conn.WriteCommand(Nop())
+}
+
+// SetLogger assigns the logger to use as well as a level
+//
+// The logger parameter is an interface that requires the following
+// method to be implemented (such as the the stdlib log.Logger):
+//
+//    Output(calldepth int, s string)
+//
+func (w *Producer) SetLogger(l logger, lvl LogLevel) {
+	w.logGuard.Lock()
+	defer w.logGuard.Unlock()
+
+	w.logger = l
+	w.logLvl = lvl
+}
+
+func (w *Producer) getLogger() (logger, LogLevel) {
+	w.logGuard.RLock()
+	defer w.logGuard.RUnlock()
+
+	return w.logger, w.logLvl
+}
+
+// String returns the address of the Producer
+func (w *Producer) String() string {
+	return w.addr
+}
+
+// Stop initiates a graceful stop of the Producer (permanent)
+//
+// NOTE: this blocks until completion
+func (w *Producer) Stop() {
+	w.guard.Lock()
+	if !atomic.CompareAndSwapInt32(&w.stopFlag, 0, 1) {
+		w.guard.Unlock()
+		return
+	}
+	w.log(LogLevelInfo, "stopping")
+	close(w.exitChan)
+	w.close()
+	w.guard.Unlock()
+	w.wg.Wait()
+}
+
+// PublishAsync publishes a message body to the specified topic
+// but does not wait for the response from `nsqd`.
+//
+// When the Producer eventually receives the response from `nsqd`,
+// the supplied `doneChan` (if specified)
+// will receive a `ProducerTransaction` instance with the supplied variadic arguments
+// and the response error if present
+func (w *Producer) PublishAsync(topic string, body []byte, doneChan chan *ProducerTransaction,
+	args ...interface{}) error {
+	return w.sendCommandAsync(Publish(topic, body), doneChan, args)
+}
+
+// MultiPublishAsync publishes a slice of message bodies to the specified topic
+// but does not wait for the response from `nsqd`.
+//
+// When the Producer eventually receives the response from `nsqd`,
+// the supplied `doneChan` (if specified)
+// will receive a `ProducerTransaction` instance with the supplied variadic arguments
+// and the response error if present
+func (w *Producer) MultiPublishAsync(topic string, body [][]byte, doneChan chan *ProducerTransaction,
+	args ...interface{}) error {
+	cmd, err := MultiPublish(topic, body)
+	if err != nil {
+		return err
+	}
+	return w.sendCommandAsync(cmd, doneChan, args)
+}
+
+// Publish synchronously publishes a message body to the specified topic, returning
+// an error if publish failed
+func (w *Producer) Publish(topic string, body []byte) error {
+	return w.sendCommand(Publish(topic, body))
+}
+
+// MultiPublish synchronously publishes a slice of message bodies to the specified topic, returning
+// an error if publish failed
+func (w *Producer) MultiPublish(topic string, body [][]byte) error {
+	cmd, err := MultiPublish(topic, body)
+	if err != nil {
+		return err
+	}
+	return w.sendCommand(cmd)
+}
+
+// DeferredPublish synchronously publishes a message body to the specified topic
+// where the message will queue at the channel level until the timeout expires, returning
+// an error if publish failed
+func (w *Producer) DeferredPublish(topic string, delay time.Duration, body []byte) error {
+	return w.sendCommand(DeferredPublish(topic, delay, body))
+}
+
+// DeferredPublishAsync publishes a message body to the specified topic
+// where the message will queue at the channel level until the timeout expires
+// but does not wait for the response from `nsqd`.
+//
+// When the Producer eventually receives the response from `nsqd`,
+// the supplied `doneChan` (if specified)
+// will receive a `ProducerTransaction` instance with the supplied variadic arguments
+// and the response error if present
+func (w *Producer) DeferredPublishAsync(topic string, delay time.Duration, body []byte,
+	doneChan chan *ProducerTransaction, args ...interface{}) error {
+	return w.sendCommandAsync(DeferredPublish(topic, delay, body), doneChan, args)
+}
+
+func (w *Producer) sendCommand(cmd *Command) error {
+	doneChan := make(chan *ProducerTransaction)
+	err := w.sendCommandAsync(cmd, doneChan, nil)
+	if err != nil {
+		close(doneChan)
+		return err
+	}
+	t := <-doneChan
+	return t.Error
+}
+
+func (w *Producer) sendCommandAsync(cmd *Command, doneChan chan *ProducerTransaction,
+	args []interface{}) error {
+	// keep track of how many outstanding producers we're dealing with
+	// in order to later ensure that we clean them all up...
+	atomic.AddInt32(&w.concurrentProducers, 1)
+	defer atomic.AddInt32(&w.concurrentProducers, -1)
+
+	if atomic.LoadInt32(&w.state) != StateConnected {
+		err := w.connect()
+		if err != nil {
+			return err
+		}
+	}
+
+	t := &ProducerTransaction{
+		cmd:      cmd,
+		doneChan: doneChan,
+		Args:     args,
+	}
+
+	select {
+	case w.transactionChan <- t:
+	case <-w.exitChan:
+		return ErrStopped
+	}
+
+	return nil
+}
+
+func (w *Producer) connect() error {
+	w.guard.Lock()
+	defer w.guard.Unlock()
+
+	if atomic.LoadInt32(&w.stopFlag) == 1 {
+		return ErrStopped
+	}
+
+	switch state := atomic.LoadInt32(&w.state); state {
+	case StateInit:
+	case StateConnected:
+		return nil
+	default:
+		return ErrNotConnected
+	}
+
+	w.log(LogLevelInfo, "(%s) connecting to nsqd", w.addr)
+
+	logger, logLvl := w.getLogger()
+
+	w.conn = NewConn(w.addr, &w.config, &producerConnDelegate{w})
+	w.conn.SetLogger(logger, logLvl, fmt.Sprintf("%3d (%%s)", w.id))
+
+	_, err := w.conn.Connect()
+	if err != nil {
+		w.conn.Close()
+		w.log(LogLevelError, "(%s) error connecting to nsqd - %s", w.addr, err)
+		return err
+	}
+	atomic.StoreInt32(&w.state, StateConnected)
+	w.closeChan = make(chan int)
+	w.wg.Add(1)
+	go w.router()
+
+	return nil
+}
+
+func (w *Producer) close() {
+	if !atomic.CompareAndSwapInt32(&w.state, StateConnected, StateDisconnected) {
+		return
+	}
+	w.conn.Close()
+	go func() {
+		// we need to handle this in a goroutine so we don't
+		// block the caller from making progress
+		w.wg.Wait()
+		atomic.StoreInt32(&w.state, StateInit)
+	}()
+}
+
+func (w *Producer) router() {
+	for {
+		select {
+		case t := <-w.transactionChan:
+			w.transactions = append(w.transactions, t)
+			err := w.conn.WriteCommand(t.cmd)
+			if err != nil {
+				w.log(LogLevelError, "(%s) sending command - %s", w.conn.String(), err)
+				w.close()
+			}
+		case data := <-w.responseChan:
+			w.popTransaction(FrameTypeResponse, data)
+		case data := <-w.errorChan:
+			w.popTransaction(FrameTypeError, data)
+		case <-w.closeChan:
+			goto exit
+		case <-w.exitChan:
+			goto exit
+		}
+	}
+
+exit:
+	w.transactionCleanup()
+	w.wg.Done()
+	w.log(LogLevelInfo, "exiting router")
+}
+
+func (w *Producer) popTransaction(frameType int32, data []byte) {
+	t := w.transactions[0]
+	w.transactions = w.transactions[1:]
+	if frameType == FrameTypeError {
+		t.Error = ErrProtocol{string(data)}
+	}
+	t.finish()
+}
+
+func (w *Producer) transactionCleanup() {
+	// clean up transactions we can easily account for
+	for _, t := range w.transactions {
+		t.Error = ErrNotConnected
+		t.finish()
+	}
+	w.transactions = w.transactions[:0]
+
+	// spin and free up any writes that might have raced
+	// with the cleanup process (blocked on writing
+	// to transactionChan)
+	for {
+		select {
+		case t := <-w.transactionChan:
+			t.Error = ErrNotConnected
+			t.finish()
+		default:
+			// keep spinning until there are 0 concurrent producers
+			if atomic.LoadInt32(&w.concurrentProducers) == 0 {
+				return
+			}
+			// give the runtime a chance to schedule other racing goroutines
+			time.Sleep(5 * time.Millisecond)
+		}
+	}
+}
+
+func (w *Producer) log(lvl LogLevel, line string, args ...interface{}) {
+	logger, logLvl := w.getLogger()
+
+	if logger == nil {
+		return
+	}
+
+	if logLvl > lvl {
+		return
+	}
+
+	logger.Output(2, fmt.Sprintf("%-4s %3d %s", lvl, w.id, fmt.Sprintf(line, args...)))
+}
+
+func (w *Producer) onConnResponse(c *Conn, data []byte) { w.responseChan <- data }
+func (w *Producer) onConnError(c *Conn, data []byte)    { w.errorChan <- data }
+func (w *Producer) onConnHeartbeat(c *Conn)             {}
+func (w *Producer) onConnIOError(c *Conn, err error)    { w.close() }
+func (w *Producer) onConnClose(c *Conn) {
+	w.guard.Lock()
+	defer w.guard.Unlock()
+	close(w.closeChan)
+}