go语言中使用TeeReader计算文件读取速度示例代码

Go  /  管理员 发布于 2年前   734

TeeReader是io库中的一个函数，传入一个Reader 和一个 Writer ，返回一个teeReader 对象 ，

当你读取 teeReader 中的内容时，会无缓冲的将读取内容写入到 Writer 中。

TeeReader通常使用在数据流的处理中，比如计算下载速度，计算文件hash值等

今天我们就来测试一下计算文件下载速度功能

示例代码

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)
	}
}

有兴趣的可以自行测试