nsq源码review:go-nsq producer

总述

nsq是一个实时分布式的消息队列平台。

核心部分是一个叫nsqd的模块，它负责接收和转发消息。同时在go-nsq的包中，提供了consumer和producer的核心接口。在读nsq源码的时候，很好奇它的数据是怎么从producer给到了consumer的，因而从源码的层面梳理了一下代码的实现细节。这部分先记录一下producer和consumer的代码细节，方便后续再查看相关代码。后面准备把nsqd和nsqdlookup相关的东西记录一下，包括数据分发、数据缓存、服务发现等实现细节。

go-nsq里的producer和consumer实现的功能就是一句话，提供消息接受和分发的接口。可是它内部的实现确颇有意思。

nsq demo

学习源码仍是要先从demo起步，首先装好nsq，能够git下来源码编译或者执行下载的二进制文件。先写个producer.go，以下

package main

import (
	"github.com/nsqio/go-nsq"
	"log"
)

func main() {
	cfg := nsq.NewConfig()
	r := []byte("hello consumer")
	addr := "127.0.0.1:4150"
	p, err := nsq.NewProducer(addr, cfg)
	if err != nil {
		log.Print(err)
	}
	err = p.Publish("serving123", r)
	if err != nil {
		log.Println(err)
	}
}
复制代码

例子里，我建立了一个producer，而后向serving123的topic里发送消息。

下面是consumer的代码，consumer.go

import (
	"fmt"
	"github.com/nsqio/go-nsq"
	"math/rand"
	"os"
	"time"
)

type SimpleHandler struct {
}

func (sh *SimpleHandler) HandleMessage(m *nsq.Message) error {
	_, err := os.Stdout.Write(m.Body)
	if err != nil {
		fmt.Println(err)
	}
	return nil
}

func main() {
	pause := make(chan bool)
	caddr := "127.0.0.1:4150"
	cfg := nsq.NewConfig()
	channel := fmt.Sprintf("tail%06d#ephemeral", rand.Int()%999999)
	c, _ := nsq.NewConsumer("serving123", channel, cfg)
	c.AddHandler(&SimpleHandler{})
	c.ConnectToNSQD(caddr)
	<-pause
}
复制代码

consumer，我定义了一个SimpleHandler，它实现了HandleMessage的接口，功能就是向标准输出打印消息。 main函数里，我新建了一个想要消费serving123这个topic的consumer，而后链接上nsqd服务。

运行

直接执行nsqd启一个服务。先执行consumer，而后执行producer，能够看到consumer里打印出的消息。

producer源码解析

完成上面的demo以后，先来看一下producer是怎么玩的。

producer实例，会去调用Publish方法去发布消息，这个方法接收了topic和message body。

func (w *Producer) Publish(topic string, body []byte) error {
	return w.sendCommand(Publish(topic, body))
}
复制代码

内部的另一个Publish方法执行结果做为参数传入sendCommand，Publish方法，构造了一个Command的三元消息体

func Publish(topic string, body []byte) *Command {
	var params = [][]byte{[]byte(topic)}
	return &Command{[]byte("PUB"), params, body}
}
复制代码

sendCommand方法里，初始化一个名叫doneChan的ProducerTransaction的指针，而后，cmd消息三元体，和doneChan一块儿传入sendCommandAsync方法里。最后监听doneChan的管道输出，而后返回error。

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
}
复制代码

ProducerTransaction结构体颇有意思，它持有一个本身相同类型的指针。目的是将最终的内容本身保存起来。在最终clearup的时候，它会调用一个finish方法去触发上面的chan监听返回，从而最终返回退出sendCommand。

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
}
复制代码

这个finish方法会把ProducerTransaction它本身自己进行传递，也就是在t.doneChan <- t这里

func (t *ProducerTransaction) finish() {
	if t.doneChan != nil {
		t.doneChan <- t
	}
}
复制代码

看一下sendCommandAsync函数。首先，它会去调用原子操做atomic.AddInt32去记录并发producer数量。上面说的doneChan指针，做为一个新的ProducerTransaction的参数传入，最终这个新的ProducerTransaction的数据传给了w.transactionChan监听的channel。

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()//创建和nsqd的链接 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 } 复制代码

创建和nsqd的链接，以及发送数据都写在w.connect()。了解了上面channel异步操做，看下这个connect函数

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
}
复制代码

函数开始部分很好理解，就是进行一些状态验证。_, err := w.conn.Connect()这里会实际去创建和nsqd的链接，会在里面跑一个readLoop和一个writeloop去进行读写的相关操做，东西比较多就再也不赘述。这里关心的是这个w.router方法。

router方法里面开了一个for循环来监听producer的channel，包括transactionChan、responseChan、errorChan、closeChan、exitChan，若是是w.transactionChan，则把这个transaction塞进producer的transactions数组里，而后向conn里去写消息即向nsqd发送数据。若是是收到了返回信号或者是错误信号，则会弹出一个transaction。若是收到关闭或者是退出信号，则到exit里面，清理全部transaction，并退出。

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")
}
复制代码

popTransaction方法，会把transactions第一个元素弹出，保存剩下的元素。而后回去调用这个弹出的ProducerTransaction的finish()方法，也就是上面说的finish()方法。上面的sendCommand方法会受到通知退出，这样就完成了消息的发布过程。

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()
}
复制代码

总结

限于篇幅，consumer相关准备起另一篇去写，consumer和producer在和nsqd通讯的地方会复用一些代码。

关于nsq的producer，咱们能学习到的是它channel的使用技巧，以及进行数据传输的时候，是如何运载数据和如何监听和通知channel。核心的部分在于ProducerTransaction这个结构体，它负责了消息的运载。在上面提到的readLoop和writeloop，里面有许多代理的方法也值得关注和学习，具体内容能够本身看一眼。