Golang IO包的妙用

背景

以一個(gè)RPC的協(xié)議包來(lái)說(shuō)，每個(gè)包有如下結(jié)構(gòu)

type Packet struct {
    TotalSize uint32    
    Magic     [4]byte    
    Payload   []byte    
    Checksum  uint32
}

其中TotalSize是整個(gè)包除去TotalSize后的字節(jié)數(shù)， Magic是一個(gè)固定長(zhǎng)度的字串，Payload是包的實(shí)際內(nèi)容，包含業(yè)務(wù)邏輯的數(shù)據(jù)。


Checksum是對(duì)Magic和Payload的adler32校驗(yàn)和。

編碼(encode)

我們使用一個(gè)原型為func EncodePacket(w io.Writer, payload []byte) error的函數(shù)來(lái)把數(shù)據(jù)打包，結(jié)合encoding/binary （https://godoc.org/encoding/binary）我們很容易寫出第一版，演示需要，錯(cuò)誤處理方面就簡(jiǎn)化處理了。

var RPC_MAGIC = [4]byte{'p', 'y', 'x', 'i'}

func EncodePacket(w io.Writer, payload []byte) error {
    // len(Magic) + len(Checksum) == 8
    totalsize := uint32(len(payload) + 8)    
    // write total size
    binary.Write(w, binary.BigEndian, totalsize)    
    
    // write magic bytes
    binary.Write(w, binary.BigEndian, RPC_MAGIC)    
    
    // write payload
    w.Write(payload)    
    
    // calculate checksum
    var buf bytes.Buffer
    buf.Write(RPC_MAGIC[:])
    buf.Write(payload)
    checksum := adler32.Checksum(buf.Bytes())    
    
    // write checksum
    return binary.Write(w, binary.BigEndian, checksum)
}

在上面的實(shí)現(xiàn)中，為了計(jì)算 checksum，我們使用了一個(gè)內(nèi)存 buffer 來(lái)緩存數(shù)據(jù)，最后把所有的數(shù)據(jù)一次性讀出來(lái)算 checksum，考慮到計(jì)算 checksum 是一個(gè)不斷 update 地過(guò)程，我們應(yīng)該有方法直接略過(guò)內(nèi)存 buffer 而計(jì)算 checksum。

查看 hash/adler32  (http://godoc.org/hash/adler32#New) 我們得知，我們可以構(gòu)造一個(gè) Hash42 的對(duì)象，這個(gè)對(duì)象內(nèi)嵌了一個(gè) Hash 的接口，這個(gè)接口的定義如下：

type Hash interface {
    // Write (via the embedded io.Writer interface) adds more data to the running hash.
    // It never returns an error.
    io.Writer    
    
    // Sum appends the current hash to b and returns the resulting slice.
    // It does not change the underlying hash state.
    Sum(b []byte) []byte

    // Reset resets the Hash to its initial state.
    Reset()    
    
    // Size returns the number of bytes Sum will return.
    Size() int

    // BlockSize returns the hash's underlying block size.
    // The Write method must be able to accept any amount
    // of data, but it may operate more efficiently if all writes
    // are a multiple of the block size.
    BlockSize() int
}

這是一個(gè)通用的計(jì)算hash的接口，標(biāo)準(zhǔn)庫(kù)里面所有計(jì)算hash的對(duì)象都實(shí)現(xiàn)了這個(gè)接口，比如 md5, crc32等。由于Hash實(shí)現(xiàn)了io.Writer接口，因此我們可以把所有要計(jì)算的數(shù)據(jù)像寫入文件一樣寫入到這個(gè)對(duì)象中，最后調(diào)用Sum(nil)就可以得到最終的hash的byte數(shù)組。利用這個(gè)思路，第二版可以這樣寫:

func EncodePacket2(w io.Writer, payload []byte) error {
    // len(Magic) + len(Checksum) == 8
    totalsize := uint32(len(RPC_MAGIC) + len(payload) + 4)    
    // write total size
    binary.Write(w, binary.BigEndian, totalsize)    
    
    // write magic bytes
    binary.Write(w, binary.BigEndian, RPC_MAGIC)    
    
    // write payload
    w.Write(payload)   
    
    // calculate checksum
    sum := adler32.New()
    sum.Write(RPC_MAGIC[:])
    sum.Write(payload)
    checksum := sum.Sum32()    
    
    // write checksum
    return binary.Write(w, binary.BigEndian, checksum)
}

注意這次的變化，前面寫入TotalSize，Magic，Payload部分沒(méi)有變化，在計(jì)算checksum的時(shí)候去掉了bytes.Buffer，減少了一次內(nèi)存申請(qǐng)和拷貝。

考慮到sum和w都是io.Writer，利用神奇的 io.MultiWriter  （https://godoc.org/io#MultiWriter），我們可以這樣寫：

func EncodePacket(w io.Writer, payload []byte) error {
    // len(Magic) + len(Checksum) == 8
    totalsize := uint32(len(RPC_MAGIC) + len(payload) + 4)    
    // write total size
    binary.Write(w, binary.BigEndian, totalsize)

    sum := adler32.New()
    ww := io.MultiWriter(sum, w)    
    // write magic bytes
    binary.Write(ww, binary.BigEndian, RPC_MAGIC)    
   
    // write payload
    ww.Write(payload)    
   
   // calculate checksum
    checksum := sum.Sum32()   
   
   // write checksum
    return binary.Write(w, binary.BigEndian, checksum)
}

注意MultiWriter的使用，我們把w和sum利用MultiWriter綁在了一起創(chuàng)建了一個(gè)新的Writer，向這個(gè)Writer里面寫入數(shù)據(jù)就同時(shí)向w和sum里面都寫入數(shù)據(jù)，這樣就完成了發(fā)送數(shù)據(jù)和計(jì)算checksum的同步進(jìn)行，而對(duì)于binary.Write來(lái)說(shuō)沒(méi)有任何區(qū)別，因?yàn)樗枰氖且粋€(gè)實(shí)現(xiàn)了Write方法的對(duì)象。

解碼(decode)

基于上面的思想，解碼也可以把接收數(shù)據(jù)和計(jì)算checksum一起進(jìn)行，完整代碼如下

func DecodePacket(r io.Reader) ([]byte, error) {
    var totalsize uint32
    err := binary.Read(r, binary.BigEndian, &totalsize)    
    if err != nil {    
        return nil, errors.Annotate(err, "read total size")
    }    
        
    // at least len(magic) + len(checksum)
    if totalsize < 8 {    
        return nil, errors.Errorf("bad packet. header:%d", totalsize)
    }

    sum := adler32.New()
    rr := io.TeeReader(r, sum)    
    
    var magic [4]byte
    err = binary.Read(rr, binary.BigEndian, &magic)    
    if err != nil {    
        return nil, errors.Annotate(err, "read magic")
    }    
    if magic != RPC_MAGIC {    
        return nil, errors.Errorf("bad rpc magic:%v", magic)
    }

    payload := make([]byte, totalsize-8)
    _, err = io.ReadFull(rr, payload)    
    if err != nil {    
        return nil, errors.Annotate(err, "read payload")
    }    
    
    var checksum uint32
    err = binary.Read(r, binary.BigEndian, &checksum)    
    if err != nil {    
        return nil, errors.Annotate(err, "read checksum")
    }    
        
    if checksum != sum.Sum32() {    
        return nil, errors.Errorf("checkSum error, %d(calc) %d(remote)", sum.Sum32(), checksum)
    }    
    return payload, nil
}

上面代碼中，我們使用了 io.TeeReader  （http://godoc.org/io#TeeReader），這個(gè)函數(shù)的原型為func TeeReader(r Reader, w Writer) Reader，它返回一個(gè)Reader，這個(gè)Reader是參數(shù)r的代理，讀取的數(shù)據(jù)還是來(lái)自r，不過(guò)同時(shí)把讀取的數(shù)據(jù)寫入到w里面。

一切皆文件

Unix 下有一切皆文件的思想，Golang 把這個(gè)思想貫徹到更遠(yuǎn)，因?yàn)楸举|(zhì)上我們對(duì)文件的抽象就是一個(gè)可讀可寫的一個(gè)對(duì)象，也就是實(shí)現(xiàn)了io.Writer和io.Reader的對(duì)象我們都可以稱為文件，在上面的例子中無(wú)論是EncodePacket還是DecodePacket我們都沒(méi)有假定編碼后的數(shù)據(jù)是發(fā)送到 socket，還是從內(nèi)存讀取數(shù)據(jù)解碼，因此我們可以這樣調(diào)用 EncodePacket :

conn, _ := net.Dial("tcp", "127.0.0.1:8000")
EncodePacket(conn, []byte("hello"))

把數(shù)據(jù)直接發(fā)送到 socket，也可以這樣

conn, _ := net.Dial("tcp", "127.0.0.1:8000")
bufconn := bufio.NewWriter(conn)
EncodePacket(bufconn, []byte("hello"))

對(duì)socket加上一個(gè)buffer來(lái)增加吞吐量，也可以這樣

conn, _ := net.Dial("tcp", "127.0.0.1:8000")
zip := zlib.NewWriter(conn)
bufconn := bufio.NewWriter(conn)
EncodePacket(bufconn, []byte("hello"))

加上一個(gè)zip壓縮，還可以利用加上 crypto/aes 來(lái)個(gè)AES加密...

在這個(gè)時(shí)候，文件已經(jīng)不再局限于io，可以是一個(gè)內(nèi)存 buffer，也可以是一個(gè)計(jì)算hash的對(duì)象，甚至是一個(gè)計(jì)數(shù)器，流量限速器。Golang 靈活的接口機(jī)制為我們提供了無(wú)限可能。

END

我一直認(rèn)為一個(gè)好的語(yǔ)言一定有一個(gè)設(shè)計(jì)良好的標(biāo)準(zhǔn)庫(kù)，Golang的標(biāo)準(zhǔn)庫(kù)是作者們多年系統(tǒng)編程的沉淀，值得我們細(xì)細(xì)品味。