作者 丨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相關賬號發布。