用 Go 写一个简单消息队列（六）：服务器实现

头像
与昊
     江苏
    阅读 10 分钟
    1

    本篇我们完成消息队列的收尾工作，将我们的服务器搭建起来。我们的设想是消费者通过 tcp 连接到服务端，生产者则是通过 http 协议直接发送消息。

    监听消费者连接

    在 server 目录下新建一个 tcp.go 文件，用于监听消费者的 tcp 连接，代码如下：

    package server

import (
    "context"
    "log"
    "net"
)

func TcpServer(ctx context.Context, addr, port string) {
    fqAddress := addr + ":" + port
    listener, err := net.Listen("tcp", fqAddress)
    if err != nil {
        panic("tcp listen(" + fqAddress + ") failed")
    }
    log.Printf("listening for clients on %s", fqAddress)

    for {
        select {
        case <-ctx.Done():
            return
        default:
            conn, err := listener.Accept()
            if err != nil {
                panic("accept failed: " + err.Error())
            }
            client := NewClient(conn, conn.RemoteAddr().String())
            go client.Handle(ctx)
        }
    }
}

    每当监听到 tcp 连接，就新建一个 client 来处理，同时传入一个 context 方便统一进行退出管理。client 和 protocol 也都要同时加上对这个 context 的监听代码，具体可参考代码仓库。

    HTTP Server

    同样在 server 目录下新建一个 http.go 文件来编写 HTTP 服务器，我们的 HTTP 服务器只提供三个方法：测试连接、写入消息和查看所有 topic，先来看一下这三个方法的代码：

    http.go

package server

import (
    "context"
    "errors"
    "fmt"
    "github.com/yhao1206/SMQ/message"
    "github.com/yhao1206/SMQ/protocol"
    "io"
    "log"
    "net/http"
    "net/url"
    "strconv"
    "time"
)

type ReqParams struct {
    params url.Values
    body   []byte
}

func NewReqParams(req *http.Request) (*ReqParams, error) {
    reqParams, err := url.ParseQuery(req.URL.RawQuery)
    if err != nil {
        return nil, err
    }

    data, err := io.ReadAll(req.Body)
    if err != nil {
        return nil, err
    }

    return &ReqParams{reqParams, data}, nil
}

func (r *ReqParams) Query(key string) (string, error) {
    keyData := r.params[key]
    if len(keyData) == 0 {
        return "", errors.New("key not in query params")
    }
    return keyData[0], nil
}

func pingHandler(w http.ResponseWriter, req *http.Request) {
    w.Header().Set("Content-Length", "2")
    io.WriteString(w, "OK")
}

func putHandler(w http.ResponseWriter, req *http.Request) {
    reqParams, err := NewReqParams(req)
    if err != nil {
        log.Printf("HTTP: error - %s", err.Error())
        return
    }

    topicName, err := reqParams.Query("topic")
    if err != nil {
        log.Printf("HTTP: error - %s", err.Error())
        return
    }

    conn := &FakeConn{}
    client := NewClient(conn, "HTTP")
    proto := &protocol.Protocol{}
    resp, err := proto.Execute(client, "PUB", topicName, string(reqParams.body))
    if err != nil {
        log.Printf("HTTP: error - %s", err.Error())
        return
    }

    w.Header().Set("Content-Length", strconv.Itoa(len(resp)))
    w.Write(resp)
}

func statsHandler(w http.ResponseWriter, req *http.Request) {
    for topicName, _ := range message.TopicMap {
        io.WriteString(w, fmt.Sprintf("%s\n", topicName))
    }
}

    为了统一管理，我们的 putHandler 方法没有直接操作 topic，而是包装了一个假的客户端向协议中发送 PUB 指令，由协议与 topic 进行交互。fake_conn 和 PUB 方法的代码如下：

    server/fake_conn.go

package server

import (
    "io"
)

type FakeConn struct {
    io.ReadWriter
}

func (c *FakeConn) Close() error {
    return nil
}

----------

protocol/protocol.go

func (p *Protocol) PUB(client StatefulReadWriter, params []string) ([]byte, error) {
    var buf bytes.Buffer
    var err error

    //  fake clients don't get to ClientInit
    if client.GetState() != -1 {
        return nil, ClientErrInvalid
    }

    if len(params) < 3 {
        return nil, ClientErrInvalid
    }

    topicName := params[1]
    body := []byte(params[2])

    _, err = buf.Write(<-util.UuidChan)
    if err != nil {
        return nil, err
    }

    _, err = buf.Write(body)
    if err != nil {
        return nil, err
    }

    topic := message.GetTopic(topicName)
    topic.PutMessage(message.NewMessage(buf.Bytes()))

    return []byte("OK"), nil
}

    再来就是启动一个 http server：

    http.go

func HttpServer(ctx context.Context, address string, port string, endChan chan struct{}) {
    http.HandleFunc("/ping", pingHandler)
    http.HandleFunc("/put", putHandler)
    http.HandleFunc("/stats", statsHandler)

    fqAddress := address + ":" + port
    httpServer := http.Server{
        Addr: fqAddress,
    }

    go func() {
        log.Printf("listening for http requests on %s", fqAddress)
        err := http.ListenAndServe(fqAddress, nil)
        if err != nil {
            log.Fatal("http.ListenAndServe:", err)
        }
    }()

    <-ctx.Done()
    log.Printf("HTTP server on %s is shutdowning...", fqAddress)
    timeoutCtx, fn := context.WithTimeout(context.Background(), 10*time.Second)
    defer fn()
    if err := httpServer.Shutdown(timeoutCtx); err != nil {
        log.Printf("HTTP server shutdown error: %v", err)
    }
    close(endChan)
}

    作为参数传入的 context 也是为了统一退出管理，当监听到退出信号（<-ctx.Done()）后，我们再生成一个带超时时间的 context，让 http server 有足够的时间完成清理工作，实现优雅退出。退出之后关闭 endChan，调用方可以根据从 endChan 接收到数据判断优雅退出已完成。

    main 函数

    终于到了最后一步，就是实现我们的 main 函数，将我们的服务都启动起来。在项目根目录下创建 smq.go 文件，写入以下代码：

    package main

import (
    "context"
    "flag"
    "github.com/yhao1206/SMQ/message"
    "github.com/yhao1206/SMQ/server"
    "github.com/yhao1206/SMQ/util"
    "os"
    "os/signal"
    "strconv"
)

var bindAddress = flag.String("address", "", "address to bind to")
var webPort = flag.Int("web-port", 5150, "port to listen on for HTTP connections")
var tcpPort = flag.Int("tcp-port", 5151, "port to listen on for TCP connections")
var memQueueSize = flag.Int("mem-queue-size", 10000, "number of messages to keep in memory (per topic)")

func main() {
    flag.Parse()

    endChan := make(chan struct{})
    signalChan := make(chan os.Signal, 1)
    signal.Notify(signalChan, os.Interrupt)

    ctx, fn := context.WithCancel(context.Background())
    go func() {
        <-signalChan
        fn()
    }()

    go message.TopicFactory(ctx, *memQueueSize)
    go util.UuidFactory(ctx)
    go server.TcpServer(ctx, *bindAddress, strconv.Itoa(*tcpPort))
    server.HttpServer(ctx, *bindAddress, strconv.Itoa(*webPort), endChan)

    for _, topic := range message.TopicMap {
        topic.Close()
    }

    <-endChan
}

    第一步依照惯例都是解析命令行参数，让用户可以自主控制监听端口和内存队列的大小。

    第二步是监听中断信号，接收到信号后通过 context 的特性向所有的后台服务发送退出信号，然后关闭所有的 topic，等待 http server 退出完成即可。简单两步，我们的消息队列系统就搭建起来了。

    测试

    我们现在来测试一下我们的消息队列，为了方便测试，这边提供了一个简单的消费者连接库：

    client/conn.go

package client

import (
    "bytes"
    "encoding/binary"
    "fmt"
    "github.com/yhao1206/SMQ/message"
    "io"
    "net"
    "strconv"
)

type Client struct {
    conn io.ReadWriteCloser
}

type Command struct {
    name   []byte
    params [][]byte
}

type Response struct {
    FrameType int32
    Data      interface{}
}

func NewClient(conn net.Conn) *Client {
    return &Client{conn}
}

func (c *Client) Connect(address string, port int) error {
    fqAddress := address + ":" + strconv.Itoa(port)
    conn, err := net.Dial("tcp", fqAddress)
    if err != nil {
        return err
    }
    c.conn = conn
    return nil
}

func (c *Client) Version(version string) error {
    _, err := c.conn.Write([]byte(version))
    return err
}

func (c *Client) Subscribe(topic string, channel string) *Command {
    params := make([][]byte, 2)
    params[0] = []byte(topic)
    params[1] = []byte(channel)
    return &Command{[]byte("SUB"), params}
}

func (c *Client) Ready(count int) *Command {
    params := make([][]byte, 1)
    params[0] = []byte(strconv.Itoa(count))
    return &Command{[]byte("RDY"), params}
}

func (c *Client) Finish(uuid string) *Command {
    params := make([][]byte, 1)
    params[0] = []byte(uuid)
    return &Command{[]byte("FIN"), params}
}

func (c *Client) Requeue(uuid string) *Command {
    params := make([][]byte, 1)
    params[0] = []byte(uuid)
    return &Command{[]byte("REQ"), params}
}

func (c *Client) Get() *Command {
    return &Command{[]byte("GET"), nil}
}

func (c *Client) WriteCommand(cmd *Command) error {
    if len(cmd.params) > 0 {
        _, err := fmt.Fprintf(c.conn, "%s %s\n", cmd.name, string(bytes.Join(cmd.params, []byte(" "))))
        if err != nil {
            return err
        }
    } else {
        _, err := fmt.Fprintf(c.conn, "%s\n", cmd.name)
        if err != nil {
            return err
        }
    }
    return nil
}

func (c *Client) ReadResponse() (*message.Message, error) {
    err := c.WriteCommand(c.Get())
    if err != nil {
        return nil, err
    }

    var msgSize int32
    err = binary.Read(c.conn, binary.BigEndian, &msgSize)
    if err != nil {
        return nil, err
    }

    buf := make([]byte, msgSize)
    _, err = c.conn.Read(buf)
    if err != nil {
        return nil, err
    }

    msg := message.NewMessage(buf)

    return msg, nil
}

    新建 examples/cousumer/consumer.go 文件，写下我们的消费者测试代码：

    examples/cousumer/consumer.go

package main

import (
    "github.com/yhao1206/SMQ/client"
    "github.com/yhao1206/SMQ/util"
    "log"
)

func main() {
    consumeClient := client.NewClient(nil)
    err := consumeClient.Connect("127.0.0.1", 5151)
    if err != nil {
        log.Fatal(err)
    }
    consumeClient.WriteCommand(consumeClient.Subscribe("test", "ch"))

    for {
        msg, err := consumeClient.ReadResponse()
        if err != nil {
            log.Fatal(err)
        }
        log.Printf("%s - %s", util.UuidToStr(msg.Uuid()), msg.Body())
        consumeClient.WriteCommand(consumeClient.Finish(util.UuidToStr(msg.Uuid())))
    }
}

    启动服务端：go run smq.go，输出如下：

    image.png

    服务已成功启动。

    再启动我们的测试消费者：go run ./examples/consumer/consumer.go，此时服务端多了几行输出：

    image.png

    可以看出我们的消费者已成功连接到服务器并执行了订阅请求，相应的 topic 也已经创建完毕。

    通过 curl 客户端（shell 或者 postman 等）发送 curl 命令来发布一条消息：

    curl --location --request POST '127.0.0.1:5150/put?topic=test' --header 'Content-Type: text/plain' --data-raw 'hello'

    服务端输出日志：

    image.png

    查看一下消费者客户端的输出：

    image.png

    可以看到我们的消费者成功拉取到了生产者的消息。

    到这里我们整个系列就全部结束了，欢迎大家提出宝贵的意见。

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

    go程序员消息队列后端
    与昊
    225 声望636 粉丝

    IT民工，主要从事web方向，喜欢研究技术和投资之道

    « 上一篇
    用 Go 写一个简单消息队列（五）：协议与后台队列实现

    引用和评论

    被 5 篇内容引用
    推荐阅读
    头像
    前端如何入门 Go 语言

    robin赞 28阅读 8.3k评论 9

    头像
    疫情特殊时期，致 SegmentFault 思否全体员工的一封公开信

    波波Nadia赞 14阅读 21.4k评论 10

    头像
    在Java程序中监听mysql的binlog

    huan1993赞 13阅读 5.5k

    头像
    喜大普奔，适用于 VS Code 的 GitHub Copilot 全新免费版本正式推出，GitHub 全球开发者突破1.5亿

    鸣飞赞 8阅读 12.5k评论 1

    头像
    Spring 实现 3 种异步流式接口，干掉接口超时烦恼

    程序员小富赞 6阅读 2.4k

    头像
    估算工作时间的 3 条规则（在软件开发项目中）

    大发明家赞 4阅读 840

    头像
    Go 语言-计算密集型服务 性能优化

    Jesse赞 4阅读 8.7k

    0 条评论
    得票最新