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，里面有许多代理的方法也值得关注和学习，具体内容可以自己看一眼。
About The Author
