Golang处理大文件上传的优化方案 分块传输与内存管理技巧

package main

import (
    "bytes"
    "fmt"
    "io"
    "mime/multipart"
    "net/http"
    "os"
    "path/filepath"
    "strconv"
)

// 假设这是一个简单的客户端上传逻辑
func uploadFileInChunks(filePath, uploadURL string, chunkSize int64) error {
    file, err := os.Open(filePath)
    if err != nil {
        return fmt.Errorf("打开文件失败: %w", err)
    }
    defer file.Close()

    fileInfo, err := file.Stat()
    if err != nil {
        return fmt.Errorf("获取文件信息失败: %w", err)
    }

    totalChunks := (fileInfo.Size() + chunkSize - 1) / chunkSize
    fileName := filepath.Base(filePath)
    fileUID := "unique_file_id_example_" + fileName // 假设一个唯一的文件ID

    fmt.Printf("开始上传文件: %s, 大小: %d字节, 分为 %d 块\n", fileName, fileInfo.Size(), totalChunks)

    for i := int64(0); i < totalChunks; i++ {
        offset := i * chunkSize
        // 使用io.LimitReader确保只读取当前块的数据
        chunkReader := io.LimitReader(file, chunkSize)

        // 准备multipart/form-data请求
        body := &bytes.Buffer{}
        writer := multipart.NewWriter(body)

        // 添加文件元数据
        writer.WriteField("file_id", fileUID)
        writer.WriteField("file_name", fileName)
        writer.WriteField("chunk_index", strconv.FormatInt(i, 10))
        writer.WriteField("total_chunks", strconv.FormatInt(totalChunks, 10))

        // 添加文件数据块
        part, err := writer.CreateFormFile("chunk_data", fileName+"_part_"+strconv.FormatInt(i, 10))
        if err != nil {
            return fmt.Errorf("创建表单文件失败: %w", err)
        }
        if _, err = io.Copy(part, chunkReader); err != nil {
            return fmt.Errorf("复制数据块失败: %w", err)
        }
        writer.Close() // 必须关闭writer才能写入结束边界

        req, err := http.NewRequest("POST", uploadURL, body)
        if err != nil {
            return fmt.Errorf("创建请求失败: %w", err)
        }
        req.Header.Set("Content-Type", writer.FormDataContentType())

        client := &http.Client{}
        resp, err := client.Do(req)
        if err != nil {
            return fmt.Errorf("发送请求失败 (块 %d): %w", i, err)
        }
        defer resp.Body.Close()

        if resp.StatusCode != http.StatusOK {
            respBody, _ := io.ReadAll(resp.Body)
            return fmt.Errorf("上传块 %d 失败, 状态码: %d, 响应: %s", i, resp.StatusCode, string(respBody))
        }
        fmt.Printf("块 %d 上传成功\n", i)
    }
    fmt.Println("文件所有块上传完成。")
    return nil
}

// 实际使用时，uploadURL 和 chunkSize 需要根据实际情况设置
// func main() {
//     // 示例调用：
//     // err := uploadFileInChunks("path/to/your/large_file.zip", "http://localhost:8080/upload_chunk", 5*1024*1024) // 5MB chunks
//     // if err != nil {
//     //     fmt.Println("上传失败:", err)
//     // }
// }

package main

import (
    "fmt"
    "io"
    "net/http"
    "os"
    "path/filepath"
    "strconv"
    "sync"
)

const uploadDir = "./uploads" // 文件存储目录

// 用于管理文件块的锁，确保并发写入时文件不损坏
var fileLocks sync.Map // map[string]*sync.Mutex

func init() {
    // 确保上传目录存在
    if _, err := os.Stat(uploadDir); os.IsNotExist(err) {
        os.Mkdir(uploadDir, 0755)
    }
}

func uploadChunkHandler(w http.ResponseWriter, r *http.Request) {
    if r.Method != "POST" {
        http.Error(w, "只支持POST请求", http.StatusMethodNotAllowed)
        return
    }

    // 解析multipart/form-data
    // 注意：这里没有设置MaxMemory，因为我们不希望将整个文件读入内存
    // ParseMultipartForm会把文件数据写入临时文件，然后我们再从临时文件读取
    // 但对于分块上传，我们更倾向于直接从请求体流式读取
    // r.ParseMultipartForm(32 << 20) // 32MB max memory for form values, files go to disk

    fileId := r.FormValue("file_id")
    fileName := r.FormValue("file_name")
    chunkIndexStr := r.FormValue("chunk_index")
    totalChunksStr := r.FormValue("total_chunks")

    if fileId == "" || fileName == "" || chunkIndexStr == "" || totalChunksStr == "" {
        http.Error(w, "缺少必要的表单字段 (file_id, file_name, chunk_index, total_chunks)", http.StatusBadRequest)
        return
    }

    chunkIndex, err := strconv.ParseInt(chunkIndexStr, 10, 64)
    if err != nil {
        http.Error(w, "chunk_index格式错误", http.StatusBadRequest)
        return
    }
    totalChunks, err := strconv.ParseInt(totalChunksStr, 10, 64)
    if err != nil {
        http.Error(w, "total_chunks格式错误", http.StatusBadRequest)
        return
    }

    // 获取文件数据块
    file, header, err := r.FormFile("chunk_data")
    if err != nil {
        http.Error(w, fmt.Sprintf("获取文件数据失败: %v", err), http.StatusBadRequest)
        return
    }
    defer file.Close()

    // 确保文件路径安全
    safeFileName := filepath.Base(fileName)
    tempFilePath := filepath.Join(uploadDir, fileId+"_"+safeFileName+".tmp") // 使用.tmp后缀表示未完成

    // 获取或创建文件锁，防止多个块同时写入同一个文件
    mu, _ := fileLocks.LoadOrStore(fileId, &sync.Mutex{})
    mutex := mu.(*sync.Mutex)
    mutex.Lock()
    defer mutex.Unlock()

    // 以追加模式打开文件，如果文件不存在则创建
    f, err := os.OpenFile(tempFilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
    if err != nil {
        http.Error(w, fmt.Sprintf("无法打开临时文件: %v", err), http.StatusInternalServerError)
        return
    }
    defer f.Close()

    // 将接收到的数据块直接写入文件
    // 这里是内存管理的关键：我们不把整个块读到内存，而是直接流式写入磁盘
    n, err := io.Copy(f, file)
    if err != nil {
        http.Error(w, fmt.Sprintf("写入文件块失败: %v", err), http.StatusInternalServerError)
        return
    }

    fmt.Printf("接收到文件 %s (ID: %s) 的块 %d/%d, 大小: %d 字节\n", safeFileName, fileId, chunkIndex, totalChunks, n)

    // 检查是否所有块都已上传
    // 这是一个简化的检查，实际生产环境需要更健壮的机制来追踪已上传的块
    // 例如，维护一个文件ID -> []bool 的映射，或者在磁盘上记录已接收的块
    // 这里我们仅通过文件大小和预期总大小来粗略判断
    fileInfo, err := os.Stat(tempFilePath)
    if err != nil {
        fmt.Println("无法获取临时文件信息:", err)
        // 但不影响当前块的成功响应
    } else {
        // 简单的计数器，实际需要更严谨的判断，比如检查所有块是否都已就位
        // 尤其是在块可能乱序到达的情况下，需要一个位图或数据库记录
        // 这里假设块是顺序到达的，或者最终文件大小可以作为判断依据
        // 更好的做法是，客户端上传完最后一个块时，发送一个“完成”请求
        // 或者服务器定期检查已接收的块数量
        fmt.Printf("当前临时文件大小: %d\n", fileInfo.Size())
    }

    // 如果是最后一个块，并且所有块都已成功写入（这里需要更复杂的逻辑验证）
    // 通常，客户端在上传完所有块后，会发送一个“合并”请求，而不是每个块都判断
    if chunkIndex == totalChunks-1 {
        // 理论上，这里应该检查所有块是否都已写入，而不是只看当前块的索引
        // 简化处理：假设最后一个块收到即完成，并重命名文件
        finalPath := filepath.Join(uploadDir, safeFileName)
        if err := os.Rename(tempFilePath, finalPath); err != nil {
            http.Error(w, fmt.Sprintf("重命名文件失败: %v", err), http.StatusInternalServerError)
            return
        }
        fmt.Printf("文件 %s (ID: %s) 上传并合并完成，存储在 %s\n", safeFileName, fileId, finalPath)
        fileLocks.Delete(fileId) // 完成后移除锁
    }

    w.WriteHeader(http.StatusOK)
    fmt.Fprintf(w, "块 %d 接收成功", chunkIndex)
}

// func main() {
//  http.HandleFunc("/upload_chunk", uploadChunkHandler)
//  fmt.Println("服务器正在监听 :8080...")
//  http.ListenAndServe(":8080", nil)
// }