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

Posted on by hackdl

本篇我们完成消息队列的收尾工作，将我们的服务器搭建起来。我们的设想是消费者通过 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

每当监听到 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("%sn", 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)

再来就是启动一个 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)
        }
    }()

作为参数传入的 context 也是为了统一退出管理，当监听到退出信号（）后，我们再生成一个带超时时间的 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() {

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

第二步是监听中断信号，接收到信号后通过 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 %sn", cmd.name, string(bytes.Join(cmd.params, []byte(" "))))
        if err != nil {
            return err
        }
    } else {
        _, err := fmt.Fprintf(c.conn, "%sn", 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，输出如下：

服务已成功启动。

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

可以看出我们的消费者已成功连接到服务器并执行了订阅请求，相应的 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'

服务端输出日志：

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

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

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

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

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

服务器托管，北京服务器托管，服务器租用，机房机柜带宽租用

服务器托管

咨询：董先生

电话13051898268 QQ/微信93663045！

上一篇:
下一篇: