支持负载均衡下载

Author: iikira

internal/pcscommand/download.go

@@ -13,6 +13,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/cookiejar"
+	"net/url"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -47,7 +48,7 @@ type DownloadOption struct {
 	Parallel             int
 }
 
-func download(id int, downloadURL, savePath string, client *requester.HTTPClient, cfg *downloader.Config, isPrintStatus, isExecutedPermission bool) error {
+func download(id int, downloadURL, savePath string, loadBalansers []string, client *requester.HTTPClient, cfg *downloader.Config, isPrintStatus, isExecutedPermission bool) error {
 	var (
 		file     *os.File
 		writerAt io.WriterAt
@@ -87,10 +88,25 @@ func download(id int, downloadURL, savePath string, client *requester.HTTPClient
 	download := downloader.NewDownloader(downloadURL, writerAt, cfg)
 	download.SetClient(client)
 	download.TryHTTP(!pcsconfig.Config.EnableHTTPS())
+	download.AddLoadBalanceServer(loadBalansers...)
 
 	exitChan = make(chan struct{})
 
 	download.OnExecute(func() {
+		if isPrintStatus {
+			go func() {
+				for {
+					time.Sleep(1 * time.Second)
+					select {
+					case <-exitChan:
+						return
+					default:
+						download.PrintAllWorkers()
+					}
+				}
+			}()
+		}
+
 		if cfg.IsTest {
 			fmt.Printf("[%d] 测试下载开始\n\n", id)
 		}
@@ -126,19 +142,6 @@ func download(id int, downloadURL, savePath string, client *requester.HTTPClient
 		}
 	})
 
-	if isPrintStatus {
-		go func() {
-			for {
-				time.Sleep(1 * time.Second)
-				select {
-				case <-exitChan:
-					return
-				default:
-					download.PrintAllWorkers()
-				}
-			}
-		}()
-	}
 	err = download.Execute()
 	close(exitChan)
 	if err != nil {
@@ -161,27 +164,6 @@ func download(id int, downloadURL, savePath string, client *requester.HTTPClient
 	return nil
 }
 
-func getDownloadFunc(id int, savePath string, cfg *downloader.Config, isPrintStatus, isExecutedPermission bool) baidupcs.DownloadFunc {
-	if cfg == nil {
-		cfg = downloader.NewConfig()
-	}
-
-	return func(downloadURL string, jar *cookiejar.Jar) error {
-		h := requester.NewHTTPClient()
-		h.SetCookiejar(jar)
-		h.SetKeepAlive(true)
-		h.SetTimeout(10 * time.Minute)
-		setupHTTPClient(h)
-
-		err := download(id, downloadURL, savePath, h, cfg, isPrintStatus, isExecutedPermission)
-		if err != nil {
-			return err
-		}
-
-		return nil
-	}
-}
-
 // RunDownload 执行下载网盘内文件
 func RunDownload(paths []string, option DownloadOption) {
 	// 设置下载配置
@@ -331,10 +313,27 @@ func RunDownload(paths []string, option DownloadOption) {
 		}
 
 		// 获取直链, 或者以分享文件的方式获取下载链接来下载
-		var dlink string
+		var (
+			dlink  string
+			dlinks []string
+		)
+
 		if option.IsLocateDownload {
-			dlink = getDownloadLink(task.path)
+			// 提取直链下载
+			rawDlinks := getDownloadLinks(task.path)
+			if len(rawDlinks) > 0 {
+				dlink = rawDlinks[0].String()
+				dlinks = make([]string, 0, len(rawDlinks)-1)
+				for _, rawDlink := range rawDlinks[1:len(rawDlinks)] {
+					if rawDlink == nil {
+						continue
+					}
+
+					dlinks = append(dlinks, rawDlink.String())
+				}
+			}
 		} else if option.IsShareDownload {
+			// 分享下载
 			dlink = getShareDLink(task.path)
 		}
 
@@ -354,9 +353,22 @@ func RunDownload(paths []string, option DownloadOption) {
 			client.SetTimeout(20 * time.Minute)
 			client.SetKeepAlive(true)
 			setupHTTPClient(client)
-			err = download(task.ID, dlink, task.savePath, client, cfg, option.IsPrintStatus, option.IsExecutedPermission)
+			err = download(task.ID, dlink, task.savePath, dlinks, client, cfg, option.IsPrintStatus, option.IsExecutedPermission)
 		} else {
-			err = pcs.DownloadFile(task.path, getDownloadFunc(task.ID, task.savePath, cfg, option.IsPrintStatus, option.IsExecutedPermission))
+			err = pcs.DownloadFile(task.path, func(downloadURL string, jar *cookiejar.Jar) error {
+				h := requester.NewHTTPClient()
+				h.SetCookiejar(jar)
+				h.SetKeepAlive(true)
+				h.SetTimeout(10 * time.Minute)
+				setupHTTPClient(h)
+
+				err := download(task.ID, downloadURL, task.savePath, dlinks, h, cfg, option.IsPrintStatus, option.IsExecutedPermission)
+				if err != nil {
+					return err
+				}
+
+				return nil
+			})
 		}
 
 		if err != nil {
@@ -398,21 +410,21 @@ func RunLocateDownload(pcspaths ...string) {
 	}
 }
 
-func getDownloadLink(pcspath string) string {
+func getDownloadLinks(pcspath string) (dlinks []*url.URL) {
 	pcs := GetBaiduPCS()
 	dInfo, pcsError := pcs.LocateDownload(pcspath)
 	if pcsError != nil {
 		pcsCommandVerbose.Warn(pcsError.Error())
-		return ""
+		return
 	}
 
-	u := dInfo.SingleURL(pcsconfig.Config.EnableHTTPS())
-	if u == nil {
+	us := dInfo.URLStrings(pcsconfig.Config.EnableHTTPS())
+	if len(us) == 0 {
 		pcsCommandVerbose.Warn("no any url")
-		return ""
+		return
 	}
 
-	return u.String()
+	return us
 }
 
 // fileExist 检查文件是否存在,

requester/downloader/download_test.go

@@ -14,6 +14,12 @@ var (
 	url2 = "https://git.oschina.net/lufenping/pixabay_img/raw/master/tiny-20170712/lizard-2427248_1920.jpg"
 )
 
+func TestRandomNumber(t *testing.T) {
+	for i := 0; i < 10; i++ {
+		fmt.Println(RandomNumber(0, 5))
+	}
+}
+
 func TestExample(t *testing.T) {
 	DoDownload(url2, "lizard-2427248_1920.jpg", nil)
 }

requester/downloader/downloader.go

@@ -5,9 +5,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/iikira/BaiduPCS-Go/pcsutil/waitgroup"
 	"github.com/iikira/BaiduPCS-Go/pcsverbose"
 	"github.com/iikira/BaiduPCS-Go/requester"
 	"io"
+	"net/http"
 	"sync"
 	"time"
 )
@@ -28,6 +30,7 @@ type Downloader struct {
 	executeTime   time.Time
 	executed      bool
 	durl          string
+	loadBalansers []string
 	tryHTTP       bool
 	writer        io.WriterAt
 	client        *requester.HTTPClient
@@ -103,19 +106,62 @@ func (der *Downloader) Execute() error {
 	status.totalSize = resp.ContentLength
 
 	var (
-		req           = resp.Request
-		durl, referer string
-	)
+		loadBalancerResponses = make([]*LoadBalancerResponse, 0, len(der.loadBalansers)+1)
+		handleLoadBalancer    = func(req *http.Request) {
+			if req != nil {
+				if der.tryHTTP {
+					req.URL.Scheme = "http"
+				}
 
-	if req != nil {
-		referer = req.Referer()
+				loadBalancer := &LoadBalancerResponse{
+					URL:     req.URL.String(),
+					Referer: req.Referer(),
+				}
 
-		if der.tryHTTP {
-			req.URL.Scheme = "http"
+				loadBalancerResponses = append(loadBalancerResponses, loadBalancer)
+				pcsverbose.Verbosef("DEBUG: download task: URL: %s, Referer: %s\n", loadBalancer.URL, loadBalancer.Referer)
+			}
 		}
-		durl = req.URL.String()
-		pcsverbose.Verbosef("DEBUG: download task: URL: %s, Referer: %s\n", durl, referer)
+	)
+
+	handleLoadBalancer(resp.Request)
+
+	// 负载均衡
+	wg := waitgroup.NewWaitGroup(10)
+	privTimeout := der.client.Client.Timeout
+	der.client.SetTimeout(5 * time.Second)
+	for _, loadBalanser := range der.loadBalansers {
+		wg.AddDelta()
+		go func(loadBalanser string) {
+			defer wg.Done()
+
+			subResp, subErr := der.client.Req("HEAD", loadBalanser, nil, nil)
+			if subResp != nil {
+				defer subResp.Body.Close()
+			}
+			if subErr != nil {
+				pcsverbose.Verbosef("DEBUG: loadBalanser Error: %s\n", subErr)
+				return
+			}
+
+			if !ServerEqual(resp, subResp) {
+				pcsverbose.Verbosef("DEBUG: loadBalanser not equal to main server: %s\n", subErr)
+				return
+			}
+
+			if subResp.Request != nil {
+				loadBalancerResponses = append(loadBalancerResponses, &LoadBalancerResponse{
+					URL: subResp.Request.URL.String(),
+				})
+			}
+			handleLoadBalancer(subResp.Request)
+
+		}(loadBalanser)
 	}
+	wg.Wait()
+	der.client.SetTimeout(privTimeout)
+
+	loadBalancerResponseList := NewLoadBalancerResponseList(loadBalancerResponses)
 
 	//load breakpoint
 	err = der.initInstanceState()
@@ -178,16 +224,17 @@ func (der *Downloader) Execute() error {
 		writerAt = der.writer
 	}
 
-	workerInit := func(wer *Worker) {
-		wer.SetClient(der.client)
-		wer.SetCacheSize(der.config.cacheSize)
-		wer.SetWriteMutex(writeMu)
-		wer.SetReferer(referer)
-	}
-
 	for i := 0; i < der.config.parallel; i++ {
-		worker := NewWorker(int32(i), durl, writerAt)
-		workerInit(worker)
+		loadBalancer := loadBalancerResponseList.SequentialGet()
+		if loadBalancer == nil {
+			continue
+		}
+
+		worker := NewWorker(i, loadBalancer.URL, writerAt)
+		worker.SetClient(der.client)
+		worker.SetCacheSize(der.config.cacheSize)
+		worker.SetWriteMutex(writeMu)
+		worker.SetReferer(loadBalancer.Referer)
 
 		// 分配线程
 		if isRange {

requester/downloader/loadbalance.go

@@ -0,0 +1,73 @@
+package downloader
+
+import (
+	"net/http"
+	"sync/atomic"
+)
+
+// LoadBalancerResponse 负载均衡响应状态
+type LoadBalancerResponse struct {
+	URL     string
+	Referer string
+}
+
+// LoadBalancerResponseList 负载均衡列表
+type LoadBalancerResponseList struct {
+	lbr    []*LoadBalancerResponse
+	cursor int32
+}
+
+// NewLoadBalancerResponseList 初始化负载均衡列表
+func NewLoadBalancerResponseList(lbr []*LoadBalancerResponse) *LoadBalancerResponseList {
+	return &LoadBalancerResponseList{
+		lbr: lbr,
+	}
+}
+
+// SequentialGet 顺序获取
+func (lbrl *LoadBalancerResponseList) SequentialGet() *LoadBalancerResponse {
+	if len(lbrl.lbr) == 0 {
+		return nil
+	}
+
+	if int(lbrl.cursor) >= len(lbrl.lbr) {
+		lbrl.cursor = 0
+	}
+
+	lbr := lbrl.lbr[int(lbrl.cursor)]
+	atomic.AddInt32(&lbrl.cursor, 1)
+	return lbr
+}
+
+// RandomGet 随机获取
+func (lbrl *LoadBalancerResponseList) RandomGet() *LoadBalancerResponse {
+	return lbrl.lbr[RandomNumber(0, len(lbrl.lbr))]
+}
+
+// AddLoadBalanceServer 增加负载均衡服务器
+func (der *Downloader) AddLoadBalanceServer(urls ...string) {
+	der.loadBalansers = append(der.loadBalansers, urls...)
+}
+
+// ServerEqual 检测负载均衡的服务器是否一致
+func ServerEqual(resp, subResp *http.Response) bool {
+	if resp == nil || subResp == nil {
+		return false
+	}
+
+	if resp.ContentLength != subResp.ContentLength {
+		return false
+	}
+
+	header, subHeader := resp.Header, subResp.Header
+	if header.Get("Content-MD5") != subHeader.Get("Content-MD5") {
+		return false
+	}
+	if header.Get("Content-Type") != subHeader.Get("Content-Type") {
+		return false
+	}
+	if header.Get("x-bs-meta-crc32") != subHeader.Get("x-bs-meta-crc32") {
+		return false
+	}
+	return true
+}