[Go语言] Golang 实现 Redis系列(六)如何实现 pipeline 模式的 redis 客户端

type Client struct {
  conn    net.Conn // 与服务端的 tcp 连接
  pendingReqs chan *Request // 等待发送的请求
  waitingReqs chan *Request // 等待服务器响应的请求
  ticker    *time.Ticker // 用于触发心跳包的计时器
  addr    string
  ctx    context.Context
  cancelFunc context.CancelFunc
  writing  *sync.WaitGroup // 有请求正在处理不能立即停止，用于实现 graceful shutdown
}
type Request struct {
  id    uint64 // 请求id
  args    [][]byte // 上行参数
  reply   redis.Reply // 收到的返回值
  heartbeat bool // 标记是否是心跳请求
  waiting   *wait.Wait // 调用协程发送请求后通过 waitgroup 等待请求异步处理完成
  err     error
}

func (client *Client) Send(args [][]byte) redis.Reply {
request := &request{
args:    args,
heartbeat: false,
waiting:   &wait.Wait{},
}
request.waiting.Add(1)
client.working.Add(1)
defer client.working.Done()
client.pendingReqs <- request // 请求入队
timeout := request.waiting.WaitWithTimeout(maxWait) // 等待响应或者超时
if timeout {
return reply.MakeErrReply("server time out")
}
if request.err != nil {
return reply.MakeErrReply("request failed")
}
return request.reply
}

// 写协程入口
func (client *Client) handleWrite() {
for req := range client.pendingReqs {
client.doRequest(req)
}
}
// 发送请求
func (client *Client) doRequest(req *request) {
if req == nil || len(req.args) == 0 {
return
}
  // 序列化请求
re := reply.MakeMultiBulkReply(req.args)
bytes := re.ToBytes()
_, err := client.conn.Write(bytes)
i := 0
  // 失败重试
for err != nil && i < 3 {
err = client.handleConnectionError(err)
if err == nil {
_, err = client.conn.Write(bytes)
}
i++
}
if err == nil {
    // 发送成功等待服务器响应
client.waitingReqs <- req
} else {
req.err = err
req.waiting.Done()
}
}

// 收到服务端的响应
func (client *Client) finishRequest(reply redis.Reply) {
defer func() {
if err := recover(); err != nil {
debug.PrintStack()
logger.Error(err)
}
}()
request := <-client.waitingReqs
if request == nil {
return
}
request.reply = reply
if request.waiting != nil {
request.waiting.Done()
}
}
// 读协程是个 RESP 协议解析器
func (client *Client) handleRead() error {
ch := parser.ParseStream(client.conn)
for payload := range ch {
if payload.Err != nil {
client.finishRequest(reply.MakeErrReply(payload.Err.Error()))
continue
}
client.finishRequest(payload.Data)
}
return nil
}

func MakeClient(addr string) (*Client, error) {
  conn, err := net.Dial("tcp", addr)
  if err != nil {
    return nil, err
  }
  ctx, cancel := context.WithCancel(context.Background())
  return &Client{
    addr:    addr,
    conn:    conn,
    sendingReqs: make(chan *Request, chanSize),
    waitingReqs: make(chan *Request, chanSize),
    ctx:     ctx,
    cancelFunc:  cancel,
    writing:   &sync.WaitGroup{},
  }, nil
}
func (client *Client) Start() {
  client.ticker = time.NewTicker(10 * time.Second)
  go client.handleWrite()
  go func() {
    err := client.handleRead()
    logger.Warn(err)
  }()
  go client.heartbeat()
}

func (client *Client) Close() {
  // 先阻止新请求进入队列
  close(client.sendingReqs)
  // 等待处理中的请求完成
  client.writing.Wait()
  // 释放资源
  _ = client.conn.Close() // 关闭与服务端的连接，连接关闭后读协程会退出
  client.cancelFunc() // 使用 context 关闭读协程
  close(client.waitingReqs) // 关闭队列
}

func TestClient(t *testing.T) {
  client, err := MakeClient("localhost:6379")
  if err != nil {
    t.Error(err)
  }
  client.Start()
  result = client.Send([][]byte{
    []byte("SET"),
    []byte("a"),
    []byte("a"),
  })
  if statusRet, ok := result.(*reply.StatusReply); ok {
    if statusRet.Status != "OK" {
      t.Error("`set` failed, result: " + statusRet.Status)
    }
  }
  result = client.Send([][]byte{
    []byte("GET"),
    []byte("a"),
  })
  if bulkRet, ok := result.(*reply.BulkReply); ok {
    if string(bulkRet.Arg) != "a" {
      t.Error("`get` failed, result: " + string(bulkRet.Arg))
    }
  }
}