浅析NSQ客户端consume过程源码

NSQ是一个用go语言实现的分布式消息系统，本篇博客不介绍NSQ的整体架构，这方面可以自己去了解。

因为曾经的工作运用到了NSQ，我这两天闲下来研究了一下NSQ的客户端是如何完成consume这整个过程的。
下面是一个关于客户端消费NSQ队列里消息的示例：

package main

import (
    "fmt"
    "time"

    "github.com/nsqio/go-nsq"
)

// nsq发布消息
func Producer() {
    p, err := nsq.NewProducer("127.0.0.1:4150", nsq.NewConfig())                // 新建生产者
    if err != nil {
        panic(err)
    }   
　　
    if err := p.Publish("test", []byte("hello NSQ!!!")); err != nil {           // 发布消息
        panic(err)
    }   
}

// nsq订阅消息
type ConsumerT struct{}

func (*ConsumerT) HandleMessage(msg *nsq.Message) error {
    fmt.Println(string(msg.Body))
    return nil 
}

func Consumer() {
    c, err := nsq.NewConsumer("test", "test-channel", nsq.NewConfig())   // 新建一个消费者
    if err != nil {
        panic(err)
    }   
    c.AddHandler(&ConsumerT{})                                           // 添加消息处理
    if err := c.ConnectToNSQD("127.0.0.1:4150"); err != nil {            // 建立连接
        panic(err)
    }   
}
// 主函数
func main() {
    Producer()
    Consumer()
    time.Sleep(time.Second * 3)
}
// 运行将会打印： hello NSQ!!!

上面的代码中，HandleMessage()是Handler接口中的一个函数，这个接口主要是提供给用户让他们完成自己对消息的处理（比如加密，数据存储等等业务上的处理）客户端需要自己实现这个接口的具体逻辑。

消费端主要的几个步骤就是先获得一个consumer对象然后添加消息处理逻辑，最后通过connectToNSQD()连接NSQ服务端，发起消费这个流程。

好，接下来我们来具体看看这几个函数。以下是nsq客户端的源码。
这是源码链接：https://github.com/nsqio/go-nsq

因为真正开始发起消费这一过程的是connectToNSQD()，我们先看这个函数的源码

func (r *Consumer) ConnectToNSQD(addr string) error {
    if atomic.LoadInt32(&r.stopFlag) == 1 {
        return errors.New("consumer stopped")
    }

    if atomic.LoadInt32(&r.runningHandlers) == 0 {
        return errors.New("no handlers")
    }

    atomic.StoreInt32(&r.connectedFlag, 1)

    logger, logLvl := r.getLogger()

    conn := NewConn(addr, &r.config, &consumerConnDelegate{r})//*****************************
    conn.SetLogger(logger, logLvl,
        fmt.Sprintf("%3d [%s/%s] (%%s)", r.id, r.topic, r.channel))

    r.mtx.Lock()
    _, pendingOk := r.pendingConnections[addr]
    _, ok := r.connections[addr]
    if ok || pendingOk {
        r.mtx.Unlock()
        return ErrAlreadyConnected
    }
    r.pendingConnections[addr] = conn
    if idx := indexOf(addr, r.nsqdTCPAddrs); idx == -1 {
        r.nsqdTCPAddrs = append(r.nsqdTCPAddrs, addr)
    }
    r.mtx.Unlock()

    r.log(LogLevelInfo, "(%s) connecting to nsqd", addr)

    cleanupConnection := func() {
        r.mtx.Lock()
        delete(r.pendingConnections, addr)
        r.mtx.Unlock()
        conn.Close()
    }

    resp, err := conn.Connect()//***********************
    if err != nil {
        cleanupConnection()
        return err
    }

    if resp != nil {
        if resp.MaxRdyCount < int64(r.getMaxInFlight()) {
            r.log(LogLevelWarning,
                "(%s) max RDY count %d < consumer max in flight %d, truncation possible",
                conn.String(), resp.MaxRdyCount, r.getMaxInFlight())
        }
    }

    cmd := Subscribe(r.topic, r.channel)
    err = conn.WriteCommand(cmd)
    if err != nil {
        cleanupConnection()
        return fmt.Errorf("[%s] failed to subscribe to %s:%s - %s",
            conn, r.topic, r.channel, err.Error())
    }

    r.mtx.Lock()
    delete(r.pendingConnections, addr)
    r.connections[addr] = conn
    r.mtx.Unlock()

    // pre-emptive signal to existing connections to lower their RDY count
    for _, c := range r.conns() {
        r.maybeUpdateRDY(c)
    }

    return nil
}

大部分代码都是一些参数和日志的处理，不用全部看，我们主要是追踪消息是如何被消费的，所以主要看我在上面用*注释出来的代码。

这里首先是建立了一个conn对象，然后启动这个对象的connect()方法发起网络IO，从而获取一个包含消息的response。

接下来，继续追踪connect()函数：

func (c *Conn) Connect() (*IdentifyResponse, error) {
    dialer := &net.Dialer{
        LocalAddr: c.config.LocalAddr,
        Timeout:   c.config.DialTimeout,
    }

    conn, err := dialer.Dial("tcp", c.addr)
    if err != nil {
        return nil, err
    }
    c.conn = conn.(*net.TCPConn)
    c.r = conn
    c.w = conn

    _, err = c.Write(MagicV2)
    if err != nil {
        c.Close()
        return nil, fmt.Errorf("[%s] failed to write magic - %s", c.addr, err)
    }

    resp, err := c.identify()
    if err != nil {
        return nil, err
    }

    if resp != nil && resp.AuthRequired {
        if c.config.AuthSecret == "" {
            c.log(LogLevelError, "Auth Required")
            return nil, errors.New("Auth Required")
        }
        err := c.auth(c.config.AuthSecret)
        if err != nil {
            c.log(LogLevelError, "Auth Failed %s", err)
            return nil, err
        }
    }

    c.wg.Add(2)
    atomic.StoreInt32(&c.readLoopRunning, 1)
    go c.readLoop()
    go c.writeLoop()
    return resp, nil
}

这段代码主要是开启一个TCP链接，并从NSQ服务端读取消息数据，一样，大部分代码只是一些message的参数处理，并不需要过多理解。这里我们来看一下go c.readLoop()这个函数的内部细节。

func (c *Conn) readLoop() {
    delegate := &connMessageDelegate{c}
    for {
        if atomic.LoadInt32(&c.closeFlag) == 1 {
            goto exit
        }

/*看这里*/frameType, data, err := ReadUnpackedResponse(c)
        if err != nil {
            if !strings.Contains(err.Error(), "use of closed network connection") {
                c.log(LogLevelError, "IO error - %s", err)
                c.delegate.OnIOError(c, err)
            }
            goto exit
        }

        if frameType == FrameTypeResponse && bytes.Equal(data, []byte("_heartbeat_")) {
            c.log(LogLevelDebug, "heartbeat received")
            c.delegate.OnHeartbeat(c)
            err := c.WriteCommand(Nop())
            if err != nil {
                c.log(LogLevelError, "IO error - %s", err)
                c.delegate.OnIOError(c, err)
                goto exit
            }
            continue
        }

        switch frameType {
        case FrameTypeResponse:
            c.delegate.OnResponse(c, data)
        case FrameTypeMessage:
        /*解码*/msg, err := DecodeMessage(data)
            if err != nil {
                c.log(LogLevelError, "IO error - %s", err)
                c.delegate.OnIOError(c, err)
                goto exit
            }
            msg.Delegate = delegate
            msg.NSQDAddress = c.String()

            atomic.AddInt64(&c.rdyCount, -1)
            atomic.AddInt64(&c.messagesInFlight, 1)
            atomic.StoreInt64(&c.lastMsgTimestamp, time.Now().UnixNano())

        /*看这里*/c.delegate.OnMessage(c, msg)
        case FrameTypeError:
            c.log(LogLevelError, "protocol error - %s", data)
            c.delegate.OnError(c, data)
        default:
            c.log(LogLevelError, "IO error - %s", err)
            c.delegate.OnIOError(c, fmt.Errorf("unknown frame type %d", frameType))
        }
    }

exit:
    atomic.StoreInt32(&c.readLoopRunning, 0)
    // start the connection close
    messagesInFlight := atomic.LoadInt64(&c.messagesInFlight)
    if messagesInFlight == 0 {
        // if we exited readLoop with no messages in flight
        // we need to explicitly trigger the close because
        // writeLoop won't
        c.close()
    } else {
        c.log(LogLevelWarning, "delaying close, %d outstanding messages", messagesInFlight)
    }
    c.wg.Done()
    c.log(LogLevelInfo, "readLoop exiting")
}

我在上述代码中标出需要细读的部分，首先是

frameType, data, err := ReadUnpackedResponse(c)

这一行代码是从NSQ服务端中读取message的真正实现的函数，ReadUnpackedResponse()这个函数的入参类型为io.writer，通过这个函数我们从connect中读取message，并以[]byte的形式返回（就是上述的data）
然后通过
msg, err := DecodeMessage(data)
将代码解码为*Message类型，最后通过
c.delegate.OnMessage(c, msg)
将代码写进consumer端中。

func (d *consumerConnDelegate) OnMessage(c *Conn, m *Message)         { d.r.onConnMessage(c, m) }

上面的onMessage()实际上是ConnDelegate 接口里的函数，我们可以去看看它的具体实现，如下：

func (r *Consumer) onConnMessage(c *Conn, msg *Message) {
    atomic.AddInt64(&r.totalRdyCount, -1)
    atomic.AddUint64(&r.messagesReceived, 1)
    r.incomingMessages <- msg//*****
    r.maybeUpdateRDY(c)
}

上述代码中的incomingMessages便是consumer端的一个全局变量，一个存储着*Message的channel。
如此，一个整体的consume流程就完成了。

写这篇博客其实并不是想写下对NSQ源码深度的解读，所以很多细节部分也没有跟大家讲解清楚。
主要是想记录自己看源码的思路和历程，主要有时候自己看的源码很多时候有些比较重要的流程过后可能就忘记了，所以在此记录一下思路，再次读的时候就能有一个很好的思路流程了。

很久没更新技术博客了，看来今后要勤快点/(ㄒoㄒ)/~~