TeeReader使用筆記

TeeReader 是 io庫中的一個函式，傳入一個Reader 和一個 Writer ，返回一個teeReader 物件 ，當你讀取 teeReader 中的內容時，會無緩衝的將讀取內容寫入到 Writer 中。

// TeeReader returns a Reader that writes to w what it reads from r.
// All reads from r performed through it are matched with
// corresponding writes to w. There is no internal buffering -
// the write must complete before the read completes.
// Any error encountered while writing is reported as a read error.
func TeeReader(r Reader, w Writer) Reader {
    return &teeReader{r, w}
}

TeeReader 的使用

TeeReader 通常使用在資料流的處理中，比如計算下載速度，計算檔案hash值等。

簡單示例

讀取 Reader 中的內容，並同步寫入到 Writer 中。

func main() {
    // 建立一個reader
    r := strings.NewReader("Hello World!")
    var buf bytes.Buffer
    // 建立一個teeReader
    reader := io.TeeReader(r, &buf)

    // 讀取TeeReader中的內容 會同步寫入到buf中
    fmt.Println(io.ReadAll(reader))

    // 讀取buf中的內容
    fmt.Println(io.ReadAll(&buf))
}

注意：這裡必須先讀取teeReader 中的內容，才會將資料寫入 buf 中，如果先讀取 buf 將讀取不到資料。

使用 TeeReader 計算讀取速度

示例：

package main

import (
    "fmt"
    "io"
    "log"
    "net/http"
    "os"
    "sync/atomic"
    "time"
)

func main() {
    out, err := os.Create("test.tmp")
    if err != nil {
        log.Fatal(err)
    }
    defer out.Close()
    resp, err := http.Get("https://dl.google.com/go/go1.17.1.src.tar.gz")

    if err != nil {
        log.Fatal(err)
    }

    defer resp.Body.Close()
    teeReader := &Speeder{}
  // 列印讀取速度
    go teeReader.Show()
    io.Copy(out, io.TeeReader(resp.Body, teeReader))
}

// Speeder 用於記錄時間段內讀取的位元組數
type Speeder struct {
    count int64
}
// Write 實現Writer介面，記錄讀取的位元組數
func (s *Speeder) Write(b []byte) (int, error) {
    c := len(b)
    atomic.AddInt64(&s.count, int64(c))
    return c, nil
}

// Show 列印讀取速度
func (s *Speeder) Show() {
    ticker := time.NewTicker(time.Second)
    defer ticker.Stop()
    for range ticker.C {
        fmt.Printf("\r%.2fkb/s", float64(atomic.LoadInt64(&s.count))/float64(1024))
        atomic.StoreInt64(&s.count, 0)
    }
}

使用 TeeReader 計算檔案hash值

package main

import (
    "crypto/sha256"
    "fmt"
    "io"
    "log"
    "net/http"
    "os"
)

func main() {
    out, err := os.Create("test.tmp")
    if err != nil {
        log.Fatal(err)
    }
    defer out.Close()
    resp, err := http.Get("https://golang.google.cn/dl/go1.18.2.src.tar.gz")

    if err != nil {
        log.Fatal(err)
    }

    defer resp.Body.Close()
    h := sha256.New()
    tee := io.TeeReader(resp.Body, h)
    io.Copy(out, tee)
    fmt.Printf("%x", h.Sum(nil))
}

輸出結果：

2c44d03ea2c34092137ab919ba602f2c261a038d08eb468528a3f3a28e5667e2

總結

1. 為什麼使用TeeReader 來計算檔案hash值，直接讀取檔案資料，然後計算不是也可以嗎？

   對比了一下透過TeeReader 和 io.ReadAll() 讀取檔案後再 sha256 佔用的記憶體來看， TeeReader 佔用的記憶體更少。

本文首發於我的部落格 喜四點

本作品採用《CC 協議》，轉載必須註明作者和本文連結