Go语言中通过通道高效传输压缩字节流

package main

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

// BytesWithError 结构体用于在通道中传递字节数据块及可能发生的错误。
type BytesWithError struct {
    Data []byte
    Err  error
}

// ChanWriter 是一个实现了 io.Writer 接口的通道类型。
// 任何写入到 ChanWriter 的数据都会被封装成 BytesWithError 并发送到其内部通道。
type ChanWriter chan BytesWithError

// Write 方法实现了 io.Writer 接口。
// 当 zlib.Writer 调用此方法时，它会将压缩后的数据 p 写入到 ChanWriter。
// 为了避免并发问题（如果 p 的底层数组被 zlib.Writer 重用），
// 我们会创建一个 p 的副本并发送到通道。
func (cw ChanWriter) Write(p []byte) (n int, err error) {
    // 创建 p 的副本以防止数据竞争，因为 p 的底层数组可能被 zlib.Writer 重用。
    buf := make([]byte, len(p))
    copy(buf, p)

    // 将数据块发送到通道。
    // 注意：这里我们假设写入操作本身不会立即产生错误，
    // 真正的写入错误（如通道关闭）将在发送时由 Go runtime 处理，
    // 或者通过更复杂的 select 逻辑来捕获。
    cw <- BytesWithError{Data: buf}
    return len(p), nil // 返回写入的字节数
}

// Compress 函数从 io.Reader 读取数据，进行 zlib 压缩，
// 并通过返回的通道流式传输压缩后的字节数据。
func Compress(r io.Reader) <-chan BytesWithError {
    // 创建一个带缓冲的通道，以提高生产者和消费者之间的吞吐量。
    // 缓冲大小可以根据实际应用场景进行调整。
    outputChan := make(chan BytesWithError, 100) 

    go func() {
        defer close(outputChan) // 确保在Goroutine退出时关闭通道

        // 创建 ChanWriter 实例，它会将数据写入 outputChan。
        cw := ChanWriter(outputChan)

        // 使用 zlib.NewWriter 创建一个 zlib 写入器，
        // 它会将压缩后的数据写入到我们的 ChanWriter (cw)。
        zlibWriter := zlib.NewWriter(cw)
        defer func() {
            // 确保 zlibWriter 被关闭，这会刷新所有剩余的压缩数据到 cw。
            // 如果在关闭时发生错误，也通过通道发送。
            if err := zlibWriter.Close(); err != nil {
                outputChan <- BytesWithError{Err: fmt.Errorf("zlib writer close error: %w", err)}
            }
        }()

        // 用于从输入 io.Reader 读取数据的缓冲区。
        readBuffer := make([]byte, 4096) // 较大的缓冲区可以提高读取效率

        for {
            n, readErr := r.Read(readBuffer)
            if n > 0 {
                // 将读取到的未压缩数据写入 zlibWriter。
                // zlibWriter 会自动压缩数据，并通过其底层 io.Writer (cw) 写入。
                _, writeErr := zlibWriter.Write(readBuffer[:n])
                if writeErr != nil {
                    // 如果写入 zlibWriter 发生错误，通过通道发送错误并退出。
                    outputChan <- BytesWithError{Err: fmt.Errorf("zlib writer write error: %w", writeErr)}
                    return
                }
            }

            if readErr != nil {
                if readErr != io.EOF {
                    // 如果读取发生非 EOF 错误，通过通道发送错误并退出。
                    outputChan <- BytesWithError{Err: fmt.Errorf("reader error: %w", readErr)}
                }
                // 无论是 EOF 还是其他读取错误，都表示输入已结束或发生问题，Goroutine应退出。
                return 
            }
        }
    }()

    return outputChan
}

func main() {
    // 示例：压缩一个字符串
    inputString := "Hello, Go channels and zlib compression! " +
        "This is a sample string to demonstrate streaming compressed bytes." +
        "We are sending data through a channel efficiently." +
        "Repeating some content to make it longer for better compression ratio testing." +
        "Hello, Go channels and zlib compression! This is a sample string."

    // 将字符串转换为 io.Reader
    reader := bytes.NewBufferString(inputString)

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

    // 模拟消费者，从通道读取压缩数据
    var receivedCompressedData bytes.Buffer
    for dataWithError := range compressedBytesChan {
        if dataWithError.Err != nil {
            log.Fatalf("Error during compression: %v", dataWithError.Err)
        }
        if dataWithError.Data != nil {
            receivedCompressedData.Write(dataWithError.Data)
            // fmt.Printf("Received %d compressed bytes\n", len(dataWithError.Data))
        }
    }

    fmt.Printf("Original data length: %d bytes\n", len(inputString))
    fmt.Printf("Compressed data length: %d bytes\n", receivedCompressedData.Len())

    // 可选：解压验证
    zlibReader, err := zlib.NewReader(&receivedCompressedData)
    if err != nil {
        log.Fatalf("Failed to create zlib reader: %v", err)
    }
    defer zlibReader.Close()

    decompressedData, err := io.ReadAll(zlibReader)
    if err != nil {
        log.Fatalf("Failed to decompress data: %v", err)
    }

    fmt.Printf("Decompressed data length: %d bytes\n", len(decompressedData))
    if string(decompressedData) == inputString {
        fmt.Println("Decompression successful! Data matches original.")
    } else {
        fmt.Println("Decompression failed! Data does not match original.")
    }
}