Golang：高效地通过Channel传输压缩字节流

package main

import (
    "bytes"
    "compress/zlib"
    "fmt"
    "io"
)

// BytesWithError 结构体用于通过通道传输字节切片及相关错误
type BytesWithError struct {
    Data []byte
    Err  error
}

// ChanWriter 是一个自定义的 io.Writer，它将写入的数据发送到其内部的通道
type ChanWriter chan BytesWithError

// Write 方法实现了 io.Writer 接口。
// 当 zlib.Writer 调用此方法时，它会将压缩后的数据块传递给 p。
// ChanWriter 负责将这些数据块通过其内部通道发送出去。
func (cw ChanWriter) Write(p []byte) (n int, err error) {
    // 重要的并发安全考虑：
    // p 是一个切片，其底层数组可能在调用者侧被重用或修改。
    // 为了确保发送到通道的数据是独立的且不会被后续操作污染，
    // 我们需要创建一个 p 的副本。
    buf := make([]byte, len(p))
    copy(buf, p)

    // 将复制的字节切片发送到通道
    // 如果通道已关闭或接收方不再接收，此操作可能导致 panic。
    // 在本例中，我们将确保通道在发送前是开放的。
    cw <- BytesWithError{Data: buf, Err: nil}
    return len(p), nil
}

// Compress 函数接收一个 io.Reader 作为输入，并返回一个接收压缩字节流的通道。
// 压缩操作在一个独立的 Goroutine 中执行。
func Compress(r io.Reader) <-chan BytesWithError {
    // 创建一个 ChanWriter 实例，它本身就是一个通道。
    // 这个通道将用于发送压缩后的字节切片。
    cw := make(ChanWriter)

    // 启动一个 Goroutine 来执行压缩逻辑，实现异步处理。
    go func() {
        // defer 确保在 Goroutine 退出时关闭通道。
        // 关闭通道是通知消费者不再有数据的重要信号。
        defer close(cw)

        // zlib.NewWriter 接收一个 io.Writer 接口作为其输出目标。
        // 这里我们将自定义的 ChanWriter 传递给它，
        // 这样 zlib.Writer 就会将压缩后的数据写入到 ChanWriter，
        // 进而通过 cw 的通道发送出去。
        zlibWriter := zlib.NewWriter(cw)
        // defer 确保 zlib writer 在 Goroutine 退出时被关闭。
        // 关闭 zlib writer 会刷新所有内部缓冲区中剩余的压缩数据。
        defer zlibWriter.Close()

        // 定义一个缓冲区用于从输入 io.Reader 中读取数据。
        // 使用较大的缓冲区可以提高文件读取和压缩的效率。
        rBuff := make([]byte, 4096) // 示例缓冲区大小，可根据实际情况调整

        for {
            // 从输入 io.Reader 中读取数据。
            n, err := r.Read(rBuff)
            if n > 0 {
                // 如果成功读取到数据，将其写入 zlibWriter。
                // zlibWriter 会将压缩后的数据通过 cw (ChanWriter) 发送。
                if _, writeErr := zlibWriter.Write(rBuff[:n]); writeErr != nil {
                    // 如果写入 zlibWriter 失败，通过通道发送错误并退出 Goroutine。
                    cw <- BytesWithError{Data: nil, Err: writeErr}
                    return
                }
            }

            // 检查读取操作的错误。
            if err != nil {
                if err != io.EOF {
                    // 如果是除 io.EOF 之外的其他错误，通过通道发送错误并退出。
                    cw <- BytesWithError{Data: nil, Err: err}
                }
                break // 读取结束 (无论是 EOF 还是其他错误)
            }
        }
    }()

    // 立即返回 ChanWriter 的通道。
    // 消费者可以立即开始从这个通道接收数据，而压缩操作在后台 Goroutine 中进行。
    return cw
}

func main() {
    // 示例：一个字符串作为输入源数据
    originalData := "This is some sample data that we want to compress and send through a channel. It should be long enough to demonstrate compression and channel usage in Go. We'll ensure the data integrity by decompressing it afterwards."
    reader := bytes.NewBufferString(originalData)

    fmt.Println("--- 压缩过程开始 ---")

    // 调用 Compress 函数，获取一个接收压缩数据的通道
    compressedChan := Compress(reader)

    var receivedCompressedData bytes.Buffer
    // 从通道中读取压缩数据块
    for chunk := range compressedChan {
        if chunk.Err != nil {
            fmt.Printf("压缩过程中发生错误: %v\n", chunk.Err)
            return
        }
        if chunk.Data != nil {
            // 将接收到的压缩数据块写入缓冲区
            receivedCompressedData.Write(chunk.Data)
        }
    }

    fmt.Println("--- 压缩过程结束 ---")
    fmt.Printf("原始数据长度: %d 字节\n", len(originalData))
    fmt.Printf("压缩后数据长度: %d 字节\n", receivedCompressedData.Len())

    // 验证：解压缩数据并与原始数据对比
    // 创建一个新的 zlib.Reader 来解压缩接收到的数据
    zlibReader, err := zlib.NewReader(&receivedCompressedData)
    if err != nil {
        fmt.Printf("创建 zlib 解压器失败: %v\n", err)
        return
    }
    defer zlibReader.Close() // 确保解压器被关闭

    // 读取所有解压缩后的数据
    decompressedData, err := io.ReadAll(zlibReader)
    if err != nil {
        fmt.Printf("读取解压缩数据失败: %v\n", err)
        return
    }

    fmt.Printf("解压缩后数据长度: %d 字节\n", len(decompressedData))
    // 比较解压缩后的数据与原始数据是否一致
    fmt.Printf("解压缩数据与原始数据匹配: %t\n", string(decompressedData) == originalData)
    // fmt.Println("解压缩数据:", string(decompressedData)) // 可选：打印解压缩数据
}