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閱讀 228,119評論 6 531
  • 序言:濱河連續發生了三起死亡事件,死亡現場離奇詭異,居然都是意外死亡,警方通過查閱死者的電腦和手機,發現死者居然都...
    沈念sama閱讀 98,382評論 3 415
  • 文/潘曉璐 我一進店門,熙熙樓的掌柜王于貴愁眉苦臉地迎上來,“玉大人,你說我怎么就攤上這事。” “怎么了?”我有些...
    開封第一講書人閱讀 176,038評論 0 373
  • 文/不壞的土叔 我叫張陵,是天一觀的道長。 經常有香客問我,道長,這世上最難降的妖魔是什么? 我笑而不...
    開封第一講書人閱讀 62,853評論 1 309
  • 正文 為了忘掉前任,我火速辦了婚禮,結果婚禮上,老公的妹妹穿的比我還像新娘。我一直安慰自己,他們只是感情好,可當我...
    茶點故事閱讀 71,616評論 6 408
  • 文/花漫 我一把揭開白布。 她就那樣靜靜地躺著,像睡著了一般。 火紅的嫁衣襯著肌膚如雪。 梳的紋絲不亂的頭發上,一...
    開封第一講書人閱讀 55,112評論 1 323
  • 那天,我揣著相機與錄音,去河邊找鬼。 笑死,一個胖子當著我的面吹牛,可吹牛的內容都是我干的。 我是一名探鬼主播,決...
    沈念sama閱讀 43,192評論 3 441
  • 文/蒼蘭香墨 我猛地睜開眼,長吁一口氣:“原來是場噩夢啊……” “哼!你這毒婦竟也來了?” 一聲冷哼從身側響起,我...
    開封第一講書人閱讀 42,355評論 0 288
  • 序言:老撾萬榮一對情侶失蹤,失蹤者是張志新(化名)和其女友劉穎,沒想到半個月后,有當地人在樹林里發現了一具尸體,經...
    沈念sama閱讀 48,869評論 1 334
  • 正文 獨居荒郊野嶺守林人離奇死亡,尸身上長有42處帶血的膿包…… 初始之章·張勛 以下內容為張勛視角 年9月15日...
    茶點故事閱讀 40,727評論 3 354
  • 正文 我和宋清朗相戀三年,在試婚紗的時候發現自己被綠了。 大學時的朋友給我發了我未婚夫和他白月光在一起吃飯的照片。...
    茶點故事閱讀 42,928評論 1 369
  • 序言:一個原本活蹦亂跳的男人離奇死亡,死狀恐怖,靈堂內的尸體忽然破棺而出,到底是詐尸還是另有隱情,我是刑警寧澤,帶...
    沈念sama閱讀 38,467評論 5 358
  • 正文 年R本政府宣布,位于F島的核電站,受9級特大地震影響,放射性物質發生泄漏。R本人自食惡果不足惜,卻給世界環境...
    茶點故事閱讀 44,165評論 3 347
  • 文/蒙蒙 一、第九天 我趴在偏房一處隱蔽的房頂上張望。 院中可真熱鬧,春花似錦、人聲如沸。這莊子的主人今日做“春日...
    開封第一講書人閱讀 34,570評論 0 26
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽。三九已至,卻和暖如春,著一層夾襖步出監牢的瞬間,已是汗流浹背。 一陣腳步聲響...
    開封第一講書人閱讀 35,813評論 1 282
  • 我被黑心中介騙來泰國打工, 沒想到剛下飛機就差點兒被人妖公主榨干…… 1. 我叫王不留,地道東北人。 一個月前我還...
    沈念sama閱讀 51,585評論 3 390
  • 正文 我出身青樓,卻偏偏與公主長得像,于是被迫代替她去往敵國和親。 傳聞我的和親對象是個殘疾皇子,可洞房花燭夜當晚...
    茶點故事閱讀 47,892評論 2 372

推薦閱讀更多精彩內容