用 Go 写一个简单消息队列（二）：客户端处理和 channel 设计

上一篇文章我们定义了消息体和基础工具，这一篇我们开始着手客户端的处理函数和 channel 的基础设计。

客户端处理函数

这里所谓的客户端指的是消费者，处理函数也就是处理消费者同我们服务之间的 tcp 连接。我们定义一个结构体 Client，里面包含有连接和状态字段，然后就是编写读写状态和 tcp 连接的相关函数。

client.go

package server

import (
    "encoding/binary"
    "io"
    "log"
)

type Client struct {
    conn  io.ReadWriteCloser
    name  string
    state int
}

func NewClient(conn io.ReadWriteCloser, name string) *Client {
    return &Client{conn, name, -1}
}

func (c *Client) String() string {
    return c.name
}

func (c *Client) GetState() int {
    return c.state
}

func (c *Client) SetState(state int) {
    c.state = state
}

func (c *Client) Read(data []byte) (int, error) {
    return c.conn.Read(data)
}

func (c *Client) Write(data []byte) (int, error) {
    var err error

    err = binary.Write(c.conn, binary.BigEndian, int32(len(data)))
    if err != nil {
        return 0, err
    }

    n, err := c.conn.Write(data)
    if err != nil {
        return 0, err
    }

    return n + 4, nil
}

func (c *Client) Close() {
    log.Printf("CLIENT(%s): closing", c.String())
    c.conn.Close()
}

这里的逻辑比较简单，唯一值得一提的是 Write 方法。在给消费者写消息之前，我们先往连接中写入消息体的长度，固定为 4 个字节，这样客户端读取的时候就可以先读取长度，然后按长度读取消息。

channel

从上篇文章中我们可以知道，channel 是我们这个消息队列中的核心数据结构之一，因此它的设计尤为重要。

维护消费者信息

首先，因为我们的消费者是从 channel 中读取消息的，所以 channel 中需要维护消费者的信息，并且可以增删消费者。因此我们先在 channel 结构中维护一个 consumer 数组和两个管道用来接收增删 consumer 的消息：

type Consumer interface {
    Close()
}

type Channel struct {
    name                string
    addClientChan       chan util.ChanReq
    removeClientChan    chan util.ChanReq
    clients             []Consumer
}

func (c *Channel) AddClient(client Consumer) {
    log.Printf("Channel(%s): adding client...", c.name)
    doneChan := make(chan interface{})
    c.addClientChan <- util.ChanReq{
        Variable: client,
        RetChan:  doneChan,
    }
    <-doneChan
}

func (c *Channel) RemoveClient(client Consumer) {
    log.Printf("Channel(%s): removing client...", c.name)
    doneChan := make(chan interface{})
    c.removeClientChan <- util.ChanReq{
        Variable: client,
        RetChan:  doneChan,
    }
    <-doneChan
}

值得注意的是，这里我们没有直接绑定上面的 Client 结构体，而是抽象出了一个 Consumer 接口。这样做的好处是倒转依赖关系，而且可以避免包循环引用。

既然有了接收消息的管道，那么我们需要一个常驻后台的 goroutine 来处理这些消息，可以称之为事件处理循环，也就是一个 for + select 组合：

// Router handles the events of Channel
func (c *Channel) Router() {
    var clientReq util.ChanReq

    for {
        select {
        case clientReq = <-c.addClientChan:
            client := clientReq.Variable.(Consumer)
            c.clients = append(c.clients, client)
            log.Printf("CHANNEL(%s) added client %#v", c.name, client)
            clientReq.RetChan <- struct{}{}
        case clientReq = <-c.removeClientChan:
            client := clientReq.Variable.(Consumer)
            indexToRemove := -1
            for k, v := range c.clients {
                if v == client {
                    indexToRemove = k
                    break
                }
            }
            if indexToRemove == -1 {
                log.Printf("ERROR: could not find client(%#v) in clients(%#v)", client, c.clients)
            } else {
                c.clients = append(c.clients[:indexToRemove], c.clients[indexToRemove+1:]...)
                log.Printf("CHANNEL(%s) removed client %#v", c.name, client)
            }
            clientReq.RetChan <- struct{}{}
        }
    }
}

收发消息

对于收发消息，这里我们使用三个管道来实现：

• msgChan：这是一个有缓冲管道，用来暂存消息，超过长度则丢弃消息（后续会加上持久化到磁盘的功能）
• incomingMessageChan：用来接收生产者的消息
• clientMessageChan：消息会被发送到这个管道，后续会由消费者拉取

代码如下：

type Channel struct {
    ...
    incomingMessageChan chan *Message
    msgChan             chan *Message
    clientMessageChan   chan *Message
}

func (c *Channel) PutMessage(msg *Message) {
    c.incomingMessageChan <- msg
}

func (c *Channel) PullMessage() *Message {
    return <-c.clientMessageChan
}

func (c *Channel) Router() {
    var clientReq util.ChanReq

    go c.MessagePump()

    for {
        select {
        ...
        case msg := <-c.incomingMessageChan:
            // 防止因 msgChan 缓冲填满时造成阻塞，加上一个 default 分支直接丢弃消息
            select {
            case c.msgChan <- msg:
                log.Printf("CHANNEL(%s) wrote message", c.name)
            default:
            }
        }
    }
}

// MessagePump send messages to ClientMessageChan
func (c *Channel) MessagePump() {
    var msg *Message

    for {
        select {
        case msg = <-c.msgChan:
        }

        c.clientMessageChan <- msg
    }
}

关闭

当 channel 关闭的时候，我们需要做一些清理的工作，首先我们增加一个接收关闭信号的管道，在接收到信号时关闭发送消息的 MessagePump 协程和消费者连接，代码如下：

type Channel struct {
    ...
    exitChan            chan util.ChanReq
}

func (c *Channel) Router() {
    var (
        ...
        closeChan = make(chan struct{})
    )

    go c.MessagePump(closeChan)

    for {
        select {
        ...
        case closeReq := <-c.exitChan:
            log.Printf("CHANNEL(%s) is closing", c.name)
            close(closeChan)

            for _, consumer := range c.clients {
                consumer.Close()
            }

            closeReq.RetChan <- nil
        }
    }
}

// MessagePump send messages to ClientMessageChan
func (c *Channel) MessagePump(closeChan chan struct{}) {
    var msg *Message

    for {
        select {
        ...
        case <-closeChan:
            return
        }
        ...
    }
}

func (c *Channel) Close() error {
    errChan := make(chan interface{})
    c.exitChan <- util.ChanReq{
        RetChan: errChan,
    }

    err, _ := (<-errChan).(error)
    return err
}

我们在事件处理循环中初始化一个管道，并作为参数传递给 MessagePump 协程，当接收到关闭信号时关闭此管道，然后依次关闭消费者连接，关闭逻辑就结束了。

channel 的完整代码如下：

channel.go

---

项目地址：https://github.com/yhao1206/SMQ
相关阅读：

• 用 Go 写一个简单消息队列（一）：定义消息和基础工具
• 用 Go 写一个简单消息队列（三）：增加完成确认和重入队列功能
• 用 Go 写一个简单消息队列（四）：topic 设计
• 用 Go 写一个简单消息队列（五）：协议与后台队列实现
• 用 Go 写一个简单消息队列（六）：服务器实现