Go言語(golang) io.TeeReader を使ってダウンロードの進捗状況を出力しながらファイルに書き込む

Go言語で大きなファイル(ISOなど)をダウンロードしながら標準出力にプログレスバーなど、
進捗状況を出力したい場合に、io.TeeReader が便利です。
また、Content-Length をダウンロードが完了する前に知ることができます。

ファイルに書き込みながら進捗状況を表示

package main

import (
    "fmt"
    "io"
    "net/http"
    "os"
)

type Progress struct {
    total int64
    size  int64
}

// io.Writer の要件を満たす
func (p *Progress) Write(data []byte) (int, error) {
    n := len(data)
    p.size += int64(n)

    fmt.Fprintf(os.Stdout, "%d/%d\n", p.size, p.total)

    return n, nil
}

func main() {
    fp, err := os.OpenFile("filename", os.O_CREATE|os.O_WRONLY, 0664)
    if err != nil {
        fmt.Fprintf(os.Stderr, "error: %s\n", err)
        os.Exit(1)
    }
    defer fp.Close()

    var p Progress

    // linux の ISO ファイルなどの大きなファイルのURL
    resp, err := http.Get("http://example.com")
    if err != nil {
        fmt.Fprintf(os.Stderr, "error: %s\n", err)
        os.Exit(1)
    }
    defer resp.Body.Close()

    // resp.Body の合計
    p.total = resp.ContentLength

    // ファイルに書き込みながら p.Write 内も実行される
    io.Copy(fp, io.TeeReader(resp.Body, &p))
}

レスポンスボディが大きい場合は、中断する

package main

import (
    "bytes"
    "fmt"
    "io"
    "io/ioutil"
    "net/http"
    "os"
)

func main() {
    resp, err := http.Get("http://example.com")
    if err != nil {
        fmt.Fprintf(os.Stderr, "error: %s\n", err)
        os.Exit(1)
    }
    defer resp.Body.Close()

    // resp.Body の合計
    fmt.Println("Content-Length:", resp.ContentLength)
    if resp.ContentLength > 10000 {
        // context を使って大きなファイルの場合は、
        // ダウンロードを中断する処理などができる
    }

    // resp.Body に戻す
    var b bytes.Buffer
    resp.Body = ioutil.NopCloser(io.TeeReader(resp.Body, &b))
    body, _ := ioutil.ReadAll(resp.Body)
    fmt.Println(string(body))

    // 戻す必要がない場合
    var b bytes.Buffer
    body, _ := ioutil.ReadAll(io.TeeReader(resp.Body, &b))
    fmt.Println(string(body))
}