Go io.TeeReader 怎样统计流经读取器的字节数

在上传、解压或解析大文件时，常见需求不是把整个输入读进内存，而是想知道“已经流经读取链多少字节”。Go 的 io.TeeReader 很适合做这件事：把原始 io.Reader 包装起来，每次读取到的数据同步写入一个统计用 io.Writer，主流程仍然从包装后的 Reader 读取。

官方文档：https://pkg.go.dev/io#TeeReader

核心规则是“读多少，旁路 Writer 就收到多少”。TeeReader 没有内部缓冲，旁路写入完成前本次读取不会完成；如果旁路写入报错，这个错误会通过 Read 返回。

先把读取链变成可计数的 Reader

io.TeeReader(r, w) 返回的仍然是一个 io.Reader。它不会主动启动协程，也不会预读；只有调用包装器的 Read 时，才先从 r 取数据，再把本次取到的有效字节交给 w。因此可以直接把它传给原来接收 Reader 的函数。

package main

import (
	"fmt"
	"io"
	"strings"
)

func main() {
	input := strings.NewReader("alpha-beta-gamma")
	var count countingWriter

	// 包装后仍实现 io.Reader，原有的 io.Copy 等调用可以继续复用。
	tee := io.TeeReader(input, &count)
	if _, err := io.Copy(io.Discard, tee); err != nil {
		// 旁路 Writer 的错误会从 tee 的 Read 链路返回。
		panic(err)
	}
	fmt.Println(count.n)
}

type countingWriter struct {
	n int64
}

func (w *countingWriter) Write(p []byte) (int, error) {
	// 这里只累计本次收到的长度，不保留数据副本。
	w.n += int64(len(p))
	return len(p), nil
}

这个计数器的优点是内存开销与输入大小无关。它只保存一个整数，读完后即可得到旁路收到的总字节数。若主流程提前停止，计数值也只代表“已经通过 TeeReader 的字节”，不是原始资源的声明大小。

计数 Writer 之外，还能接入什么旁路

统计只是最小用法。只要对象实现 io.Writer，就能把同一批有效字节交给它。例如 hash.Hash 可以持续计算摘要，bytes.Buffer 可以保留小输入副本，自定义 Writer 则可以把累计值投递给进度回调。选择旁路对象时，要先明确是否需要保存全部内容。

package stream

import (
	"crypto/sha256"
	"encoding/hex"
	"io"
)

func digestAndCount(src io.Reader) (int64, string, error) {
	hashWriter := sha256.New()
	count := &countingWriter{}

	// 先统计，再把同一批读取结果交给摘要 Writer；不需要保存完整正文。
	tee := io.TeeReader(src, io.MultiWriter(count, hashWriter))
	if _, err := io.Copy(io.Discard, tee); err != nil {
		// 读取错误或任一旁路写入错误都必须交给调用方处理。
		return count.n, "", err
	}
	return count.n, hex.EncodeToString(hashWriter.Sum(nil)), nil
}

type countingWriter struct{ n int64 }

func (w *countingWriter) Write(p []byte) (int, error) {
	// 返回完整长度，才能让 MultiWriter 把这次写入视为成功。
	w.n += int64(len(p))
	return len(p), nil
}

这里的 io.MultiWriter 只是把一次旁路写入分发给两个 Writer；文章的核心仍是 TeeReader 负责把“读取动作”连接到“旁路写入”。如果摘要 Writer 或计数 Writer 返回错误，调用方应停止继续消费并保留已返回的错误。

最容易误判的三个边界

• 没有内部缓冲：TeeReader 不会把数据暂存后异步写出。旁路 Writer 较慢时，主读取也会等待，不能把它当作后台进度通道。
• 按有效字节计数：Writer 收到的是本次 Read 返回的 p[:n]，不是缓冲区容量。读取返回 n > 0 且同时带错误时，先处理这 n 个字节，再处理错误。
• 写入错误会中断读取语义：旁路写失败后，TeeReader 从 Read 返回写入错误。计数值可能已经增加，不能仅凭计数值就判断整个输入成功完成。

如果只想显示进度，可以让 Writer 只更新原子计数或发送受控事件；如果需要摘要，不要另起一条读取原始 Reader 的链路，否则两条读取会争抢同一数据源。对网络请求体、文件上传这类一次性流，TeeReader 的“边读边旁路”通常比先 ReadAll 更稳妥。

常见问题

Q：TeeReader 会把数据缓存起来吗？
不会。它没有内部缓冲，旁路 Writer 的写入必须在本次 Read 返回前完成。

Q：为什么统计值小于文件大小？
它统计的是已经被调用方读取并通过旁路的字节。如果上层提前停止、读取报错或只读取了限制范围，统计值自然可能小于资源总大小。

Q：旁路 Writer 返回短写怎么办？
Writer 应在返回少于输入长度时同时返回非空错误；不要把短写当作成功，并让错误回到调用方。