Carpe Diem

備忘録

Goでのstreamの扱い方を学ぶ

概要

結論から言うと、Streamで扱っているものはStreamのまま扱うです。

具体的にはio.Readerを毎回ioutil.ReadAll[]byteに変換せずにそのまま使いましょうです。

なぜStreamを使うべきか

Node.jsの例ですが、こちらで非常に分かりやすく説明されています。

yosuke-furukawa.hatenablog.com

それを踏まえて考えてみると、Goの場合以下の2つが大きいと思います。

1. メモリの効率化

ioutil.ReadAllなどで一旦全て[]byteに変換すると、その分メモリを消費しますし、アロケーションGCに依る速度低下が起きます。

一方io.Readerio.Writerは各chunkの処理に同じバイトを使いまわすので、メモリの効率が良いです。

2. 標準パッケージの多くがio.Readerをサポートしてる

io.Readerio.Writerがメソッドの少ない非常に良いインターフェースであるため、多くの標準パッケージがインタフェースを実装していたり、引数として扱える形でサポートしています。

推測するな、計測せよ

Goの精神で推測するな、計測せよという言葉があります。
頭でこねくり回したり言葉で説明するよりも、実際に計測した方が早いですし確実です。
以下に例を示しました。

S3などからデータをとり、それを渡すProxyを考えた場合を想定しています。S3からio.Readerでデータを受け取った後、それをResponseWriterにWriteしています。

サンプルコード

func buf1(body io.Reader, w http.ResponseWriter) {
    b, err := ioutil.ReadAll(body)
    if err != nil {
        panic(err)
    }
    w.Write(b)
}

func buf2(body io.Reader, w http.ResponseWriter) {
    _, err := io.Copy(w, body)
    if err != nil {
        panic(err)
    }
}

ベンチマークコード

var image []byte

func init() {
    body, _ := os.Open("./image.jpg")
    image, _ = ioutil.ReadAll(body)
    body.Close()
}

func BenchmarkImage1(b *testing.B) {
    for i := 0; i < b.N; i++ {
        w := httptest.NewRecorder()
        r := bytes.NewReader(image)
        buf1(r, w)
    }
}

func BenchmarkImage2(b *testing.B) {
    for i := 0; i < b.N; i++ {
        w := httptest.NewRecorder()
        r := bytes.NewReader(image)
        buf2(r, w)
    }
}

結果

26KB画像
BenchmarkImage1-8         100000         18950 ns/op       92960 B/op          16 allocs/op
BenchmarkImage2-8         200000          6397 ns/op       29472 B/op          10 allocs/op
940KB画像
BenchmarkImage1-8           3000        549892 ns/op     3062820 B/op          21 allocs/op
BenchmarkImage2-8          10000        172470 ns/op      967713 B/op          10 allocs/op
3MB画像
BenchmarkImage1-8           1000       1906131 ns/op    11844643 B/op          23 allocs/op
BenchmarkImage2-8           3000        568504 ns/op     3458080 B/op          10 allocs/op

このようにio.Copyを使ってStreamのまま扱う場合はallocsが変わらない&処理が早いのに対し、一旦ioutil.ReadAllで変換する場合はメモリの使用量やallocsがどんどん増えていきます。

やはりStreamはStreamのまま扱った方が良さそうです。

具体的な実装例

以下具体的な実装例を書いていきます。

json.Unmarshalでなくjson.NewDecoderを使おう

こちらはよく知られている例で、HTTPのリクエストをサーバ内でJSONに変換する際はUnmarshalでなくNewDecoderを使う、というやり方です。

go - Decoding JSON using json.Unmarshal vs json.NewDecoder.Decode - Stack Overflow

NG

http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    var u User
    data, _ := ioutil.ReadAll(req.Body)
    err = json.Unmarshal(data, &u)
    ...
})

OK

http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    var u User
    err = json.NewDecoder(req.Body).Decode(&u)
    ...
})

io.Copyを使おう

io.Readerからio.Writerにデータを渡したい時はio.Copyを使います。先ほどと同じですが、外部APIなどから受け取ったレスポンスを特に加工せずにProxy的に返したい場合は以下のようにします。

NG

func buf1(body io.Reader, w http.ResponseWriter) {
    b, err := ioutil.ReadAll(body)
    if err != nil {
        panic(err)
    }
    w.Write(b)
}

OK

func buf2(body io.Reader, w http.ResponseWriter) {
    _, err := io.Copy(w, body)
    if err != nil {
        panic(err)
    }
}

ただio.Copyは自動で32KBの内部バッファを持つので、小さいデータをやり取りする場合は無駄な確保になります。その場合はio.CopyBufferで再利用するバッファの容量を調整しましょう。

bytes.Bufferよりio.Pipeを使おう

bytes.Bufferはio.Readerio.Writerの両方共を実装しているので、I/O系の処理をする時に非常に便利です。

ただし「このstructをio.Readerとして渡したいなぁ」と思った時に、bytes.Bufferエンコードすると一時的にそのデータを保持することになります。
小さいデータならまだ良いですが、大きいデータを扱う場合はやはり無駄にメモリを消費してしまいます。
そこでio.Pipeを使うと特に内部バッファを保つ必要もなく、io.Readerとして渡すことが可能になります。

NG

func pipe1(v User) {
    var buf bytes.Buffer

    err := json.NewEncoder(&buf).Encode(&v)

    resp, err := http.Post("example.com", "application/json", &buf)
}

OK

func pipe2(v User) {
    pr, pw := io.Pipe()

    go func() {
        err := json.NewEncoder(pw).Encode(&v)
        pw.Close()
    }()

    resp, err := http.Post("example.com", "application/json", pr)
}

ここでgoroutineを使っているのは、io.Pipe()で生成されたWriterは内部バッファを持たないので同じ関数内で使うとデッドロックが発生するからです。
また使い終わったあとは必ずCloseしましょう。

ただしこれについては可読性が落ちることもあるので、ケースバイケースで利用すると良いです。

io.Readerのまま使おう

ioutil.ReadAlljson.Unmarshalを使う癖が付いていると以下のように書いてしまいがちですが、

func LoadGzippedJSON(r io.Reader, v interface{}) error {
    data, err := ioutil.ReadAll(r)
    if err != nil {
        return err
    }
    // oh wait, we need a Reader again.. 
    raw := bytes.NewBuffer(data)
    unz, err := gzip.NewReader(raw)
    if err != nil {
        return err
    }
    buf, err := ioutil.ReadAll(unz)
    if err != nil {
        return err
    }
    return json.Unmarshal(buf, &v)
}

となりますが、io.Readerのまま扱えば

func LoadGzippedJSON(r io.Reader, v interface{}) error {
    raw, err := gzip.NewReader(r)
    if err != nil {
        return err
    }
    return json.NewDecoder(raw).Decode(&v)
}

のように非常にシンプルに書けます。

まとめ

io.Reader, io.Writerを引数としているものはStreamのままで扱う。
ioutil.ReadAllなどを使ってわざわざ[]byteには変換しない。

ソース