io.CopyにはなるべくWriteTo関数を渡してあげた方が良さそう #Go

io.Readerのファイルタイプを判定するを拝見してちょっと気になったので調べました。

io.MultiReaderではWrite関数しか実装されておらず、 WriteTo(w Writer) (n int64, err error)関数が存在しません。
そのため、みんな大好きio.Copyで積極利用されるWriteTo(w Writer) (n int64, err error)関数が働かず少し残念な気持ちがありました。

もしかすると気持ちだけの問題かもしれないのでBenchmarkしてみました。

bufio.Readerには WriteToが実装されいたので、こちらを元のReaderに被せてる事にしました。そうすると都合よくまさにリード位置を進めず先頭から指定バイト数だけ読むbufio.Peek関数が用意されていたのでそちらを利用しました。
元のisGzip関数を書き直したのが以下になります。

var head = []byte{0x1F, 0x8B, 0x8}

func isGzip2(input io.Reader) (io.Reader, bool, error) {
    r := bufio.NewReader(input)
    buf, err := r.Peek(len(head))
    if err != nil {
        return nil, false, err
    }
    return r, bytes.Equal(buf, head), nil
}

コード全文:
https://github.com/masahide/golang-peek-test/blob/master/peek_test.go

そして、このコードをGithub ActionsでBenchmark実行した結果が以下です。

2020-01-09T13:28:03.6328460Z goos: linux
2020-01-09T13:28:03.6328821Z goarch: amd64
2020-01-09T13:28:04.8143891Z BenchmarkIsGzip-2           100      11298715 ns/op         248 B/op          5 allocs/op
2020-01-09T13:28:05.9750862Z BenchmarkIsGzip2-2       918237          1158 ns/op        4240 B/op          3 allocs/op
2020-01-09T13:28:05.9767115Z PASS
2020-01-09T13:28:05.9780877Z ok     _/home/runner/work/golang-peek-test/golang-peek-test    2.475s

テスト結果全文はこちら

入力がbytes.NewBufferで、出力先もioutil.Discardなのでおそらく極端な例にはなってると思うのでかなり差が出てしまってるのだと思うのですが、少なからずio.Copyを利用する場合はWriteToを意識しておく方が良さそうです。

ちなみに、標準ライブラリでWriteToの実装があるパッケージはここで一覧されるので眺めておくと良さそう。
(ReadFromもおそらくは同様だと思われるので一緒にこちらも)