nprasad2077/gopeed
1
0 Fork 0
mirror of https://github.com/GopeedLab/gopeed.git synced 2026-11-07 09:00:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: bt improve

Browse Source
This commit is contained in:
levi committed 2022-10-26 17:18:13 +08:00
1 parent c6aa83194f
commit 0fc133e46d
28 files changed
+1537 -260
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/protocol/bt/fetcher.go
+136 -23
View File
@@ -2,12 +2,16 @@ package bt
import (
"github.com/anacrolix/torrent"
"github.com/anacrolix/torrent/metainfo"
"github.com/anacrolix/torrent/storage"
"github.com/monkeyWie/gopeed-core/internal/controller"
"github.com/monkeyWie/gopeed-core/internal/fetcher"
"github.com/monkeyWie/gopeed-core/pkg/base"
"github.com/monkeyWie/gopeed-core/pkg/util"
"path/filepath"
"sync"
"sync/atomic"
"time"
)
var client *torrent.Client
@@ -16,83 +20,184 @@ type Fetcher struct {
Ctl controller.Controller
torrent *torrent.Torrent
res *base.Resource
opts *base.Options
ready atomic.Bool
progress fetcher.Progress
torrentPaths map[string]string
}
func (f *Fetcher) Setup(ctl controller.Controller) (err error) {
f.Ctl = ctl
var once sync.Once
once.Do(func() {
client, err = torrent.NewClient(nil)
cfg := torrent.NewDefaultClientConfig()
cfg.ListenPort = 0
// Support custom download path for each torrent
pieceCompletion, err := storage.NewDefaultPieceCompletionForDir(cfg.DataDir)
if err != nil {
pieceCompletion = storage.NewMapPieceCompletion()
}
clientImplCloser := storage.NewFileOpts(storage.NewFileClientOpts{
ClientBaseDir: cfg.DataDir,
TorrentDirMaker: func(baseDir string, info *metainfo.Info, infoHash metainfo.Hash) string {
if dir, ok := f.torrentPaths[infoHash.String()]; ok {
return dir
}
return baseDir
},
PieceCompletion: pieceCompletion,
})
cfg.DefaultStorage = clientImplCloser
client, err = torrent.NewClient(cfg)
if err != nil {
return
}
f.torrentPaths = make(map[string]string)
})
return
}
func (f *Fetcher) Wait() (err error) {
client.WaitAll()
for {
if f.ready.Load() {
done := true
for _, selectIndex := range f.opts.SelectFiles {
file := f.torrent.Files()[selectIndex]
if file.BytesCompleted() < file.Length() {
done = false
break
}
}
if done {
// remove unselected files
for i, file := range f.torrent.Files() {
selected := false
for _, selectIndex := range f.opts.SelectFiles {
if i == selectIndex {
selected = true
break
}
}
if !selected {
util.SafeRemove(filepath.Join(f.opts.Path, file.Path()))
}
}
break
}
}
time.Sleep(time.Millisecond * 500)
}
return nil
}
func (f *Fetcher) Resolve(req *base.Request) (res *base.Resource, err error) {
schema := util.ParseSchema(req.URL)
if schema == "MAGNET" {
f.torrent, err = client.AddMagnet(req.URL)
} else {
f.torrent, err = client.AddTorrentFromFile(req.URL)
}
if err != nil {
if err = f.addTorrent(req.URL, true); err != nil {
return
}
<-f.torrent.GotInfo()
defer f.torrent.Drop()
res = &base.Resource{
Req: req,
Size: f.torrent.Length(),
Name: f.torrent.Name(),
Range: true,
Files: make([]*base.FileInfo, len(f.torrent.Files())),
Extra: map[string]any{
"infoHash": f.torrent.InfoHash().String(),
},
}
for i, file := range f.torrent.Files() {
res.Files[i] = &base.FileInfo{
Name: filepath.Base(file.DisplayPath()),
Path: filepath.Dir(file.Path()),
Path: util.Dir(file.Path()),
Size: file.Length(),
}
res.Size += file.Length()
}
return
}
func (f *Fetcher) Create(res *base.Resource, opts *base.Options) (err error) {
f.res = res
f.opts = opts
if len(f.opts.SelectFiles) == 0 {
f.opts.SelectFiles = make([]int, 0)
for i := range f.torrent.Files() {
f.opts.SelectFiles = append(f.opts.SelectFiles, i)
}
}
if f.opts.Path != "" {
if extra, ok := res.Extra.(map[string]any); ok {
f.torrentPaths[extra["infoHash"].(string)] = f.opts.Path
}
}
f.progress = make(fetcher.Progress, len(f.opts.SelectFiles))
f.ready.Store(false)
return nil
}
func (f *Fetcher) Start() (err error) {
if len(f.opts.SelectFiles) == 0 {
if err = f.addTorrent(f.res.Req.URL, false); err != nil {
return
}
files := f.torrent.Files()
if len(f.opts.SelectFiles) == len(files) {
f.torrent.DownloadAll()
} else {
for _, index := range f.opts.SelectFiles {
f.torrent.Files()[index].Download()
for _, selectIndex := range f.opts.SelectFiles {
file := files[selectIndex]
file.Download()
}
}
return
}
func (f *Fetcher) Pause() (err error) {
f.torrent.DisallowDataDownload()
f.ready.Store(false)
f.torrent.Drop()
return
}
func (f *Fetcher) Continue() (err error) {
f.torrent.AllowDataDownload()
return
return f.Start()
}
func (f *Fetcher) Close() (err error) {
f.torrent.Drop()
return nil
}
func (f *Fetcher) Progress() fetcher.Progress {
p := make(fetcher.Progress, len(f.opts.SelectFiles))
for i := range p {
file := f.torrent.Files()[i]
p[i] = file.BytesCompleted()
if !f.ready.Load() {
return f.progress
}
return p
for i := range f.progress {
selectIndex := f.opts.SelectFiles[i]
file := f.torrent.Files()[selectIndex]
f.progress[i] = file.BytesCompleted()
}
return f.progress
}
func (f *Fetcher) addTorrent(url string, resolve bool) (err error) {
schema := util.ParseSchema(url)
if schema == "MAGNET" {
f.torrent, err = client.AddMagnet(url)
} else {
f.torrent, err = client.AddTorrentFromFile(url)
}
if err != nil {
return
}
<-f.torrent.GotInfo()
if resolve {
return
}
f.ready.Store(true)
return
}
type FetcherBuilder struct {
@@ -107,3 +212,11 @@ func (fb *FetcherBuilder) Schemes() []string {
func (fb *FetcherBuilder) Build() fetcher.Fetcher {
return &Fetcher{}
}
func (fb *FetcherBuilder) Store(f fetcher.Fetcher) (data any, err error) {
return nil, nil
}
func (fb *FetcherBuilder) Restore() (v any, f func(res *base.Resource, opts *base.Options, v any) fetcher.Fetcher) {
return nil, nil
}
internal/protocol/bt/fetcher_test.go
+4 -1
View File
@@ -19,15 +19,18 @@ func TestFetcher_Resolve_Torrent(t *testing.T) {
want := &base.Resource{
Req: res.Req,
Name: "ubuntu-22.04-live-server-amd64.iso",
Size: 1466714112,
Range: true,
Files: []*base.FileInfo{
{
Name: "ubuntu-22.04-live-server-amd64.iso",
Path: ".",
Size: 1466714112,
},
},
Extra: map[string]any{
"infoHash": "8a55cfbd5ca5d11507364765936c4f9e55b253ed",
},
}
if !reflect.DeepEqual(want, res) {
t.Errorf("Resolve() got = %v, want %v", res, want)
internal/protocol/http/fetcher.go
+71 -55
View File
@@ -6,12 +6,14 @@ import (
"fmt"
"github.com/monkeyWie/gopeed-core/internal/fetcher"
"github.com/monkeyWie/gopeed-core/pkg/base"
"github.com/monkeyWie/gopeed-core/pkg/util"
"golang.org/x/sync/errgroup"
"io"
"mime"
"net/http"
"net/http/cookiejar"
"net/url"
"os"
"path"
"path/filepath"
"strconv"
@@ -38,8 +40,9 @@ type Fetcher struct {
res *base.Resource
opts *base.Options
status base.Status
chunks []*Chunk
chunks []*chunk
file *os.File
ctx context.Context
cancel context.CancelFunc
pauseCh chan interface{}
@@ -95,20 +98,21 @@ func (f *Fetcher) Resolve(req *base.Request) (*base.Resource, error) {
contentDisposition := httpResp.Header.Get(base.HttpHeaderContentDisposition)
if contentDisposition != "" {
_, params, _ := mime.ParseMediaType(contentDisposition)
filename := params["filename"]
filename := params["filePath"]
if filename != "" {
file.Name = filename
}
}
// Get file filename by URL
// Get file filePath by URL
if file.Name == "" && strings.Count(req.URL, "/") > 2 {
file.Name = filepath.Base(req.URL)
}
// unknown file filename
// unknown file filePath
if file.Name == "" {
file.Name = "unknown"
}
res.Files = append(res.Files, file)
res.Name = file.Name
return res, nil
}
@@ -121,35 +125,13 @@ func (f *Fetcher) Create(res *base.Resource, opts *base.Options) error {
func (f *Fetcher) Start() (err error) {
// 创建文件
name := f.filename()
_, err = f.Ctl.Touch(name, f.res.Size)
name := f.filepath()
f.file, err = f.Ctl.Touch(name, f.res.Size)
if err != nil {
return err
}
f.status = base.DownloadStatusRunning
if f.res.Range {
// 每个连接平均需要下载的分块大小
chunkSize := f.res.Size / int64(f.opts.Connections)
f.chunks = make([]*Chunk, f.opts.Connections)
for i := 0; i < f.opts.Connections; i++ {
var (
begin = chunkSize * int64(i)
end int64
)
if i == f.opts.Connections-1 {
// 最后一个分块需要保证把文件下载完
end = f.res.Size - 1
} else {
end = begin + chunkSize - 1
}
chunk := NewChunk(begin, end)
f.chunks[i] = chunk
}
} else {
// 只支持单连接下载
f.chunks = make([]*Chunk, 1)
f.chunks[0] = NewChunk(0, 0)
}
f.chunks = splitChunk(f.res, f.opts)
f.fetch()
return
}
@@ -169,8 +151,7 @@ func (f *Fetcher) Continue() (err error) {
return
}
f.status = base.DownloadStatusRunning
var name = f.filename()
_, err = f.Ctl.Open(name)
f.file, err = os.OpenFile(f.filepath(), os.O_RDWR, os.ModeAppend)
if err != nil {
return err
}
@@ -178,6 +159,13 @@ func (f *Fetcher) Continue() (err error) {
return
}
func (f *Fetcher) Close() (err error) {
if err = f.Pause(); err != nil {
return
}
return
}
func (f *Fetcher) Progress() fetcher.Progress {
p := make(fetcher.Progress, 0)
if len(f.chunks) > 0 {
@@ -190,13 +178,8 @@ func (f *Fetcher) Progress() fetcher.Progress {
return p
}
func (f *Fetcher) filename() string {
// 创建文件
var filename = f.opts.Name
if filename == "" {
filename = f.res.Files[0].Name
}
return filepath.Join(f.opts.Path, filename)
func (f *Fetcher) filepath() string {
return util.Filepath(f.opts.Path, f.res.Files[0].Name, f.opts.Name)
}
func (f *Fetcher) fetch() {
@@ -212,7 +195,7 @@ func (f *Fetcher) fetch() {
go func() {
err := eg.Wait()
// 下载停止，关闭文件句柄
f.Ctl.Close(f.filename())
f.file.Close()
if f.status == base.DownloadStatusPause {
f.pauseCh <- nil
} else {
@@ -227,7 +210,6 @@ func (f *Fetcher) fetch() {
}
func (f *Fetcher) fetchChunk(index int) (err error) {
filename := f.filename()
chunk := f.chunks[index]
httpReq, err := buildRequest(f.ctx, f.res.Req)
@@ -280,7 +262,7 @@ func (f *Fetcher) fetchChunk(index int) (err error) {
for {
n, err := resp.Body.Read(buf)
if n > 0 {
_, err := f.Ctl.Write(filename, chunk.Begin+chunk.Downloaded, buf[:n])
_, err := f.file.WriteAt(buf[:n], chunk.Begin+chunk.Downloaded)
if err != nil {
return false, err
}
@@ -316,6 +298,33 @@ func (f *Fetcher) fetchChunk(index int) (err error) {
return
}
func splitChunk(res *base.Resource, opts *base.Options) (chunks []*chunk) {
if res.Range {
// 每个连接平均需要下载的分块大小
chunkSize := res.Size / int64(opts.Connections)
chunks = make([]*chunk, opts.Connections)
for i := 0; i < opts.Connections; i++ {
var (
begin = chunkSize * int64(i)
end int64
)
if i == opts.Connections-1 {
// 最后一个分块需要保证把文件下载完
end = res.Size - 1
} else {
end = begin + chunkSize - 1
}
chunk := newChunk(begin, end)
chunks[i] = chunk
}
} else {
// 只支持单连接下载
chunks = make([]*chunk, 1)
chunks[0] = newChunk(0, 0)
}
return
}
func buildClient() *http.Client {
// Cookie handle
jar, _ := cookiejar.New(nil)
@@ -339,7 +348,7 @@ func buildRequest(ctx context.Context, req *base.Request) (httpReq *http.Request
if req.Extra == nil {
method = http.MethodGet
} else {
extra := req.Extra.(Extra)
extra := req.Extra.(extra)
if extra.Method != "" {
method = extra.Method
} else {
@@ -367,8 +376,8 @@ func buildRequest(ctx context.Context, req *base.Request) (httpReq *http.Request
return httpReq, nil
}
type FetcherData struct {
Chunks []*Chunk
type fetcherData struct {
Chunks []*chunk
}
type FetcherBuilder struct {
@@ -387,19 +396,26 @@ func (fb *FetcherBuilder) Build() fetcher.Fetcher {
}
}
func (fb *FetcherBuilder) Store(f fetcher.Fetcher) interface{} {
func (fb *FetcherBuilder) Store(f fetcher.Fetcher) (data any, err error) {
_f := f.(*Fetcher)
return &FetcherData{
return &fetcherData{
Chunks: _f.chunks,
}
}, nil
}
func (fb *FetcherBuilder) Resume(res *base.Resource, opts *base.Options, data interface{}) fetcher.Fetcher {
fetcher := new(FetcherBuilder).Build().(*Fetcher)
fs := data.(*FetcherData)
fetcher.res = res
fetcher.opts = opts
fetcher.chunks = fs.Chunks
fetcher.status = base.DownloadStatusPause
return fetcher
func (fb *FetcherBuilder) Restore() (v any, f func(res *base.Resource, opts *base.Options, v any) fetcher.Fetcher) {
return &fetcherData{}, func(res *base.Resource, opts *base.Options, v any) fetcher.Fetcher {
fd := v.(*fetcherData)
fb := &FetcherBuilder{}
fetcher := fb.Build().(*Fetcher)
fetcher.res = res
fetcher.opts = opts
fetcher.status = base.DownloadStatusPause
if len(fd.Chunks) == 0 {
fetcher.chunks = splitChunk(res, opts)
} else {
fetcher.chunks = fd.Chunks
}
return fetcher
}
}
internal/protocol/http/fetcher_test.go
+17 -8
View File
@@ -4,8 +4,8 @@ import (
"fmt"
"github.com/monkeyWie/gopeed-core/internal/controller"
"github.com/monkeyWie/gopeed-core/internal/fetcher"
"github.com/monkeyWie/gopeed-core/internal/test"
"github.com/monkeyWie/gopeed-core/pkg/base"
"github.com/monkeyWie/gopeed-core/pkg/test"
"net"
"reflect"
"testing"
@@ -14,6 +14,7 @@ import (
func TestFetcher_Resolve(t *testing.T) {
testResolve(test.StartTestFileServer, &base.Resource{
Name: test.BuildName,
Size: test.BuildSize,
Range: true,
Files: []*base.FileInfo{
@@ -24,6 +25,7 @@ func TestFetcher_Resolve(t *testing.T) {
},
}, t)
testResolve(test.StartTestChunkedServer, &base.Resource{
Name: test.BuildName,
Size: 0,
Range: false,
Files: []*base.FileInfo{
@@ -148,7 +150,7 @@ func downloadNormal(listener net.Listener, connections int, t *testing.T) {
func downloadPost(listener net.Listener, connections int, t *testing.T) {
fetcher, _, _ := downloadReady(listener, connections, t)
fetcher.(*Fetcher).res.Req.Extra = Extra{
fetcher.(*Fetcher).res.Req.Extra = extra{
Method: "POST",
Header: map[string]string{
"Authorization": "Bearer 123456",
@@ -176,11 +178,11 @@ func downloadContinue(listener net.Listener, connections int, t *testing.T) {
if err != nil {
t.Fatal(err)
}
time.Sleep(time.Millisecond * 200)
time.Sleep(time.Millisecond * 50)
if err := fetcher.Pause(); err != nil {
t.Fatal(err)
}
time.Sleep(time.Millisecond * 200)
time.Sleep(time.Millisecond * 50)
if err := fetcher.Continue(); err != nil {
t.Fatal(err)
}
@@ -215,12 +217,19 @@ func downloadResume(listener net.Listener, connections int, t *testing.T) {
}
fb := new(FetcherBuilder)
time.Sleep(time.Millisecond * 200)
data := fb.Store(fetcher)
time.Sleep(time.Millisecond * 200)
time.Sleep(time.Millisecond * 50)
data, err := fb.Store(fetcher)
if err != nil {
t.Fatal(err)
}
time.Sleep(time.Millisecond * 50)
fetcher.Pause()
fetcher = fb.Resume(res, opts, data)
_, f := fb.Restore()
f(res, opts, data)
if err != nil {
t.Fatal(err)
}
fetcher.Setup(controller.NewController())
fetcher.Continue()
internal/protocol/http/model.go
+4 -4
View File
@@ -2,22 +2,22 @@ package http
import "github.com/monkeyWie/gopeed-core/pkg/base"
type Chunk struct {
type chunk struct {
Status base.Status
Begin int64
End int64
Downloaded int64
}
func NewChunk(begin int64, end int64) *Chunk {
return &Chunk{
func newChunk(begin int64, end int64) *chunk {
return &chunk{
Status: base.DownloadStatusReady,
Begin: begin,
End: end,
}
}
type Extra struct {
type extra struct {
Method string
Header map[string]string
Body string