Golang如何实现高效的文件并发下载 讲解sync.WaitGroup与goroutine的配合使用

package main

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

func downloadChunk(url string, start int64, end int64, chunkID int, wg *sync.WaitGroup, filename string) error {
    defer wg.Done()

    req, err := http.NewRequest("GET", url, nil)
    if err != nil {
        return err
    }

    rangeHeader := fmt.Sprintf("bytes=%d-%d", start, end)
    req.Header.Set("Range", rangeHeader)

    resp, err := http.DefaultClient.Do(req)
    if err != nil {
        return err
    }
    defer resp.Body.Close()

    if resp.StatusCode != http.StatusPartialContent {
        return fmt.Errorf("unexpected status code %d", resp.StatusCode)
    }

    chunkFilename := fmt.Sprintf("%s.part%d", filename, chunkID)
    file, err := os.Create(chunkFilename)
    if err != nil {
        return err
    }
    defer file.Close()

    _, err = io.Copy(file, resp.Body)
    if err != nil {
        return err
    }

    fmt.Printf("Chunk %d downloaded successfully\n", chunkID)
    return nil
}

func mergeChunks(filename string, numChunks int) error {
    outputFile, err := os.Create(filename)
    if err != nil {
        return err
    }
    defer outputFile.Close()

    for i := 0; i < numChunks; i++ {
        chunkFilename := fmt.Sprintf("%s.part%d", filename, i)
        chunkFile, err := os.Open(chunkFilename)
        if err != nil {
            return err
        }
        defer chunkFile.Close()

        _, err = io.Copy(outputFile, chunkFile)
        if err != nil {
            return err
        }

        os.Remove(chunkFilename) // 删除临时文件
    }

    return nil
}

func main() {
    url := "YOUR_FILE_URL" // 替换为你的文件URL
    filename := "downloaded_file.dat"
    numChunks := 4 // 并发数

    resp, err := http.Head(url)
    if err != nil {
        fmt.Println("Error getting file size:", err)
        return
    }
    defer resp.Body.Close()

    fileSize, err := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64)
    if err != nil {
        fmt.Println("Error parsing content length:", err)
        return
    }

    chunkSize := fileSize / int64(numChunks)

    var wg sync.WaitGroup
    for i := 0; i < numChunks; i++ {
        start := int64(i) * chunkSize
        end := start + chunkSize - 1
        if i == numChunks-1 {
            end = fileSize - 1
        }

        wg.Add(1)
        go func(i int, start int64, end int64) {
            err := downloadChunk(url, start, end, i, &wg, filename)
            if err != nil {
                fmt.Printf("Error downloading chunk %d: %v\n", i, err)
            }
        }(i, start, end)
    }

    wg.Wait()

    err = mergeChunks(filename, numChunks)
    if err != nil {
        fmt.Println("Error merging chunks:", err)
        return
    }

    fmt.Println("File downloaded successfully!")
}

func downloadChunkWithRetry(url string, start int64, end int64, chunkID int, wg *sync.WaitGroup, filename string, maxRetries int) error {
    defer wg.Done()

    for i := 0; i <= maxRetries; i++ {
        err := downloadChunk(url, start, end, chunkID, wg, filename)
        if err == nil {
            return nil // 下载成功
        }
        fmt.Printf("Attempt %d failed for chunk %d: %v\n", i+1, chunkID, err)
        if i < maxRetries {
            // 等待一段时间后重试
            time.Sleep(time.Duration(i+1) * time.Second)
        }
    }
    return fmt.Errorf("failed to download chunk %d after %d retries", chunkID, maxRetries)
}

// 在main函数中调用
// wg.Add(1)
// go func(i int, start int64, end int64) {
//     err := downloadChunkWithRetry(url, start, end, i, &wg, filename, 3) // 最大重试3次
//     if err != nil {
//         fmt.Printf("Error downloading chunk %d: %v\n", i, err)
//     }
// }(i, start, end)

func downloadChunk(url string, start int64, end int64, chunkID int, wg *sync.WaitGroup, filename string) error {
    defer wg.Done()

    chunkFilename := fmt.Sprintf("%s.part%d", filename, chunkID)
    var currentSize int64 = 0

    // 检查是否存在部分下载的文件
    fileInfo, err := os.Stat(chunkFilename)
    if err == nil {
        currentSize = fileInfo.Size()
        start += currentSize // 从上次中断的位置继续下载
        fmt.Printf("Resuming chunk %d from %d bytes\n", chunkID, currentSize)
    } else if !os.IsNotExist(err) {
        return err // 发生其他错误
    }

    if start > end {
        fmt.Printf("Chunk %d already fully downloaded\n", chunkID)
        return nil
    }

    req, err := http.NewRequest("GET", url, nil)
    if err != nil {
        return err
    }

    rangeHeader := fmt.Sprintf("bytes=%d-%d", start, end)
    req.Header.Set("Range", rangeHeader)

    resp, err := http.DefaultClient.Do(req)
    if err != nil {
        return err
    }
    defer resp.Body.Close()

    if resp.StatusCode != http.StatusPartialContent && resp.StatusCode != http.StatusOK {
        return fmt.Errorf("unexpected status code %d", resp.StatusCode)
    }

    file, err := os.OpenFile(chunkFilename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) // 以追加模式打开文件
    if err != nil {
        return err
    }
    defer file.Close()

    _, err = io.Copy(file, resp.Body)
    if err != nil {
        return err
    }

    fmt.Printf("Chunk %d downloaded successfully\n", chunkID)
    return nil
}