如何统计从io.Reader流式传输至cmd.Stdin的字节总量？

方法一：用io.TeeReader快速实现 #

io.TeeReader能在读取数据的同时，把内容复制到指定Writer，我们可以借助它做字节计数。

修改后的代码示例：

resp, err := httpClient.Get(url)
if err != nil {
    err = errors.Wrap(err, "HTTP request failed")
    return
}
defer resp.Body.Close() // 务必关闭响应体，避免资源泄漏

var totalBytes int64
// 自定义一个计数Writer，每次写入就累加字节数
countWriter := io.WriterFunc(func(p []byte) (n int, err error) {
    totalBytes += int64(len(p))
    return len(p), nil
})

// 用TeeReader包装原响应体，读取时同步计数
teeReader := io.TeeReader(resp.Body, countWriter)

args = append(args, "-")
cmd := exec.Command("exiftool", args...)

stdout, err := cmd.StdoutPipe()
if err != nil {
    return
}
cmd.Stdin = teeReader

err = cmd.Start()
if err != nil {
    return
}

// 等待命令执行结束，确保所有数据都传输完成
if err := cmd.Wait(); err != nil {
    // 按需处理命令执行错误
    return
}

// 此时totalBytes就是实际传输的总字节数
fmt.Printf("传输总字节数: %d\n", totalBytes)

方法二：自定义计数Reader #

如果需要更灵活的控制，可以自己实现io.Reader接口，在Read方法里统计字节：

// 自定义计数Reader
type CountReader struct {
    Source    io.Reader
    TotalRead int64
}

func (cr *CountReader) Read(p []byte) (n int, err error) {
    n, err = cr.Source.Read(p)
    cr.TotalRead += int64(n)
    return
}

// 使用示例
resp, err := httpClient.Get(url)
if err != nil {
    err = errors.Wrap(err, "HTTP request failed")
    return
}
defer resp.Body.Close()

countReader := &CountReader{Source: resp.Body}

args = append(args, "-")
cmd := exec.Command("exiftool", args...)

stdout, err := cmd.StdoutPipe()
if err != nil {
    return
}
cmd.Stdin = countReader

err = cmd.Start()
if err != nil {
    return
}

if err := cmd.Wait(); err != nil {
    // 处理命令执行错误
    return
}

fmt.Printf("传输总字节数: %d\n", countReader.TotalRead)

关键注意点 #

• 必须调用cmd.Wait()，等命令执行完再读取计数，不然会统计到不完整的字节数。
• 不管哪种方法，都要记得关闭resp.Body，防止资源泄漏。
• 如果命令中途异常退出，计数器只会统计到实际传输的字节数，而非响应体的总大小。

内容的提问来源于stack exchange，提问作者frederix