Go语言中通过通道高效传递压缩字节流的实践

// BytesWithError 结构体用于在通道中传递字节切片和可能的错误
type BytesWithError struct {
    Bytes []byte
    Err   error
}

// ChanWriter 是一个实现了 io.Writer 接口的通道，用于发送 BytesWithError 结构体
type ChanWriter chan BytesWithError

// Write 方法将数据 p 包装成 BytesWithError 并发送到通道中。
// 注意：为了避免并发修改问题，这里需要对传入的 p 进行复制。
func (cw ChanWriter) Write(p []byte) (n int, err error) {
    // 创建 p 的副本，以确保发送到通道的数据是独立的，
    // 避免 p 在外部被修改导致通道中的数据不一致。
    dataCopy := make([]byte, len(p))
    copy(dataCopy, p)

    cw <- BytesWithError{Bytes: dataCopy, Err: nil}
    return len(p), nil
}

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

// Compress 函数通过通道传递压缩后的字节流。
// 它接收一个 io.Reader 作为输入，并返回一个只读的 BytesWithError 通道。
func Compress(r io.Reader) <-chan BytesWithError {
    // 创建一个带缓冲的通道，以提高生产者和消费者之间的解耦程度
    // 缓冲区大小可根据实际需求调整
    c := make(chan BytesWithError, 10) 

    go func() {
        defer close(c) // 确保在 Goroutine 结束时关闭通道

        // 创建 ChanWriter 实例，作为 zlib.NewWriter 的目标
        cw := ChanWriter(c)

        // 创建 zlib 写入器，将压缩数据写入 cw
        zw := zlib.NewWriter(cw)
        defer func() {
            if err := zw.Close(); err != nil {
                // 如果关闭 zlib 写入器时发生错误，通过通道发送
                c <- BytesWithError{Err: err}
            }
        }()

        // 使用 io.Copy 将输入读取器的数据复制到 zlib 写入器中
        // io.Copy 会自动处理分块读取和写入
        if _, err := io.Copy(zw, r); err != nil {
            // 如果在复制过程中发生错误，通过通道发送
            c <- BytesWithError{Err: err}
        }
    }()

    return c
}

import (
    "bytes"
    "fmt"
    "io"
    "log"
)

func main() {
    // 示例输入数据
    originalData := "This is a long string that will be compressed and sent through a channel. " +
        "We are testing the efficiency and correctness of the compression and channel transmission mechanism. " +
        "Go channels are powerful for concurrent programming, and combining them with io.Writer " +
        "allows for flexible data pipeline construction."

    reader := bytes.NewBufferString(originalData)

    // 调用 Compress 函数，获取一个只读通道
    compressedStream := Compress(reader)

    // 模拟消费者接收并处理压缩数据
    var receivedCompressedBytes bytes.Buffer
    for bwe := range compressedStream {
        if bwe.Err != nil {
            log.Printf("Error receiving compressed data: %v", bwe.Err)
            return
        }
        if bwe.Bytes != nil {
            receivedCompressedBytes.Write(bwe.Bytes)
            // fmt.Printf("Received %d compressed bytes\n", len(bwe.Bytes))
        }
    }

    fmt.Printf("Original data length: %d\n", len(originalData))
    fmt.Printf("Total compressed data length received: %d\n", receivedCompressedBytes.Len())

    // 可选：验证解压缩后的数据
    decompressReader, err := zlib.NewReader(&receivedCompressedBytes)
    if err != nil {
        log.Fatalf("Failed to create zlib reader: %v", err)
    }
    defer decompressReader.Close()

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

    fmt.Printf("Decompressed data length: %d\n", len(decompressedData))
    fmt.Printf("Decompressed data matches original: %t\n", string(decompressedData) == originalData)
    // fmt.Printf("Decompressed data: %s\n", string(decompressedData))
}