Golang IO包的妙用

背景

以一个RPC的协议包来说，每个包有如下结构

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

其中TotalSize是整个包除去TotalSize后的字节数， Magic是一个固定长度的字串，Payload是包的实际内容，包含业务逻辑的数据。


Checksum是对Magic和Payload的adler32校验和。

编码(encode)

我们使用一个原型为func EncodePacket(w io.Writer, payload []byte) error的函数来把数据打包，结合encoding/binary （https://godoc.org/encoding/binary）我们很容易写出第一版，演示需要，错误处理方面就简化处理了。

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)
}

在上面的实现中，为了计算 checksum，我们使用了一个内存 buffer 来缓存数据，最后把所有的数据一次性读出来算 checksum，考虑到计算 checksum 是一个不断 update 地过程，我们应该有方法直接略过内存 buffer 而计算 checksum。

查看 hash/adler32  (http://godoc.org/hash/adler32#New) 我们得知，我们可以构造一个 Hash42 的对象，这个对象内嵌了一个 Hash 的接口，这个接口的定义如下：

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
}

这是一个通用的计算hash的接口，标准库里面所有计算hash的对象都实现了这个接口，比如 md5, crc32等。由于Hash实现了io.Writer接口，因此我们可以把所有要计算的数据像写入文件一样写入到这个对象中，最后调用Sum(nil)就可以得到最终的hash的byte数组。利用这个思路，第二版可以这样写:

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部分没有变化，在计算checksum的时候去掉了bytes.Buffer，减少了一次内存申请和拷贝。

考虑到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绑在了一起创建了一个新的Writer，向这个Writer里面写入数据就同时向w和sum里面都写入数据，这样就完成了发送数据和计算checksum的同步进行，而对于binary.Write来说没有任何区别，因为它需要的是一个实现了Write方法的对象。

解码(decode)

基于上面的思想，解码也可以把接收数据和计算checksum一起进行，完整代码如下

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），这个函数的原型为func TeeReader(r Reader, w Writer) Reader，它返回一个Reader，这个Reader是参数r的代理，读取的数据还是来自r，不过同时把读取的数据写入到w里面。

一切皆文件

Unix 下有一切皆文件的思想，Golang 把这个思想贯彻到更远，因为本质上我们对文件的抽象就是一个可读可写的一个对象，也就是实现了io.Writer和io.Reader的对象我们都可以称为文件，在上面的例子中无论是EncodePacket还是DecodePacket我们都没有假定编码后的数据是发送到 socket，还是从内存读取数据解码，因此我们可以这样调用 EncodePacket :

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

把数据直接发送到 socket，也可以这样

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

对socket加上一个buffer来增加吞吐量，也可以这样

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

加上一个zip压缩，还可以利用加上 crypto/aes 来个AES加密…

在这个时候，文件已经不再局限于io，可以是一个内存 buffer，也可以是一个计算hash的对象，甚至是一个计数器，流量限速器。Golang 灵活的接口机制为我们提供了无限可能。

END

我一直认为一个好的语言一定有一个设计良好的标准库，Golang的标准库是作者们多年系统编程的沉淀，值得我们细细品味。