Golang中实现文件下载实时进度监控：自定义io.Reader的实践

package main

import (
    "bytes"
    "fmt"
    "io"
    "net/http"
    "os"
    "strings"
    "time" // 用于控制打印频率
)

// ProgressReader 包装一个 io.Reader，用于跟踪读取的字节数。
type ProgressReader struct {
    io.Reader
    totalRead int64 // 已读取的总字节数
    name      string // 阅读器名称，用于打印区分
    totalSize int64 // 总文件大小 (可选，用于计算百分比)
    lastPrint time.Time // 上次打印时间，用于控制打印频率
}

// NewProgressReader 创建一个新的 ProgressReader 实例。
func NewProgressReader(reader io.Reader, name string, totalSize int64) *ProgressReader {
    return &ProgressReader{
        Reader: reader,
        name: name,
        totalSize: totalSize,
        lastPrint: time.Now(),
    }
}

// Read 方法 '重写' 了底层 io.Reader 的 Read 方法。
// io.Copy 会调用此方法，我们在这里跟踪字节计数并转发调用。
func (pr *ProgressReader) Read(p []byte) (int, error) {
    n, err := pr.Reader.Read(p)
    pr.totalRead += int64(n)

    // 控制打印频率，避免频繁输出导致性能问题或日志过多
    // 或者在文件末尾 (io.EOF) 时强制打印最后一次进度
    if time.Since(pr.lastPrint) > 100*time.Millisecond || err == io.EOF {
        if pr.totalSize > 0 {
            progress := float64(pr.totalRead) / float64(pr.totalSize) * 100
            fmt.Printf("\r[%s] 已下载: %d / %d 字节 (%.2f%%)", pr.name, pr.totalRead, pr.totalSize, progress)
        } else {
            fmt.Printf("\r[%s] 已下载: %d 字节", pr.name, pr.totalRead)
        }
        pr.lastPrint = time.Now()
    }

    // 下载完成时换行，使最终的完成信息显示在新行
    if err == io.EOF {
        fmt.Println()
    }

    return n, err
}

func main() {
    // 示例1: 使用内存数据模拟进度跟踪
    fmt.Println("--- 示例1: 内存数据进度跟踪 ---")
    var src io.Reader
    var dst bytes.Buffer

    // 创建一些随机输入数据作为源
    testData := strings.Repeat("Some random input data for testing progress. ", 100)
    src = bytes.NewBufferString(testData)
    dataSize := int64(len(testData))

    // 使用我们的自定义 ProgressReader 包装源 Reader
    progressSrc := NewProgressReader(src, "内存源", dataSize)

    count, err := io.Copy(&dst, progressSrc)
    if err != nil {
        fmt.Println("复制内存数据时出错:", err)
        os.Exit(1)
    }
    fmt.Printf("内存数据传输完成，总计 %d 字节\n\n", count)

    // 示例2: 应用于实际文件下载
    fmt.Println("--- 示例2: 实际文件下载进度跟踪 ---")
    // 替换为实际可用的下载链接，例如一个测试文件。
    // 注意：某些链接可能需要处理重定向或SSL证书问题。
    downloadURL := "https://speed.hetzner.de/100MB.bin" // 一个100MB的测试文件
    outputFileName := "downloaded_file.bin"

    fmt.Printf("开始下载文件: %s\n", downloadURL)

    // 1. 发起HTTP GET请求
    resp, err := http.Get(downloadURL)
    if err != nil {
        fmt.Println("发起HTTP请求失败:", err)
        return
    }
    defer resp.Body.Close() // 确保关闭响应体

    if resp.StatusCode != http.StatusOK {
        fmt.Printf("下载失败，HTTP状态码: %d\n", resp.StatusCode)
        return
    }

    // 2. 获取文件总大小 (如果可用)
    var totalSize int64
    if resp.ContentLength > 0 {
        totalSize = resp.ContentLength
        fmt.Printf("文件总大小: %d 字节\n", totalSize)
    } else {
        fmt.Println("无法获取文件总大小，将显示已下载字节数。")
    }

    // 3. 创建输出文件
    outFile, err := os.Create(outputFileName)
    if err != nil {
        fmt.Println("创建输出文件失败:", err)
        return
    }
    defer outFile.Close() // 确保关闭文件

    // 4. 使用 ProgressReader 包装响应体
    progressReader := NewProgressReader(resp.Body, "下载器", totalSize)

    // 5. 将数据从 ProgressReader 复制到输出文件
    downloadedBytes, err := io.Copy(outFile, progressReader)
    if err != nil {
        fmt.Println("下载文件时出错:", err)
        return
    }

    fmt.Printf("文件下载完成，总计 %d 字节保存到 %s\n", downloadedBytes, outputFileName)
}