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

func Compress(r io.Reader) (<-chan byte) {
    c := make(chan byte)
    go func(){
        var wBuff bytes.Buffer // 这是一个问题，zlib.NewWriter需要io.Writer，而非bytes.Buffer的指针
        rBuff := make([]byte, 1024)
        // zlib.NewWriter期望一个io.Writer，此处传入*wBuff是错误的，因为wBuff是值类型
        // 且即使传入正确的io.Writer，wBuff也会累积所有数据，而非实时发送
        writer := zlib.NewWriter(*wBuff) 
        for {
            n, err := r.Read(rBuff)
            if err != nil && err != io.EOF { panic(err) }
            if n == 0 { break }
            writer.Write(rBuff) // 压缩并写入压缩数据
            // 如何通过通道发送已写入的压缩字节？
            // wBuff最终会包含所有压缩数据，无法实现流式发送
        }
        writer.Close()
        close(c) // 表示没有更多数据
    }()
    return c
}

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

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

// Write 方法实现了io.Writer接口
func (cw ChanWriter) Write(p []byte) (n int, err error) {
    // 为了避免并发访问时数据被修改，发送一个切片的副本
    // 否则，如果p在发送后被上游重用，接收方可能会看到不一致的数据
    bufCopy := make([]byte, len(p))
    copy(bufCopy, p)

    // 将数据发送到通道
    cw <- BytesWithError{Bytes: bufCopy, Err: nil}
    return len(p), nil // 假设写入总是成功，实际中可能需要处理通道阻塞等情况
}

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

// CompressStream 接收一个io.Reader，返回一个只读的BytesWithError通道
// 它在一个goroutine中执行压缩，并通过通道流式发送压缩后的字节切片。
func CompressStream(r io.Reader) <-chan BytesWithError {
    // 创建一个带缓冲的通道，以避免在生产者和消费者之间产生过多的阻塞
    // 缓冲区大小可以根据实际需求调整
    outputChan := make(chan BytesWithError, 10) 

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

        // 创建一个ChanWriter，它会将数据写入到outputChan
        chanWriter := ChanWriter(outputChan)

        // 使用zlib.NewWriter将压缩数据写入到我们的chanWriter中
        // zlib库会调用chanWriter.Write方法来发送压缩数据块
        zlibWriter := zlib.NewWriter(chanWriter)
        defer func() {
            // 在关闭zlibWriter之前，需要确保它将所有内部缓冲的数据都刷新到chanWriter
            if err := zlibWriter.Close(); err != nil {
                // 如果关闭时发生错误，通过通道发送错误
                outputChan <- BytesWithError{Err: err}
            }
        }()

        // 从输入io.Reader中读取数据并写入zlibWriter进行压缩
        // io.Copy是一个高效的复制函数
        if _, err := io.Copy(zlibWriter, r); err != nil {
            // 如果复制过程中发生错误，通过通道发送错误
            outputChan <- BytesWithError{Err: err}
            return // 发生错误后退出goroutine
        }
        // io.Copy完成后，zlibWriter内部可能还有未刷新数据
        // defer中的zlibWriter.Close()会负责刷新并关闭
    }()

    return outputChan
}

package main

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

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

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

// Write 方法实现了io.Writer接口
func (cw ChanWriter) Write(p []byte) (n int, err error) {
    // 为了避免并发访问时数据被修改，发送一个切片的副本
    bufCopy := make([]byte, len(p))
    copy(bufCopy, p)

    // 将数据发送到通道
    cw <- BytesWithError{Bytes: bufCopy, Err: nil}
    return len(p), nil
}

// CompressStream 接收一个io.Reader，返回一个只读的BytesWithError通道
// 它在一个goroutine中执行压缩，并通过通道流式发送压缩后的字节切片。
func CompressStream(r io.Reader) <-chan BytesWithError {
    outputChan := make(chan BytesWithError, 10) 

    go func() {
        defer close(outputChan) 

        chanWriter := ChanWriter(outputChan)
        zlibWriter := zlib.NewWriter(chanWriter)
        defer func() {
            if err := zlibWriter.Close(); err != nil {
                outputChan <- BytesWithError{Err: err}
            }
        }()

        if _, err := io.Copy(zlibWriter, r); err != nil {
            outputChan <- BytesWithError{Err: err}
            return 
        }
    }()

    return outputChan
}

func main() {
    // 模拟一个大的输入数据
    inputData := bytes.Repeat([]byte("This is some sample data to be compressed. "), 1000)
    inputReader := bytes.NewReader(inputData)

    fmt.Printf("原始数据大小: %d 字节\n", len(inputData))

    // 调用CompressStream获取压缩数据通道
    compressedDataChan := CompressStream(inputReader)

    var compressedBuffer bytes.Buffer
    var totalCompressedBytes int

    // 从通道接收压缩数据
    fmt.Println("开始接收压缩数据...")
    for dataWithError := range compressedDataChan {
        if dataWithError.Err != nil {
            log.Fatalf("压缩过程中发生错误: %v", dataWithError.Err)
        }
        if dataWithError.Bytes != nil {
            compressedBuffer.Write(dataWithError.Bytes)
            totalCompressedBytes += len(dataWithError.Bytes)
            // fmt.Printf("接收到 %d 字节的压缩数据块\n", len(dataWithError.Bytes))
        }
    }
    fmt.Println("压缩数据接收完毕。")

    fmt.Printf("总计接收压缩数据大小: %d 字节\n", totalCompressedBytes)

    // 可选：验证解压缩
    fmt.Println("\n开始解压缩验证...")
    zlibReader, err := zlib.NewReader(&compressedBuffer)
    if err != nil {
        log.Fatalf("创建zlib解压器失败: %v", err)
    }
    defer zlibReader.Close()

    decompressedBuffer := new(bytes.Buffer)
    _, err = io.Copy(decompressedBuffer, zlibReader)
    if err != nil {
        log.Fatalf("解压缩失败: %v", err)
    }

    fmt.Printf("解压缩数据大小: %d 字节\n", decompressedBuffer.Len())

    if bytes.Equal(inputData, decompressedBuffer.Bytes()) {
        fmt.Println("解压缩数据与原始数据一致。验证成功！")
    } else {
        fmt.Println("解压缩数据与原始数据不一致。验证失败！")
    }

    // 演示通道的非阻塞性（如果消费者处理慢）
    // 在实际应用中，消费者通常会尽快处理数据，或者通道有足够大的缓冲区
    time.Sleep(100 * time.Millisecond) // 给予goroutine一些时间完成
}

原始数据大小: 40000 字节
开始接收压缩数据...
压缩数据接收完毕。
总计接收压缩数据大小: 121 字节

开始解压缩验证...
解压缩数据大小: 40000 字节
解压缩数据与原始数据一致。验证成功！