使用aws-go-sdk从Amazon S3下载CSV时如何暂停/恢复下载？

这个问题我之前在处理大S3文件的时候也遇到过——当下载速度远超处理速度时，内存里堆的未处理数据会越来越多，直到爆内存。核心思路是让下载速度适配处理速度，而不是让下载一直跑满带宽。这里有两个实用的方案，结合aws-go-sdk就能实现：

方案1：利用流式读取+有界通道实现自动“暂停”（速率匹配） #

AWS Go SDK的GetObject返回的resp.Body是一个io.ReadCloser，我们可以通过有界通道来控制读取速率：当处理goroutine跟不上下载速度时，通道会被填满，读取操作会自动阻塞，从而暂停下载；当处理完部分数据后，通道有空位，下载会自动恢复。

代码示例 #

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

func main() {
	// 初始化S3客户端（假设已完成配置）
	cfg, err := aws.LoadDefaultConfig(context.TODO())
	if err != nil {
		panic(fmt.Sprintf("加载AWS配置失败: %v", err))
	}
	s3Client := s3.NewFromConfig(cfg)

	bucket := "your-target-bucket"
	key := "large-data.csv"

	// 获取S3对象的流式读取器
	resp, err := s3Client.GetObject(context.TODO(), &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		panic(fmt.Sprintf("获取S3对象失败: %v", err))
	}
	defer resp.Body.Close()

	// 创建有界通道，限制同时待处理的chunk数量（根据内存调整，比如10个）
	chunkChan := make(chan []byte, 10)
	var wg sync.WaitGroup

	// 启动CSV处理goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		var leftover []byte // 保存上一个chunk的不完整行
		reader := csv.NewReader(nil)

		for chunk := range chunkChan {
			// 拼接上一个chunk的剩余内容，避免行被截断
			fullData := append(leftover, chunk...)
			reader = csv.NewReader(bufio.NewReader(bytes.NewReader(fullData)))

			for {
				record, err := reader.Read()
				if err == io.EOF {
					// 保存当前未完成的行，用于下一个chunk拼接
					tempLeftover, _ := reader.ReadAll()
					if len(tempLeftover) > 0 {
						leftover = []byte(tempLeftover[0])
					} else {
						leftover = nil
					}
					break
				}
				if err != nil {
					fmt.Printf("解析CSV行失败: %v\n", err)
					continue
				}
				// 执行你的业务处理逻辑（这里模拟慢处理）
				processRecord(record)
			}
		}

		// 处理最后剩余的不完整行
		if len(leftover) > 0 {
			reader = csv.NewReader(bytes.NewReader(leftover))
			if record, err := reader.Read(); err == nil {
				processRecord(record)
			}
		}
	}()

	// 从S3流读取数据，发送到通道
	buf := make([]byte, 1024*1024) // 1MB的读取缓冲
	for {
		n, err := resp.Body.Read(buf)
		if n > 0 {
			// 通道满时会自动阻塞，暂停下载
			chunkChan <- buf[:n]
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			fmt.Printf("读取S3流失败: %v\n", err)
			break
		}
	}

	close(chunkChan)
	wg.Wait()
	fmt.Println("所有数据处理完成")
}

// 模拟业务处理逻辑（这里加了休眠模拟慢处理）
func processRecord(record []string) {
	time.Sleep(10 * time.Millisecond)
	fmt.Printf("已处理记录: %v\n", record)
}

方案2：分段下载+手动控制实现主动暂停/恢复 #

如果需要手动触发暂停/恢复（比如用户输入、外部信号），可以将大文件拆分为多个小分片，每次只下载一个分片，处理完成后再请求下一个。通过原子变量或context控制下载流程，实现主动暂停。

代码示例 #

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"sync/atomic"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

// 原子变量控制暂停状态
var isPaused atomic.Bool

func main() {
	// 初始化S3客户端
	cfg, err := aws.LoadDefaultConfig(context.TODO())
	if err != nil {
		panic(fmt.Sprintf("加载AWS配置失败: %v", err))
	}
	s3Client := s3.NewFromConfig(cfg)

	bucket := "your-target-bucket"
	key := "large-data.csv"

	// 获取文件总大小
	headResp, err := s3Client.HeadObject(context.TODO(), &s3.HeadObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		panic(fmt.Sprintf("获取S3文件元数据失败: %v", err))
	}
	fileSize := *headResp.ContentLength

	// 分片大小设置为10MB（可根据内存调整）
	chunkSize := int64(10 * 1024 * 1024)
	var start int64 = 0
	var leftover []byte

	// 启动监听输入的goroutine，用于手动暂停/恢复
	go func() {
		scanner := bufio.NewScanner(os.Stdin)
		fmt.Println("输入 'pause' 暂停，输入 'resume' 恢复")
		for scanner.Scan() {
			input := scanner.Text()
			switch input {
			case "pause":
				isPaused.Store(true)
				fmt.Println("下载已暂停")
			case "resume":
				isPaused.Store(false)
				fmt.Println("下载已恢复")
			}
		}
	}()

	for start < fileSize {
		// 检查是否处于暂停状态，等待恢复
		for isPaused.Load() {
			time.Sleep(500 * time.Millisecond)
		}

		// 计算当前分片的结束位置
		end := start + chunkSize - 1
		if end >= fileSize {
			end = fileSize - 1
		}

		// 请求指定范围的分片数据
		resp, err := s3Client.GetObject(context.TODO(), &s3.GetObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
			Range:  aws.String(fmt.Sprintf("bytes=%d-%d", start, end)),
		})
		if err != nil {
			fmt.Printf("下载分片失败: %v\n", err)
			break
		}

		// 读取分片内容
		chunkData, err := io.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			fmt.Printf("读取分片内容失败: %v\n", err)
			break
		}

		// 处理分片数据（同样需要处理跨分片的行）
		fullData := append(leftover, chunkData...)
		reader := csv.NewReader(bufio.NewReader(bytes.NewReader(fullData)))
		var tempLeftover []string

		for {
			record, err := reader.Read()
			if err == io.EOF {
				tempLeftover, _ = reader.ReadAll()
				break
			}
			if err != nil {
				fmt.Printf("解析CSV行失败: %v\n", err)
				continue
			}
			processRecord(record)
		}

		// 更新剩余的不完整行
		if len(tempLeftover) > 0 {
			leftover = []byte(tempLeftover[0])
		} else {
			leftover = nil
		}

		start = end + 1
		fmt.Printf("已处理到字节位置: %d\n", start)
	}

	// 处理最后剩余的行
	if len(leftover) > 0 {
		reader := csv.NewReader(bytes.NewReader(leftover))
		if record, err := reader.Read(); err == nil {
			processRecord(record)
		}
	}
	fmt.Println("所有数据处理完成")
}

func processRecord(record []string) {
	time.Sleep(10 * time.Millisecond)
	fmt.Printf("已处理记录: %v\n", record)
}

关键注意事项 #

1. CSV行截断处理：无论哪种方案，都要保存上一个分片/chunk的不完整行，和下一部分数据拼接后再解析，避免行被截断导致解析错误。
2. 分片大小调整：根据你的内存容量和处理速度调整分片/chunk的大小，平衡内存占用和处理效率。
3. 错误重试：实际生产中建议给S3请求添加重试逻辑，避免网络波动导致的下载失败。

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