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 协议》，转载必须注明作者和本文链接