go下载文件-支持断点续载

package utils

import (
    "bufio"
    "context"
    "errors"
    log "zaplog" //todo替换日志库
    "github.com/rs/xid"
    "golang.org/x/sync/semaphore"
    "io"
    "net/http"
    "net/url"
    "os"
    "strconv"
    "strings"
    "sync"
    "time"
)

/**
 * fileUri: 文件下载地址
 * fileName: 下载后的文件存储(绝对)路径+文件名
 * timeout: 超时时间
 */
func Down(ctx context.Context, fileUri, fileName string, timeout time.Duration) error {
    log.InfoWithCtx(ctx, "Down fileName "+fileName)
    pathSli := strings.Split(fileName, "/")
    pathDir := strings.TrimSuffix(fileName, pathSli[len(pathSli)-1]) //去掉文件名
    pathDir = strings.TrimSuffix(pathDir, "/")                       //去掉最后的/
    err := os.MkdirAll(pathDir, 0777)                                //创建目录
    if err != nil {
        log.InfoWithCtx(ctx, "Down os.MkdirAll err:%s", err)
        return err
    }

    var file *os.File
    var size int64
    if FileIsExist(fileName) {
        log.InfoWithCtx(ctx, "down utils.FileIsExist")
        fi, err := os.OpenFile(fileName, os.O_RDWR, os.ModePerm)
        if err != nil {
            log.InfoWithCtx(ctx, "down open err:%s", err)
            return err
        }
        stat, _ := fi.Stat()
        size = stat.Size()
        log.InfoWithCtx(ctx, "down stat.Size: %v", size)
        sk, err := fi.Seek(size, 0)
        log.InfoWithCtx(ctx, "down fi.Seek: %v", sk)
        if err != nil {
            log.InfoWithCtx(ctx, "down seek err:%s", err)
            _ = fi.Close()
            return err
        }
        if sk != size {
            log.InfoWithCtx(ctx, "down seek length not equal file size,seek=%d,size=%d", sk, size)
            _ = fi.Close()
            return errors.New("seek length not equal file size")
        }
        file = fi
    } else {
        create, err := os.Create(fileName)
        if err != nil {
            log.InfoWithCtx(ctx, "down create err:%s", err)
            return err
        }
        file = create
    }
    client := &http.Client{}
    client.Timeout = timeout
    request := http.Request{}
    request.Method = http.MethodGet
    if size != 0 {
        header := http.Header{}
        header.Set("Range", "bytes="+strconv.FormatInt(size, 10)+"-")
        request.Header = header
    }
    parse, err := url.Parse(fileUri)
    if err != nil {
        log.InfoWithCtx(ctx, "down url err:%s", err)
        return err
    }
    request.URL = parse
    get, err := client.Do(&request)
    if err != nil {
        log.InfoWithCtx(ctx, "down get err:%s", err)
        return err
    }
    defer func() {
        err := get.Body.Close()
        if err != nil {
            log.InfoWithCtx(ctx, "down body close: %s", err)
        }
        err = file.Close()
        if err != nil {
            log.InfoWithCtx(ctx, "down file close: %s", err)
        }
    }()
    log.InfoWithCtx(ctx, "down content-length: %v", get.ContentLength)
    if get.ContentLength == 0 || get.ContentLength == size {
        log.InfoWithCtx(ctx, "down already downloaded")
        return nil
    }
    body := get.Body
    writer := bufio.NewWriter(file)
    bs := make([]byte, 5*1024*1024) //每次读取的最大字节数,不可为0
    for {
        var read int
        read, err = body.Read(bs)
        log.DebugWithCtx(ctx, "down read: %v, err:%s", read, err)
        if err != nil {
            if err != io.EOF {
                log.DebugWithCtx(ctx, "down read err:%s", err)
            } else {
                err = nil
                if read > 0 {
                    _, _ = writer.Write(bs[:read])
                }
            }
            break
        }
        _, err = writer.Write(bs[:read])
        if err != nil {
            log.InfoWithCtx(ctx, "down write err:%s", err)
            break
        }
    }
    if err != nil {
        return err
    }
    err = writer.Flush()
    if err != nil {
        log.InfoWithCtx(ctx, "down writer flush:%s", err)
        return err
    }
    log.InfoWithCtx(ctx, "down download success")
    return nil
}

func FileIsExist(fileName string) bool {
    _, err := os.Stat(fileName)
    return !os.IsNotExist(err)
}

func Download2File(ctx context.Context, fileUrl, path string) error {
    log.DebugWithCtx(ctx, "Download2File file "+path)
    pathSli := strings.Split(path, "/")
    pathDir := strings.TrimSuffix(path, pathSli[len(pathSli)-1]) //去掉文件名
    pathDir = strings.TrimSuffix(pathDir, "/")                   //去掉最后的/
    err := os.MkdirAll(pathDir, 0777)                            //创建目录
    if err != nil {
        log.DebugWithCtx(ctx, "Download2File os.MkdirAll err", err.Error())
        return err
    }

    out, err := os.Create(path) //创建文件
    if err != nil {
        log.DebugWithCtx(ctx, "Download2File os.Create err", err.Error())
        return err
    }
    defer out.Close()

    resp, err := http.Get(fileUrl) //下载文件
    if err != nil {
        log.DebugWithCtx(ctx, "Download2File http.Get err", err.Error())
        return err
    }
    defer resp.Body.Close()

    _, err = io.Copy(out, resp.Body) //文件内容拷贝
    if err != nil {
        log.DebugWithCtx(ctx, "Download2File io.Copy err", err.Error())
        return err
    }
    return nil
}

type DownLoadItem struct {
    Path        string //下载后的存储路径
    DownLoadUrl string //下载地址
}

// 文件下载示例
func BatchDownload(ctx context.Context, list []*DownLoadItem) error {
    //var list []*DownLoadItem
    //list = append(list, &DownLoadItem{}) //补充待下载资源
    var wg sync.WaitGroup
    var errs []error
    sem := semaphore.NewWeighted(3) //用于限制goroutine数量
    for _, v := range list {
        if v.DownLoadUrl == "" {
            return errInfo.ErrInternalServer
        }
        log.DebugWithCtx(ctx, "download uri:%v", v.DownLoadUrl)
        wg.Add(1)
        go func(downLoadUrl, path string) {
            defer wg.Done()
            _ = sem.Acquire(context.Background(), 1)
            var errf error
            isExists := FileIsExist(path)
            if !isExists { //如果文件不存在,则直接采用 io.Copy
                errf = Download2File(ctx, downLoadUrl, path)
            }
            if isExists || errf != nil { //文件存在 || 第一次下载失败
                for i := 0; i < 5; i++ {
                    errf = Down(ctx, downLoadUrl, path, 10*time.Minute) //支持断点续传
                    if errf == nil {
                        break
                    }
                }
            }
            if errf != nil {
                log.ErrorWithCtx(ctx, "download utils.Down error %s, uri:%v", errf, downLoadUrl)
                errs = append(errs, errf)
            }
            sem.Release(1)
        }(v.DownLoadUrl, v.Path)
    }
    wg.Wait()
    if len(errs) > 0 {
        return errs[0]
    }
    return nil
}

 

posted @ 2025-06-04 20:34  划水的猫  阅读(38)  评论(0)    收藏  举报