Golang 實現Redis 協定解析器的解決方案

$3rnSETrn

$4arnb

*2
$3
foo
$3
bar

*3
$3
SET
$3
key
$5
value

*3rn$3rnSETrn$3rnkeyrn$5rnvaluern

*3  
$3
SET
$4
arnb 
$7
myvalue

*3  
$3
SET
$4
a  // 錯誤的分行
b // 錯誤的分行
$7
myvalue

msg = make([]byte, 4 + 2) // 正文長度4 + 換行符長度2
_, err = io.ReadFull(reader, msg)

// Payload stores redis.Reply or error
type Payload struct {
	Data redis.Reply
	Err  error
}

// ParseStream 通過 io.Reader 讀取資料並將結果通過 channel 將結果返回給呼叫者
// 流式處理的介面適合供使用者端/伺服器端使用
func ParseStream(reader io.Reader) <-chan *Payload {
	ch := make(chan *Payload)
	go parse0(reader, ch)
	return ch
}

// ParseOne 解析 []byte 並返回 redis.Reply 
func ParseOne(data []byte) (redis.Reply, error) {
	ch := make(chan *Payload)
	reader := bytes.NewReader(data)
	go parse0(reader, ch)
	payload := <-ch // parse0 will close the channel
	if payload == nil {
		return nil, errors.New("no reply")
	}
	return payload.Data, payload.Err
}

func parse0(reader io.Reader, ch chan<- *Payload) {
    // 初始化讀取狀態
    readingMultiLine := false
    expectedArgsCount := 0
    var args [][]byte
    var bulkLen int64
    for {
        // 上文中我們提到 RESP 是以行為單位的
        // 因為行分為簡單字串和二進位制安全的BulkString，我們需要封裝一個 readLine 函數來相容
        line, err = readLine(reader, bulkLen)
        if err != nil { 
            // 處理錯誤
            return
        }
        // 接下來我們對剛剛讀取的行進行解析
        // 我們簡單的將 Reply 分為兩類:
        // 單行: StatusReply, IntReply, ErrorReply
        // 多行: BulkReply, MultiBulkReply

        if !readingMultiLine {
            if isMulitBulkHeader(line) {
                // 我們收到了 MulitBulkReply 的第一行
                // 獲得 MulitBulkReply 中 BulkString 的個數
                expectedArgsCount = parseMulitBulkHeader(line)
                // 等待 MulitBulkReply 後續行
                readingMultiLine = true
            } else if isBulkHeader(line) {
                // 我們收到了 BulkReply 的第一行
                // 獲得 BulkReply 第二行的長度, 通過 bulkLen 告訴 readLine 函數下一行 BulkString 的長度
                bulkLen = parseBulkHeader()
                // 這個 Reply 中一共有 1 個 BulkString
                expectedArgsCount = 1 
                // 等待 BulkReply 後續行
                readingMultiLine = true
            } else {
                // 處理 StatusReply, IntReply, ErrorReply 等單行 Reply
                reply := parseSingleLineReply(line)
                // 通過 ch 返回結果
                emitReply(ch)
            }
        } else {
            // 進入此分支說明我們正在等待 MulitBulkReply 或 BulkReply 的後續行
            // MulitBulkReply 的後續行有兩種，BulkHeader 或者 BulkString
            if isBulkHeader(line) {
                bulkLen = parseBulkHeader()
            } else {
                // 我們正在讀取一個 BulkString, 它可能是 MulitBulkReply 或 BulkReply 
                args = append(args, line)
            }
            if len(args) == expectedArgsCount { // 我們已經讀取了所有後續行
                // 通過 ch 返回結果
                emitReply(ch)
                // 重置狀態, 準備解析下一條 Reply
                readingMultiLine = false
                expectedArgsCount = 0
                args = nil
                bulkLen = 0
            }
        }
    }
}

func readLine(bufReader *bufio.Reader, state *readState) ([]byte, bool, error) {
	var msg []byte
	var err error
	if state.bulkLen == 0 { // read simple line
		msg, err = bufReader.ReadBytes('n')
		if err != nil {
			return nil, true, err
		}
		if len(msg) == 0 || msg[len(msg)-2] != 'r' {
			return nil, false, errors.New("protocol error: " + string(msg))
		}
	} else { // read bulk line (binary safe)
		msg = make([]byte, state.bulkLen+2)
		_, err = io.ReadFull(bufReader, msg)
		if err != nil {
			return nil, true, err
		}
		if len(msg) == 0 ||
			msg[len(msg)-2] != 'r' ||
			msg[len(msg)-1] != 'n' {
			return nil, false, errors.New("protocol error: " + string(msg))
		}
		state.bulkLen = 0
	}
	return msg, false, nil
}

func parseMultiBulkHeader(msg []byte, state *readState) error {
	var err error
	var expectedLine uint64
	expectedLine, err = strconv.ParseUint(string(msg[1:len(msg)-2]), 10, 32)
	if err != nil {
		return errors.New("protocol error: " + string(msg))
	}
	if expectedLine == 0 {
		state.expectedArgsCount = 0
		return nil
	} else if expectedLine > 0 {
		// first line of multi bulk reply
		state.msgType = msg[0]
		state.readingMultiLine = true
		state.expectedArgsCount = int(expectedLine)
		state.args = make([][]byte, 0, expectedLine)
		return nil
	} else {
		return errors.New("protocol error: " + string(msg))
	}
}

func parseBulkHeader(msg []byte, state *readState) error {
	var err error
	state.bulkLen, err = strconv.ParseInt(string(msg[1:len(msg)-2]), 10, 64)
	if err != nil {
		return errors.New("protocol error: " + string(msg))
	}
	if state.bulkLen == -1 { // null bulk
		return nil
	} else if state.bulkLen > 0 {
		state.msgType = msg[0]
		state.readingMultiLine = true
		state.expectedArgsCount = 1
		state.args = make([][]byte, 0, 1)
		return nil
	} else {
		return errors.New("protocol error: " + string(msg))
	}
}

func parseSingleLineReply(msg []byte) (redis.Reply, error) {
	str := strings.TrimSuffix(string(msg), "n")
	str = strings.TrimSuffix(str, "r")
	var result redis.Reply
	switch msg[0] {
	case '+': // status reply
		result = reply.MakeStatusReply(str[1:])
	case '-': // err reply
		result = reply.MakeErrReply(str[1:])
	case ':': // int reply
		val, err := strconv.ParseInt(str[1:], 10, 64)
		if err != nil {
			return nil, errors.New("protocol error: " + string(msg))
		}
		result = reply.MakeIntReply(val)
	default:
		// parse as text protocol
		strs := strings.Split(str, " ")
		args := make([][]byte, len(strs))
		for i, s := range strs {
			args[i] = []byte(s)
		}
		result = reply.MakeMultiBulkReply(args)
	}
	return result, nil
}