Go中能否实现Writer到Reader流式传输？如何规避临时文件

需求背景 #

当前处理无法放入Worker内存的大文件，方案是将处理结果写入/tmp临时文件后再上传至S3。希望改用有限大小的字节缓冲区，实现边写入数据边流式上传，彻底避免临时文件的使用。

当前实现代码 #

import (
    "bufio"
    "bytes"
    "context"
    "fmt"
    "log"
    "os"
    "runtime"
    "strings"
    "sync"

    "github.com/aws/aws-sdk-go-v2/service/s3"
    "github.com/korovkin/limiter"
    "github.com/xitongsys/parquet-go/parquet"
    "github.com/xitongsys/parquet-go/writer"
)

func DownloadWarc(
    ctx context.Context,
    s3Client *s3.Client,
    warcs []*types.Warc,
    path string,
) error {
    key := fmt.Sprintf("parsed_warc/%s.parquet", path)

    filename := fmt.Sprintf("/tmp/%s", path)
    file, err := os.Create(filename)
    if err != nil {
        return fmt.Errorf("error creating file: %s", err)
    }
    defer file.Close()

    bytesWriter := bufio.NewWriter(file)
    pw, err := writer.NewParquetWriterFromWriter(bytesWriter, new(Page), 4)
    if err != nil {
        return fmt.Errorf("Can't create parquet writer: %s", err)
    }

    pw.RowGroupSize = 128 * 1024 * 1024 //128M
    pw.CompressionType = parquet.CompressionCodec_SNAPPY

    mutex := sync.Mutex{}
    numWorkers := runtime.NumCPU() * 2
    fmt.Printf("Using %d workers\n", numWorkers)
    limit := limiter.NewConcurrencyLimiter(numWorkers)

    for i, warc := range warcs {
        limit.Execute(func() {
            log.Printf("%d: %+v", i, warc)
            body, err := GetWarc(ctx, s3Client, warc)
            if err != nil {
                fmt.Printf("error getting warc: %s", err)
                return
            }

            page, err := Parse(body)
            if err != nil {
                key := fmt.Sprintf("unparsed_warc/%s.warc", path)
                s3Client.PutObject(
                    ctx,
                    &s3.PutObjectInput{
                        Body:   bytes.NewReader(body),
                        Bucket: &s3Record.Bucket.Name,
                        Key:    &key,
                    },
                )
                fmt.Printf("error getting page %s: %s", key, err)
                return
            }

            mutex.Lock()
            err = pw.Write(page)
            pw.Flush(true)
            mutex.Unlock()
            if err != nil {
                fmt.Printf("error writing page: %s", err)
                return
            }
        })
    }

    limit.WaitAndClose()
    err = pw.WriteStop()
    if err != nil {
        return fmt.Errorf("error writing stop: %s", err)
    }
    bytesWriter.Flush()

    file.Seek(0, 0)
    _, err = s3Client.PutObject(
        ctx,
        &s3.PutObjectInput{
            Body:   file,
            Bucket: &s3Record.Bucket.Name,
            Key:    &key,
        },
    )
    if err != nil {
        return fmt.Errorf("error uploading warc: %s", err)
    }

    return nil
}

解决方案：用io.Pipe()实现流式传输 #

完全可以实现，核心是利用Go标准库的io.Pipe()构建生产者-消费者模型：

• 生产者：在goroutine中处理WARC文件、生成Parquet数据，写入Pipe的Writer端
• 消费者：主goroutine从Pipe的Reader端读取数据，直接作为S3上传的Body

这种方式不需要临时文件，内存占用仅由Pipe的缓冲区大小（默认64KB，可通过自定义实现调整）决定，完全符合“有限内存”的要求。

修改后的实现代码 #

import (
    "bufio"
    "bytes"
    "context"
    "fmt"
    "io"
    "log"
    "runtime"
    "sync"

    "github.com/aws/aws-sdk-go-v2/service/s3"
    "github.com/korovkin/limiter"
    "github.com/xitongsys/parquet-go/parquet"
    "github.com/xitongsys/parquet-go/writer"
)

func DownloadWarc(
    ctx context.Context,
    s3Client *s3.Client,
    warcs []*types.Warc,
    path string,
    s3Bucket string, // 提取bucket参数，避免原代码中未定义的s3Record
) error {
    key := fmt.Sprintf("parsed_warc/%s.parquet", path)

    // 创建Pipe，实现流式读写
    pipeReader, pipeWriter := io.Pipe()
    errChan := make(chan error, 1)

    // 启动生产者goroutine：处理WARC并写入Parquet到Pipe
    go func() {
        defer pipeWriter.Close() // 必须关闭Writer，否则Reader会一直阻塞
        var err error

        // 用PipeWriter创建Parquet Writer
        bytesWriter := bufio.NewWriter(pipeWriter)
        pw, err := writer.NewParquetWriterFromWriter(bytesWriter, new(Page), 4)
        if err != nil {
            errChan <- fmt.Errorf("Can't create parquet writer: %s", err)
            return
        }
        pw.RowGroupSize = 128 * 1024 * 1024 //128M
        pw.CompressionType = parquet.CompressionCodec_SNAPPY

        mutex := sync.Mutex{}
        numWorkers := runtime.NumCPU() * 2
        fmt.Printf("Using %d workers\n", numWorkers)
        limit := limiter.NewConcurrencyLimiter(numWorkers)

        for i, warc := range warcs {
            // 检查ctx是否已取消，提前退出
            select {
            case <-ctx.Done():
                errChan <- ctx.Err()
                return
            default:
            }

            limit.Execute(func() {
                log.Printf("%d: %+v", i, warc)
                body, err := GetWarc(ctx, s3Client, warc)
                if err != nil {
                    log.Printf("error getting warc: %s", err)
                    return
                }

                page, err := Parse(body)
                if err != nil {
                    unparsedKey := fmt.Sprintf("unparsed_warc/%s.warc", path)
                    _, putErr := s3Client.PutObject(
                        ctx,
                        &s3.PutObjectInput{
                            Body:   bytes.NewReader(body),
                            Bucket: &s3Bucket,
                            Key:    &unparsedKey,
                        },
                    )
                    if putErr != nil {
                        log.Printf("error uploading unparsed warc %s: %s", unparsedKey, putErr)
                    }
                    log.Printf("error parsing page %s: %s", unparsedKey, err)
                    return
                }

                mutex.Lock()
                defer mutex.Unlock()
                if writeErr := pw.Write(page); writeErr != nil {
                    log.Printf("error writing page: %s", writeErr)
                    err = writeErr // 记录错误，后续统一返回
                    return
                }
                if flushErr := pw.Flush(true); flushErr != nil {
                    log.Printf("error flushing parquet writer: %s", flushErr)
                    err = flushErr
                    return
                }
            })
        }

        limit.WaitAndClose()
        if stopErr := pw.WriteStop(); stopErr != nil {
            err = fmt.Errorf("error writing stop: %s", stopErr)
        }
        if flushErr := bytesWriter.Flush(); flushErr != nil {
            err = fmt.Errorf("error flushing buffer: %s", flushErr)
        }

        if err != nil {
            errChan <- err
        }
        close(errChan)
    }()

    // 消费者：从Pipe读取数据并上传S3
    _, uploadErr := s3Client.PutObject(
        ctx,
        &s3.PutObjectInput{
            Body:   pipeReader,
            Bucket: &s3Bucket,
            Key:    &key,
        },
    )

    // 等待生产者完成，获取生产过程中的错误
    select {
    case produceErr := <-errChan:
        if produceErr != nil {
            return produceErr
        }
    case <-ctx.Done():
        return ctx.Err()
    }

    if uploadErr != nil {
        return fmt.Errorf("error uploading warc: %s", uploadErr)
    }

    return nil
}

关键改动说明 #

1. 替换临时文件为io.Pipe()：Pipe的Writer端由生产者goroutine写入，Reader端直接作为S3上传的数据源，实现边写边传。
2. 分离生产/消费逻辑：生产者在goroutine中处理WARC、生成Parquet，主goroutine专注于上传，两者通过Pipe同步。
3. 错误处理优化：用error channel收集生产者的错误，确保生产过程中的问题能被捕获并返回。
4. Context取消支持：在生产者循环中检查ctx状态，避免在请求取消后继续无效工作。
5. 移除临时文件操作：彻底删除了os.Create、file.Seek等临时文件相关代码，减少IO开销。

注意事项 #

• io.Pipe()是同步阻塞的：写入操作会等待读取端消费数据，读取操作会等待写入端产生数据，因此内存占用始终保持在Pipe的缓冲区范围内。
• Parquet Writer不是并发安全的：仍然需要用sync.Mutex保护pw.Write()和pw.Flush()操作。
• 必须关闭Pipe Writer：生产者完成后要调用pipeWriter.Close()，否则S3上传会一直等待数据，导致超时。

内容的提问来源于stack exchange，提问作者Jader Dias