Runner in the High

技術のことをかくこころみ

ioutil.Discardとio.CopyNでアロケーションコストを抑えてデータサイズを判定する

ioutil パッケージに Discard という /dev/null 的な io.Writer が用意されている。

最近これを使うタイミングがあったのでメモ。以下のようなコードがあるとする。

package main

import (
    "fmt"
    "io"
    "io/ioutil"
    "strings"
)

func main() {
    a := strings.NewReader("123456789")

    _, err := io.CopyN(ioutil.Discard, a, 10)
    if err == io.EOF {
        fmt.Println("more than 9")
    } else if err != nil {
        panic(err)
    }
}

上記のコードを実行すると more than 9 が出力になる。

ある io.Reader に対して io.CopyN でサイズ指定をしてデータコピーを試み、それがEOFかどうかを見ることで対象のデータのサイズが任意のサイズを超えているかどうかを疑似的に判定している、という感じ。データサイズの判定がしたいだけでコピー先のデータは捨ててしまってokなので ioutil.Discard を使っている。

このコードではイメージ付きづらいが、仮にsrc変数に300MiBとかのデータが入っているとすると、その具体的なサイズを取得するのに []byte などへ変換して len に通せばとんでもないアロケーションコストが発生することになる。一方で io.Reader のまま引き回せば []byte へ変換するアロケーションコストを抑えながらサイズ判定もできる。