Golang IO包的妙用

作者 丨icexin

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我们很容易写出第一版,演示需要,错误处理方面就简化处理了:

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 我们得知,我们可以构造一个 Hash32 的对象,这个对象内嵌了一个 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,我们可以这样写:

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,这个函数的原型为 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 灵活的接口机制为我们提供了无限可能。

结尾

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

技术交流群:426582602,本文仅授权51 Reboot相关账号发布。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 199,636评论 5 468
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 83,890评论 2 376
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 146,680评论 0 330
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 53,766评论 1 271
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 62,665评论 5 359
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,045评论 1 276
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,515评论 3 390
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,182评论 0 254
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,334评论 1 294
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,274评论 2 317
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,319评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,002评论 3 315
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,599评论 3 303
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,675评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 30,917评论 1 255
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 42,309评论 2 345
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 41,885评论 2 341

推荐阅读更多精彩内容