feat: bt fetcher reuse (#142)

go.mod

@@ -3,6 +3,7 @@ module github.com/GopeedLab/gopeed
 go 1.19
 
 require (
+	github.com/anacrolix/missinggo/v2 v2.7.2-0.20230527121029-a582b4f397b9
 	github.com/anacrolix/torrent v1.52.3
 	github.com/gorilla/handlers v1.5.1
 	github.com/gorilla/mux v1.8.0
@@ -27,7 +28,6 @@ require (
 	github.com/anacrolix/log v0.14.0 // indirect
 	github.com/anacrolix/missinggo v1.3.0 // indirect
 	github.com/anacrolix/missinggo/perf v1.0.0 // indirect
-	github.com/anacrolix/missinggo/v2 v2.7.2-0.20230527121029-a582b4f397b9 // indirect
 	github.com/anacrolix/mmsg v1.0.0 // indirect
 	github.com/anacrolix/multiless v0.3.0 // indirect
 	github.com/anacrolix/stm v0.4.0 // indirect

internal/protocol/bt/fetcher.go

@@ -7,7 +7,6 @@ import (
 	"github.com/GopeedLab/gopeed/pkg/util"
 	"github.com/anacrolix/torrent"
 	"github.com/anacrolix/torrent/metainfo"
-	"github.com/anacrolix/torrent/storage"
 	"path"
 	"path/filepath"
 	"sync"
@@ -15,7 +14,12 @@ import (
 	"time"
 )
 
-var client *torrent.Client
+var (
+	client        *torrent.Client
+	lock          sync.Mutex
+	torrentDirMap = make(map[string]string)
+	ftMap         = make(map[string]*fileTorrentImpl)
+)
 
 type Fetcher struct {
 	ctl *controller.Controller
@@ -23,10 +27,9 @@ type Fetcher struct {
 	torrent *torrent.Torrent
 	meta    *fetcher.FetcherMeta
 
-	ready    atomic.Bool
-	progress fetcher.Progress
-
-	torrentPaths map[string]string
+	torrentReady atomic.Bool
+	create       atomic.Bool
+	progress     fetcher.Progress
 }
 
 func (f *Fetcher) Name() string {
@@ -38,34 +41,29 @@ func (f *Fetcher) Setup(ctl *controller.Controller) (err error) {
 	if f.meta == nil {
 		f.meta = &fetcher.FetcherMeta{}
 	}
-	var once sync.Once
-	once.Do(func() {
-		cfg := torrent.NewDefaultClientConfig()
-		cfg.ListenPort = 0
+	return
+}
 
-		// 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
-		}
+func (f *Fetcher) initClient() (err error) {
+	lock.Lock()
+	defer lock.Unlock()
 
-		f.torrentPaths = make(map[string]string)
+	if client != nil {
+		return
+	}
+
+	cfg := torrent.NewDefaultClientConfig()
+	cfg.ListenPort = 0
+	cfg.DefaultStorage = newFileOpts(newFileClientOpts{
+		ClientBaseDir: cfg.DataDir,
+		HandleFileTorrent: func(infoHash metainfo.Hash, ft *fileTorrentImpl) {
+			if dir, ok := torrentDirMap[infoHash.String()]; ok {
+				ft.setTorrentDir(dir)
+			}
+			ftMap[infoHash.String()] = ft
+		},
 	})
+	client, err = torrent.NewClient(cfg)
 	return
 }
 
@@ -73,7 +71,14 @@ func (f *Fetcher) Resolve(req *base.Request) error {
 	if err := f.addTorrent(req.URL); err != nil {
 		return err
 	}
-	defer f.torrent.Drop()
+	go func() {
+		// recycle unused torrent resource
+		time.Sleep(time.Minute * 3)
+		if !f.create.Load() {
+			f.torrentReady.Store(false)
+			f.safeDrop()
+		}
+	}()
 	res := &base.Resource{
 		Name:    f.torrent.Name(),
 		Range:   true,
@@ -95,6 +100,7 @@ func (f *Fetcher) Resolve(req *base.Request) error {
 }
 
 func (f *Fetcher) Create(opts *base.Options) (err error) {
+	f.create.Store(true)
 	f.meta.Opts = opts
 	if len(opts.SelectFiles) == 0 {
 		opts.SelectFiles = make([]int, 0)
@@ -103,16 +109,19 @@ func (f *Fetcher) Create(opts *base.Options) (err error) {
 		}
 	}
 	if opts.Path != "" {
-		f.torrentPaths[f.meta.Res.Hash] = path.Join(f.meta.Opts.Path, f.meta.Res.RootDir)
+		torrentDirMap[f.meta.Res.Hash] = path.Join(f.meta.Opts.Path, f.meta.Res.RootDir)
+		if ft, ok := ftMap[f.meta.Res.Hash]; ok {
+			// reuse resolve fetcher
+			ft.setTorrentDir(torrentDirMap[f.meta.Res.Hash])
+		}
 	}
 
 	f.progress = make(fetcher.Progress, len(f.meta.Opts.SelectFiles))
-	f.ready.Store(false)
 	return nil
 }
 
 func (f *Fetcher) Start() (err error) {
-	if !f.ready.Load() {
+	if !f.torrentReady.Load() {
 		if err = f.addTorrent(f.meta.Req.URL); err != nil {
 			return
 		}
@@ -130,8 +139,8 @@ func (f *Fetcher) Start() (err error) {
 }
 
 func (f *Fetcher) Pause() (err error) {
-	f.ready.Store(false)
-	f.torrent.Drop()
+	f.torrentReady.Store(false)
+	f.safeDrop()
 	return
 }
 
@@ -140,13 +149,22 @@ func (f *Fetcher) Continue() (err error) {
 }
 
 func (f *Fetcher) Close() (err error) {
-	f.torrent.Drop()
+	f.safeDrop()
 	return nil
 }
 
+func (f *Fetcher) safeDrop() {
+	defer func() {
+		// ignore panic
+		_ = recover()
+	}()
+
+	f.torrent.Drop()
+}
+
 func (f *Fetcher) Wait() (err error) {
 	for {
-		if f.ready.Load() {
+		if f.torrentReady.Load() {
 			done := true
 			for _, selectIndex := range f.meta.Opts.SelectFiles {
 				file := f.torrent.Files()[selectIndex]
@@ -182,7 +200,7 @@ func (f *Fetcher) Meta() *fetcher.FetcherMeta {
 }
 
 func (f *Fetcher) Progress() fetcher.Progress {
-	if !f.ready.Load() {
+	if !f.torrentReady.Load() {
 		return f.progress
 	}
 	for i := range f.progress {
@@ -194,6 +212,9 @@ func (f *Fetcher) Progress() fetcher.Progress {
 }
 
 func (f *Fetcher) addTorrent(url string) (err error) {
+	if err = f.initClient(); err != nil {
+		return
+	}
 	schema := util.ParseSchema(url)
 	if schema == "MAGNET" {
 		f.torrent, err = client.AddMagnet(url)
@@ -216,7 +237,7 @@ func (f *Fetcher) addTorrent(url string) (err error) {
 		f.torrent.AddTrackers(announceList)
 	}
 	<-f.torrent.GotInfo()
-	f.ready.Store(true)
+	f.torrentReady.Store(true)
 	return
 }
 

internal/protocol/bt/file.go

@@ -0,0 +1,214 @@
+// modify from github.com/anacrolix/torrent/storage/file.go
+
+package bt
+
+import (
+	"fmt"
+	"github.com/anacrolix/torrent/storage"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/anacrolix/missinggo/v2"
+
+	"github.com/anacrolix/torrent/common"
+	"github.com/anacrolix/torrent/metainfo"
+	"github.com/anacrolix/torrent/segments"
+)
+
+// File-based storage for torrents, that isn't yet bound to a particular torrent.
+type fileClientImpl struct {
+	opts newFileClientOpts
+}
+
+type newFileClientOpts struct {
+	// The base directory for all downloads.
+	ClientBaseDir     string
+	HandleFileTorrent func(infoHash metainfo.Hash, ft *fileTorrentImpl)
+	FilePathMaker     storage.FilePathMaker
+	PieceCompletion   storage.PieceCompletion
+}
+
+func newFileOpts(opts newFileClientOpts) storage.ClientImplCloser {
+	if opts.FilePathMaker == nil {
+		opts.FilePathMaker = func(opts storage.FilePathMakerOpts) string {
+			var parts []string
+			if opts.Info.Name != metainfo.NoName {
+				parts = append(parts, opts.Info.Name)
+			}
+			return filepath.Join(append(parts, opts.File.Path...)...)
+		}
+	}
+	if opts.PieceCompletion == nil {
+		ret, err := storage.NewDefaultPieceCompletionForDir(opts.ClientBaseDir)
+		if err != nil {
+			ret = storage.NewMapPieceCompletion()
+		}
+		opts.PieceCompletion = ret
+	}
+	return fileClientImpl{opts}
+}
+
+func (fs fileClientImpl) Close() error {
+	return fs.opts.PieceCompletion.Close()
+}
+
+func (fs fileClientImpl) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash) (_ storage.TorrentImpl, err error) {
+	upvertedFiles := info.UpvertedFiles()
+	files := make([]file, 0, len(upvertedFiles))
+	for _, fileInfo := range upvertedFiles {
+		filePath := filepath.Join("", fs.opts.FilePathMaker(storage.FilePathMakerOpts{
+			Info: info,
+			File: &fileInfo,
+		}))
+		f := file{
+			rawPath: filePath,
+			path:    filePath,
+			length:  fileInfo.Length,
+		}
+		if f.length == 0 {
+			err = CreateNativeZeroLengthFile(f.path)
+			if err != nil {
+				err = fmt.Errorf("creating zero length file: %w", err)
+				return
+			}
+		}
+		files = append(files, f)
+	}
+	t := &fileTorrentImpl{
+		files,
+		segments.NewIndex(common.LengthIterFromUpvertedFiles(upvertedFiles)),
+		infoHash,
+		fs.opts.PieceCompletion,
+	}
+	fs.opts.HandleFileTorrent(infoHash, t)
+	return storage.TorrentImpl{
+		Piece: t.Piece,
+		Close: t.Close,
+	}, nil
+}
+
+type file struct {
+	// The safe, OS-local file path.
+	rawPath string
+	path    string
+	length  int64
+}
+
+type fileTorrentImpl struct {
+	files          []file
+	segmentLocater segments.Index
+	infoHash       metainfo.Hash
+	completion     storage.PieceCompletion
+}
+
+func (fts *fileTorrentImpl) Piece(p metainfo.Piece) storage.PieceImpl {
+	// Create a view onto the file-based torrent storage.
+	_io := fileTorrentImplIO{fts}
+	// Return the appropriate segments of this.
+	return &filePieceImpl{
+		fts,
+		p,
+		missinggo.NewSectionWriter(_io, p.Offset(), p.Length()),
+		io.NewSectionReader(_io, p.Offset(), p.Length()),
+	}
+}
+
+func (fts *fileTorrentImpl) Close() error {
+	return nil
+}
+
+func (fts *fileTorrentImpl) setTorrentDir(dir string) {
+	for i, f := range fts.files {
+		fts.files[i].path = filepath.Join(dir, f.rawPath)
+	}
+}
+
+// A helper to create zero-length files which won't appear for file-orientated storage since no
+// writes will ever occur to them (no torrent data is associated with a zero-length file). The
+// caller should make sure the file name provided is safe/sanitized.
+func CreateNativeZeroLengthFile(name string) error {
+	os.MkdirAll(filepath.Dir(name), 0o777)
+	var f io.Closer
+	f, err := os.Create(name)
+	if err != nil {
+		return err
+	}
+	return f.Close()
+}
+
+// Exposes file-based storage of a torrent, as one big ReadWriterAt.
+type fileTorrentImplIO struct {
+	fts *fileTorrentImpl
+}
+
+// Returns EOF on short or missing file.
+func (fst *fileTorrentImplIO) readFileAt(file file, b []byte, off int64) (n int, err error) {
+	f, err := os.Open(file.path)
+	if os.IsNotExist(err) {
+		// File missing is treated the same as a short file.
+		err = io.EOF
+		return
+	}
+	if err != nil {
+		return
+	}
+	defer f.Close()
+	// Limit the read to within the expected bounds of this file.
+	if int64(len(b)) > file.length-off {
+		b = b[:file.length-off]
+	}
+	for off < file.length && len(b) != 0 {
+		n1, err1 := f.ReadAt(b, off)
+		b = b[n1:]
+		n += n1
+		off += int64(n1)
+		if n1 == 0 {
+			err = err1
+			break
+		}
+	}
+	return
+}
+
+// Only returns EOF at the end of the torrent. Premature EOF is ErrUnexpectedEOF.
+func (fst fileTorrentImplIO) ReadAt(b []byte, off int64) (n int, err error) {
+	fst.fts.segmentLocater.Locate(segments.Extent{off, int64(len(b))}, func(i int, e segments.Extent) bool {
+		n1, err1 := fst.readFileAt(fst.fts.files[i], b[:e.Length], e.Start)
+		n += n1
+		b = b[n1:]
+		err = err1
+		return err == nil // && int64(n1) == e.Length
+	})
+	if len(b) != 0 && err == nil {
+		err = io.EOF
+	}
+	return
+}
+
+func (fst fileTorrentImplIO) WriteAt(p []byte, off int64) (n int, err error) {
+	// log.Printf("write at %v: %v bytes", off, len(p))
+	fst.fts.segmentLocater.Locate(segments.Extent{off, int64(len(p))}, func(i int, e segments.Extent) bool {
+		name := fst.fts.files[i].path
+		os.MkdirAll(filepath.Dir(name), 0o777)
+		var f *os.File
+		f, err = os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0o666)
+		if err != nil {
+			return false
+		}
+		var n1 int
+		n1, err = f.WriteAt(p[:e.Length], e.Start)
+		// log.Printf("%v %v wrote %v: %v", i, e, n1, err)
+		closeErr := f.Close()
+		n += n1
+		p = p[n1:]
+		if err == nil {
+			err = closeErr
+		}
+		if err == nil && int64(n1) != e.Length {
+			err = io.ErrShortWrite
+		}
+		return err == nil
+	})
+	return
+}

internal/protocol/bt/file_misc.go

@@ -0,0 +1,36 @@
+// modify from github.com/anacrolix/torrent/storage/file_misc.go
+
+package bt
+
+import "github.com/anacrolix/torrent/metainfo"
+
+type requiredLength struct {
+	fileIndex int
+	length    int64
+}
+
+func extentCompleteRequiredLengths(info *metainfo.Info, off, n int64) (ret []requiredLength) {
+	if n == 0 {
+		return
+	}
+	for i, fi := range info.UpvertedFiles() {
+		if off >= fi.Length {
+			off -= fi.Length
+			continue
+		}
+		n1 := n
+		if off+n1 > fi.Length {
+			n1 = fi.Length - off
+		}
+		ret = append(ret, requiredLength{
+			fileIndex: i,
+			length:    off + n1,
+		})
+		n -= n1
+		if n == 0 {
+			return
+		}
+		off = 0
+	}
+	return
+}

internal/protocol/bt/file_piece.go

@@ -0,0 +1,61 @@
+// modify from github.com/anacrolix/torrent/storage/file_piece.go
+
+package bt
+
+import (
+	"github.com/anacrolix/torrent/metainfo"
+	"github.com/anacrolix/torrent/storage"
+	"io"
+	"log"
+	"os"
+)
+
+type filePieceImpl struct {
+	*fileTorrentImpl
+	p metainfo.Piece
+	io.WriterAt
+	io.ReaderAt
+}
+
+var _ storage.PieceImpl = (*filePieceImpl)(nil)
+
+func (fs *filePieceImpl) pieceKey() metainfo.PieceKey {
+	return metainfo.PieceKey{InfoHash: fs.infoHash, Index: fs.p.Index()}
+}
+
+func (fs *filePieceImpl) Completion() storage.Completion {
+	c, err := fs.completion.Get(fs.pieceKey())
+	if err != nil {
+		log.Printf("error getting piece completion: %s", err)
+		c.Ok = false
+		return c
+	}
+
+	verified := true
+	if c.Complete {
+		// If it's allegedly complete, check that its constituent files have the necessary length.
+		for _, fi := range extentCompleteRequiredLengths(fs.p.Info, fs.p.Offset(), fs.p.Length()) {
+			s, err := os.Stat(fs.files[fi.fileIndex].path)
+			if err != nil || s.Size() < fi.length {
+				verified = false
+				break
+			}
+		}
+	}
+
+	if !verified {
+		// The completion was wrong, fix it.
+		c.Complete = false
+		fs.completion.Set(fs.pieceKey(), false)
+	}
+
+	return c
+}
+
+func (fs *filePieceImpl) MarkComplete() error {
+	return fs.completion.Set(fs.pieceKey(), true)
+}
+
+func (fs *filePieceImpl) MarkNotComplete() error {
+	return fs.completion.Set(fs.pieceKey(), false)
+}