mirror of
https://github.com/GopeedLab/gopeed.git
synced 2026-11-04 08:10:55 +00:00
refactor: 利用context做取消
This commit is contained in:
liwei committed
1 parent
79a93d5b50
commit
eee6112be5
4 files changed
+56
-51
No files matched your search
+50
-50
@@ -2,9 +2,11 @@ package http
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/monkeyWie/gopeed-core/download/base"
|
||||
"github.com/monkeyWie/gopeed-core/download/http/model"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"mime"
|
||||
@@ -39,6 +41,8 @@ type Fetcher struct {
|
||||
clients []*http.Response
|
||||
chunks []*model.Chunk
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
pauseCh chan interface{}
|
||||
}
|
||||
|
||||
@@ -58,7 +62,7 @@ func FetcherBuilder() ([]string, func() base.Fetcher) {
|
||||
}
|
||||
|
||||
func (f *Fetcher) Resolve(req *base.Request) (*base.Resource, error) {
|
||||
httpReq, err := buildRequest(req)
|
||||
httpReq, err := buildRequest(nil, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -176,7 +180,7 @@ func (f *Fetcher) Pause() (err error) {
|
||||
return
|
||||
}
|
||||
f.status = base.DownloadStatusPause
|
||||
f.stop()
|
||||
f.cancel()
|
||||
<-f.pauseCh
|
||||
return
|
||||
}
|
||||
@@ -205,30 +209,19 @@ func (f *Fetcher) filename() string {
|
||||
}
|
||||
|
||||
func (f *Fetcher) fetch() {
|
||||
errCh := make(chan error, f.opts.Connections)
|
||||
|
||||
f.ctx, f.cancel = context.WithCancel(context.Background())
|
||||
eg, _ := errgroup.WithContext(f.ctx)
|
||||
for i := 0; i < f.opts.Connections; i++ {
|
||||
go func(i int) {
|
||||
errCh <- f.fetchChunk(i, f.filename(), f.chunks[i])
|
||||
}(i)
|
||||
i := i
|
||||
eg.Go(func() error {
|
||||
return f.fetchChunk(i)
|
||||
})
|
||||
}
|
||||
|
||||
go func() {
|
||||
var err error
|
||||
stopFlag := false
|
||||
for i := 0; i < f.opts.Connections; i++ {
|
||||
fetchErr := <-errCh
|
||||
if fetchErr != nil && !stopFlag {
|
||||
// 有一个连接失败就立即终止下载
|
||||
err = fetchErr
|
||||
stopFlag = true
|
||||
f.stop()
|
||||
}
|
||||
}
|
||||
|
||||
err := eg.Wait()
|
||||
// 下载停止，关闭文件句柄
|
||||
f.Ctl.Close(f.filename())
|
||||
|
||||
if f.status == base.DownloadStatusPause {
|
||||
f.pauseCh <- nil
|
||||
} else {
|
||||
@@ -242,18 +235,11 @@ func (f *Fetcher) fetch() {
|
||||
}()
|
||||
}
|
||||
|
||||
func (f *Fetcher) stop() {
|
||||
if len(f.clients) > 0 {
|
||||
for _, client := range f.clients {
|
||||
if client != nil {
|
||||
client.Body.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
func (f *Fetcher) fetchChunk(index int) (err error) {
|
||||
filename := f.filename()
|
||||
chunk := f.chunks[index]
|
||||
|
||||
func (f *Fetcher) fetchChunk(index int, name string, chunk *model.Chunk) (err error) {
|
||||
httpReq, err := buildRequest(f.res.Req)
|
||||
httpReq, err := buildRequest(f.ctx, f.res.Req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -302,7 +288,7 @@ func (f *Fetcher) fetchChunk(index int, name string, chunk *model.Chunk) (err er
|
||||
for {
|
||||
n, err := resp.Body.Read(buf)
|
||||
if n > 0 {
|
||||
_, err := f.Ctl.Write(name, chunk.Begin+chunk.Downloaded, buf[:n])
|
||||
_, err := f.Ctl.Write(filename, chunk.Begin+chunk.Downloaded, buf[:n])
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
@@ -346,30 +332,44 @@ func buildClient() *http.Client {
|
||||
}
|
||||
}
|
||||
|
||||
func buildRequest(req *base.Request) (*http.Request, error) {
|
||||
func buildRequest(ctx context.Context, req *base.Request) (httpReq *http.Request, err error) {
|
||||
url, err := url.Parse(req.URL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
httpReq := &http.Request{
|
||||
URL: url,
|
||||
Header: map[string][]string{},
|
||||
return
|
||||
}
|
||||
|
||||
if req.Extra != nil {
|
||||
if extra, ok := req.Extra.(model.Extra); ok {
|
||||
if extra.Method != "" {
|
||||
httpReq.Method = extra.Method
|
||||
} else {
|
||||
httpReq.Method = http.MethodGet
|
||||
var (
|
||||
method string
|
||||
body io.Reader
|
||||
)
|
||||
headers := make(map[string][]string)
|
||||
if req.Extra == nil {
|
||||
method = http.MethodGet
|
||||
} else {
|
||||
extra := req.Extra.(model.Extra)
|
||||
if extra.Method != "" {
|
||||
method = extra.Method
|
||||
} else {
|
||||
method = http.MethodGet
|
||||
}
|
||||
if len(extra.Header) > 0 {
|
||||
for k, v := range extra.Header {
|
||||
headers[k] = []string{v}
|
||||
}
|
||||
if len(extra.Header) > 0 {
|
||||
for k, v := range extra.Header {
|
||||
httpReq.Header[k] = []string{v}
|
||||
}
|
||||
}
|
||||
httpReq.Body = ioutil.NopCloser(bytes.NewBufferString(extra.Body))
|
||||
}
|
||||
if extra.Body != "" {
|
||||
body = ioutil.NopCloser(bytes.NewBufferString(extra.Body))
|
||||
}
|
||||
}
|
||||
|
||||
if ctx != nil {
|
||||
httpReq, err = http.NewRequestWithContext(ctx, method, url.String(), body)
|
||||
} else {
|
||||
httpReq, err = http.NewRequest(method, url.String(), body)
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
httpReq.Header = headers
|
||||
return httpReq, nil
|
||||
}
|
||||
@@ -264,10 +264,11 @@ func downloadContinue(listener net.Listener, connections int, t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(time.Millisecond * 500)
|
||||
time.Sleep(time.Millisecond * 200)
|
||||
if err := fetcher.Pause(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(time.Millisecond * 200)
|
||||
if err := fetcher.Continue(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -9,4 +9,5 @@ require (
|
||||
github.com/marksamman/bencode v0.0.0-20150821143521-dc84f26e086e
|
||||
github.com/sirupsen/logrus v1.4.2
|
||||
golang.org/x/net v0.0.0-20200707034311-ab3426394381
|
||||
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
|
||||
)
|
||||
@@ -44,10 +44,13 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh
|
||||
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
|
||||
golang.org/x/net v0.0.0-20200707034311-ab3426394381 h1:VXak5I6aEWmAXeQjA+QSZzlgNrpq9mjcfDemuexIKsU=
|
||||
golang.org/x/net v0.0.0-20200707034311-ab3426394381/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA=
|
||||
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
|
||||
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20190422165155-953cdadca894 h1:Cz4ceDQGXuKRnVBDTS23GTn/pU5OE2C0WrNTOYK1Uuc=
|
||||
golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd h1:xhmwyvizuTgC2qz7ZlMluP20uW+C3Rm0FD/WLDX8884=
|
||||
golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
|
||||
Reference in new issue
Block a user