實現一個協程帶進度條下載器

Outlaws發表於2020-07-01

原文釋出在：lybc.site/golang-doawloader-with-go...

下載一個檔案

我們正常使用 go 語言下載一個檔案應該是這樣的：

package main

import (
    "net/http"
    "io/ioutil"
)

func main() {
    resourceUrl := "https://www.xxx.yyy/aaa.jpg"

    // Get the data
    resp, err := http.Get(resourceUrl)
    if err != nil {
        return err
    }
    defer resp.Body.Close()
    if data, err := ioutil.ReadAll(resp.Body); err = nil {
        ioutil.WriteFile("aaa.jpg", data, 0644)
    }
}

這種寫法處理小檔案沒有什麼問題，ioutil.ReadAll 方法會將檔案先讀取到記憶體中，一旦需要下載視訊類或其他大檔案時，很有可能造成 OOM 的問題。為了避免這個問題我們通常會使用 io.Copy

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs. It returns the number of bytes
// copied and the first error encountered while copying, if any.
//
// A successful Copy returns err == nil, not err == EOF.
// Because Copy is defined to read from src until EOF, it does
// not treat an EOF from Read as an error to be reported.
//
// If src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
// Otherwise, if dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copy(dst Writer, src Reader) (written int64, err error)

那麼我們下載檔案的程式碼可以改成這樣：

func main {
    resp, err := http.Get(resourceUrl)
    if err != nil {
        return nil
    }
    defer resp.Body.Close()
    tmpFile, err := os.Create("filename.tmp")
    if err != nil {
        tmpFile.Close()
        return err
    }
    if _, err := io.Copy(tmpFile, resp.Body; err != nil {
        return err
    }
    os.Rename("filename.tmp", "filename")

}

非同步化

同步下載檔案效率太低，無法重複利用到頻寬，我們利用協程將這一過程非同步化，最簡單的外面包一層 go 就完事了

go func() {
    // 省略下載過程
}()

go 語言開協程的開銷很低，為了避免協程開太多導致一些不可預知的意外我們需要控制一下協程的數量，實現一個簡單的協程池：

// 預設協程池的長度等於CPU的核數
pool := make(chan int, runtime.NumCPU)

for {
    go func() {
        pool <- 1
        // 省略下載過程
        <- pool
    }()
}

任務開始前將ID塞入協程池，任務結束後退出，這樣就可以控制到同時進行下載的協程數量。

為了避免協程還在執行時主程式退出，我們還需要加入 WaitGroup 等待所有協程執行結束

pool := make(chan int, runtime.NumCPU)
wg := sync.WaitGroup{}
for {
    wg.Add(1)
    go func() {
        defer wg.Done()
        pool <- 1
        // 省略下載過程
        <- pool
    }()
}
wg.Wait()

進度條

手工實現進度條有點麻煩，使用了一個開源庫 vbauerster/mpb

完整程式

將以上過程封裝在一起後是這樣的：

package utils

import (
    "fmt"
    "github.com/vbauerster/mpb/v5"
    "github.com/vbauerster/mpb/v5/decor"
    "io"
    "net/http"
    "os"
    "runtime"
    "strconv"
    "sync"
)

type Resource struct {
    Filename string
    Url string
}

type Downloader struct {
    wg *sync.WaitGroup
    pool chan *Resource
    Concurrent int
    HttpClient http.Client
    TargetDir string
    Resources []Resource
}

func NewDownloader(targetDir string) *Downloader {
    concurrent := runtime.NumCPU()
    return &Downloader{
        wg: &sync.WaitGroup{},
        TargetDir: targetDir,
        Concurrent: concurrent,
    }
}

func (d *Downloader) AppendResource(filename, url string) {
    d.Resources = append(d.Resources, Resource{
        Filename: filename,
        Url:      url,
    })
}

func (d *Downloader) Download(resource Resource, progress *mpb.Progress) error {
    defer d.wg.Done()
    d.pool <- &resource
    finalPath := d.TargetDir + "/" + resource.Filename
    // 建立臨時檔案
    target, err := os.Create(finalPath + ".tmp")
    if err != nil {
        return err
    }

    // 開始下載
    req, err := http.NewRequest(http.MethodGet, resource.Url, nil)
    if err != nil {
        return err
    }

    resp, err := http.DefaultClient.Do(req)
    if err != nil {
        target.Close()
        return err
    }
    defer resp.Body.Close()
    fileSize, _ := strconv.Atoi(resp.Header.Get("Content-Length"))
    // 建立一個進度條
    bar := progress.AddBar(
        int64(fileSize),
        // 進度條前的修飾
        mpb.PrependDecorators(
            decor.CountersKibiByte("% .2f / % .2f"), // 已下載數量
            decor.Percentage(decor.WCSyncSpace),     // 進度百分比
        ),
        // 進度條後的修飾
        mpb.AppendDecorators(
            decor.EwmaETA(decor.ET_STYLE_GO, 90),
            decor.Name(" ] "),
            decor.EwmaSpeed(decor.UnitKiB, "% .2f", 60),
        ),
    )
    reader := bar.ProxyReader(resp.Body)
    defer reader.Close()
    // 將下載的檔案流拷貝到臨時檔案
    if _, err := io.Copy(target, reader); err != nil {
        target.Close();
        return err
    }

    // 關閉臨時並修改臨時檔案為最終檔案
    target.Close()
    if err := os.Rename(finalPath + ".tmp", finalPath); err != nil {
        return err
    }
    <- d.pool
    return nil
}

func (d *Downloader) Start() error {
    d.pool = make(chan *Resource, d.Concurrent)
    fmt.Println("開始下載，當前併發：", d.Concurrent)
    p := mpb.New(mpb.WithWaitGroup(d.wg))
    for _, resource := range d.Resources {
        d.wg.Add(1)
        go d.Download(resource, p)
    }
    p.Wait()
    d.wg.Wait()
    return nil
}

用法

downloader := NewDownloader("./")
downloader.AppendResource("001.jpg", "http://222.186.12.239:10010/ksacf_20190731/001.jpg")
downloader.AppendResource("002.jpg", "http://222.186.12.239:10010/ksacf_20190731/002.jpg")
downloader.AppendResource("003.jpg", "http://222.186.12.239:10010/ksacf_20190731/003.jpg")
downloader.AppendResource("004.jpg", "http://222.186.12.239:10010/ksacf_20190731/004.jpg")
// 可自主調整協程數量，預設為CPU核數
downloader.Concurrent = 3
downloader.Start()

效果

圖片太大上傳失敗，請移步原文檢視：lybc.site/golang-doawloader-with-go...

本作品採用《CC 協議》，轉載必須註明作者和本文連結

相關文章