基于go实现并发下载工具

本文介绍如何利用Go的并发特性开发一个下载工具,通过示例命令`go run filename.go -uri=http://www.example.com/test.mp4`,演示了如何进行文件的并发下载操作。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

/*
 * @Author: xiangcai
 * @Date: 2021-11-02 19:52:51
 * @LastEditors: xiangcai
 * @LastEditTime: 2021-11-03 15:40:23
 * @Description: file content
 */
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/sgs921107/gcommon"
	"github.com/sgs921107/glogging"
)

var (
	zapLogging  = glogging.NewZapLogging(glogging.Options{})
	zapLogger   = zapLogging.GetSugaredLogger()
	barInterval = time.Millisecond * 100
)

type Headers map[string]string

type ConcurrentDownloader struct {
	// 请求头
	headers Headers
	// 并发数
	Concurrent int
	client     http.Client
	// 下载链接
	URI            string
	filePath       string
	logger         *glogging.ZapSugaredLogger
	wg             *sync.WaitGroup
	downloadedSize int64
	size           int64
	exit           bool
}

func (c ConcurrentDownloader) SendReq(
	method string,
	body io.Reader,
	headers *Headers,
) (*http.Response, error) {
	request, err := http.NewRequest(method, c.URI, body)
	for key, value := range c.headers {
		request.Header.Add(key, value)
	}
	if headers != nil {
		for key, value := range *headers {
			request.Header.Add(key, value)
		}
	}
	if err != nil {
		c.logger.Errorw("Create Request Failed",
			"errMsg", err.Error(),
		)
		return nil, err
	}
	return c.client.Do(request)
}

func (c *ConcurrentDownloader) SubDownload(num, start, end int) error {
	defer c.wg.Done()
	c.logger.Infof("Run Task %d, Range %d - %d", num, start, end)
	subFilePath := fmt.Sprintf("%s.%d", c.filePath, num)
	// 子文件已经存在
	fileInfo, err := os.Stat(subFilePath)
	if !os.IsNotExist(err) {
		fileSize := fileInfo.Size()
		atomic.AddInt64(&c.downloadedSize, fileSize)
		// 已经下载完成
		if int(fileSize) >= end-start {
			c.logger.Infof("Task %d Already Downloaded", num)
			return nil
		}
		start += int(fileSize)
	}
	headers := &Headers{
		"Range": fmt.Sprintf("bytes=%d-%d", start, end),
	}
	resp, err := c.SendReq("GET", nil, headers)
	if err != nil {
		c.logger.Errorw("Reqeust Error",
			"taskNum", num,
		)
		return err
	}
	defer resp.Body.Close()
	file, err := os.OpenFile(subFilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		c.logger.Errorw("Open SubFile Error",
			"filename", subFilePath,
			"errMsg", err.Error(),
		)
		return err
	}
	defer file.Close()
	// 流失读取
	for !c.exit {
		buffer := make([]byte, 1024)
		n, err := resp.Body.Read(buffer)
		if n != 0 {
			atomic.AddInt64(&c.downloadedSize, int64(n))
			file.Write(buffer)
		}
		if err == io.EOF {
			break
		} else if err != nil {
			c.logger.Errorw("Read Resp Body Error",
				"errMsg", err.Error(),
			)
			return err
		}
	}
	c.logger.Infof("Task %d Done", num)
	return nil
}

// MergeSubFiles 合并所有子文件
func (c ConcurrentDownloader) MergeSubFiles() error {
	file, err := os.OpenFile(c.filePath, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer file.Close()
	for i := 0; i < c.Concurrent; i++ {
		subFile, err := os.Open(fmt.Sprintf("%s.%d", c.filePath, i))
		if err != nil {
			return err
		}
		defer subFile.Close()
		if buffer, err := ioutil.ReadAll(subFile); err != nil {
			return err
		} else {
			file.Write(buffer)
		}
		os.Remove(subFile.Name())
	}
	return nil
}

// ProgressBar 进度条
func (c *ConcurrentDownloader) ProgressBar() {
	for !c.exit {
		downloadedSize := atomic.LoadInt64(&c.downloadedSize)
		percentagePoint := downloadedSize * 100 / c.size
		var downloadedBar = make([]string, percentagePoint/2)
		var remainingBar = make([]string, (100-percentagePoint)/2)
		for i := range downloadedBar {
			downloadedBar[i] = "▋"
		}
		for i := range remainingBar {
			remainingBar[i] = "."
		}
		fmt.Print("\r")
		fmt.Printf(
			"Progress: %d%% |%s%s|",
			percentagePoint,
			strings.Join(downloadedBar, ""),
			strings.Join(remainingBar, ""),
		)
		time.Sleep(barInterval)
	}
	fmt.Println()
}

func (c *ConcurrentDownloader) singalHandler() {
	//创建监听退出chan
	ch := make(chan os.Signal, 1)
	//监听指定信号 ctrl+c kill
	signal.Notify(
		ch,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT,
		// syscall.SIGUSR1,
		// syscall.SIGUSR2,
	)
	for sig := range ch {
		logger := c.logger.With(
			"signal", sig,
		)
		switch sig {
		case syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT:
			c.exit = true
			logger.Warnw("Recv Exit Signal")
		default:
			logger.Infow("Recv A Signal")
		}
	}
}

// Download 执行下载任务
func (c *ConcurrentDownloader) Download() error {
	go c.singalHandler()
	c.logger.Infow("start downlaod",
		"url", c.URI,
	)
	size, err := c.GetTargetSize()
	if err != nil {
		c.logger.Error("Get Content Length Error",
			"errMsg", err.Error(),
		)
		return err
	}
	c.size = int64(size)
	fileInfo, _ := os.Stat(c.filePath)
	if fileInfo != nil && int(fileInfo.Size()) >= size {
		c.logger.Infow("Already Downloaded",
			"url", c.URI,
		)
		return nil
	}
	c.logger.Infof("file size: %.2fM", float64(size)/(1024*1024))
	// 每个子任务下载的大小
	subSize := size/c.Concurrent + 1
	for i := 0; i < c.Concurrent; i++ {
		start, end := subSize*i, subSize*(i+1)-1
		if end > size {
			end = size
		}
		c.wg.Add(1)
		go c.SubDownload(i, start, end)
	}
	go c.ProgressBar()
	c.wg.Wait()
	time.Sleep(barInterval)
	c.exit = true
	downloadedSize := atomic.LoadInt64(&c.downloadedSize)
	if int(downloadedSize) < size {
		c.logger.Errorw("DownloadedSizeError",
			"ExpectSize", size,
			"DownloadedSize", downloadedSize,
		)
		return err
	}
	// 合并文件
	if c.MergeSubFiles() != nil {
		c.logger.Errorw("ErrMergeSubFiles",
			"errMsg", err.Error(),
		)
		return err
	}
	return nil
}

// GetTargetSize 获取下载目标的大小
func (c ConcurrentDownloader) GetTargetSize() (int, error) {
	resp, err := c.SendReq("HEAD", nil, nil)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	size, err := strconv.Atoi(resp.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return size, nil
}

// NewDownloader 实例化downloader
func NewDownloader(
	uri, saveDir string,
	concurrent int,
	timeout time.Duration,
	headers Headers,
) (*ConcurrentDownloader, error) {
	realURI, err := url.Parse(uri)
	if err != nil {
		return nil, err
	}
	if !gcommon.PathIsExist(saveDir) {
		return nil, os.ErrNotExist
	}
	_, filename := path.Split(realURI.Path)
	if filename == "" {
		return nil, errors.New("INVALID RESOURCE")
	}
	downlaoder := &ConcurrentDownloader{
		URI:        realURI.String(),
		Concurrent: concurrent,
		logger:     zapLogger,
		wg:         &sync.WaitGroup{},
		filePath:   path.Join(saveDir, filename),
		headers:    headers,
	}
	downlaoder.client = http.Client{
		Timeout: timeout,
	}
	return downlaoder, nil
}

func init() {
	// 设置最大可用的cpu数量
	runtime.GOMAXPROCS(runtime.NumCPU())
}

func main() {
	var uri, saveDir, headers string
	var concurrent, timeout int
	flag.StringVar(&uri, "uri", "", "download link, required")
	flag.StringVar(&saveDir, "dir", ".", "download file to this dir")
	flag.IntVar(&concurrent, "n", 5, "concurrent num")
	flag.IntVar(&timeout, "t", 0, "download timeout, unit: Minute")
	flag.StringVar(&headers, "headers", "{}", "Request Headers, formater: json")
	flag.Parse()
	if uri == "" {
		zapLogger.Error("Param URI Is Required")
		flag.CommandLine.Usage()
		return
	}
	reqHeaders := make(Headers)
	if json.Unmarshal([]byte(headers), &reqHeaders) != nil {
		zapLogger.Error("Invalid Headers")
		return
	}
	downlaoder, err := NewDownloader(
		uri,
		saveDir,
		concurrent,
		time.Minute*time.Duration(timeout),
		reqHeaders,
	)
	if err != nil {
		zapLogger.Errorw("Instantiation Downloader Failed",
			"errMsg", err.Error(),
		)
		return
	}
	downlaoder.Download()
}

example cmd: go run filename.go -uri=http://www.example.com/test.mp4

本书作者带你一步一步深入这些方法。你将理解 Go语言为何选定这些并发模型,这些模型又会带来什么问题,以及你如何组合利用这些模型中的原语去解决问题。学习那些让你在独立且自信的编写与实现任何规模并发系统时所需要用到的技巧和工具。 理解Go语言如何解决并发难以编写正确这一根本问题。 学习并发与并行的关键性区别。 深入到Go语言的内存同步原语。 利用这些模式中的原语编写可维护的并发代码。 将模式组合成为一系列的实践,使你能够编写大规模的分布式系统。 学习 goroutine 背后的复杂性,以及Go语言的运行时如何将所有东西连接在一起。 作者简介 · · · · · · Katherine Cox-Buday是一名计算机科学家,目前工作于 Simple online banking。她的业余爱好包括软件工程、创作、Go 语言(igo、baduk、weiquei) 以及音乐,这些都是她长期的追求,并且有着不同层面的贡献。 目录 · · · · · · 前言 1 第1章 并发概述 9 摩尔定律,Web Scale和我们所陷入的混乱 10 为什么并发很难? 12 竞争条件 13 原子性 15 内存访问同步 17 死锁、活锁和饥饿 20 确定并发安全 28 面对复杂性的简单性 31 第2章 对你的代码建模:通信顺序进程 33 并发与并行的区别 33 什么是CSP 37 如何帮助你 40 Go语言并发哲学 43 第3章 Go语言并发组件 47 goroutine 47 sync包 58 WaitGroup 58 互斥锁和读写锁 60 cond 64 once 69 池 71 channel 76 select 语句 92 GOMAXPROCS控制 97 小结 98 第4章 Go语言并发模式 99 约束 99 for-select循环103 防止goroutine泄漏 104 or-channel 109 错误处理112 pipeline 116 构建pipeline的最佳实践 120 一些便利的生成器 126 扇入,扇出 132 or-done-channel 137 tee-channel 139 桥接channel模式 140 队列排队143 context包 151 小结 168 第5章 大规模并发 169 异常传递169 超时和取消 178 心跳 184 复制请求197 速率限制199 治愈异常的goroutine 215 小结 222 第6章 goroutine和Go语言运行时 223 工作窃取223 窃取任务还是续体 231 向开发人员展示所有这些信息 240 尾声 240 附录A 241
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值