fix(bt): races between building and closing the shared client (#1587)

internal/protocol/bt/fetcher.go

@@ -67,16 +67,11 @@ func (f *Fetcher) Setup(ctl *controller.Controller) {
 	return
 }
 
+// initClient must be called with lock held.
 func (f *Fetcher) initClient() (err error) {
-	lock.Lock()
-	defer lock.Unlock()
-
 	if client != nil {
 		return
 	}
-	if closeCtx == nil {
-		closeCtx, closeFunc = context.WithCancel(context.Background())
-	}
 
 	cfg = torrent.NewDefaultClientConfig()
 	cfg.Seed = true
@@ -94,12 +89,23 @@ func (f *Fetcher) initClient() (err error) {
 	}
 
 	closeCtx, closeFunc = context.WithCancel(context.Background())
-	go func() {
+	// closeClient may reset closeCtx before the goroutine runs, so pass it as an argument.
-		dnsResolver.Run(closeCtx)
+	go dnsResolver.Run(closeCtx)
-	}()
 	return
 }
 
+// addTorrentSpec holds lock from client initialization until the torrent is added,
+// so closeClient cannot close or reset the client in between.
+func (f *Fetcher) addTorrentSpec(spec *torrent.TorrentSpec) (*torrent.Torrent, bool, error) {
+	lock.Lock()
+	defer lock.Unlock()
+
+	if err := f.initClient(); err != nil {
+		return nil, false, err
+	}
+	return client.AddTorrentSpec(spec)
+}
+
 func (f *Fetcher) Resolve(req *base.Request) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
@@ -748,9 +754,6 @@ func (f *Fetcher) addTorrent(req *base.Request) (err error) {
 	if err = base.ParseReqExtra[bt.ReqExtra](req); err != nil {
 		return
 	}
-	if err = f.initClient(); err != nil {
-		return
-	}
 	schema := util.ParseSchema(req.URL)
 	privateTorrent := false
 	var spec *torrent.TorrentSpec
@@ -829,7 +832,7 @@ func (f *Fetcher) addTorrent(req *base.Request) (err error) {
 	s := &deferredFileStorage{}
 	spec.Storage = s
 	var added bool
-	f.torrent, added, err = client.AddTorrentSpec(spec)
+	f.torrent, added, err = f.addTorrentSpec(spec)
 	if err != nil {
 		return err
 	}

internal/protocol/bt/fetcher_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"sync"
 	"testing"
 
 	"github.com/GopeedLab/gopeed/internal/controller"
@@ -71,6 +72,82 @@ func TestFetcher_ResolveWithProxy(t *testing.T) {
 	}))
 }
 
+func TestFetcher_CloseAfterInit(t *testing.T) {
+	if err := closeClient(); err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 100; i++ {
+		f := buildFetcher(t).(*Fetcher)
+		f.meta.Req = &base.Request{}
+		lock.Lock()
+		err := f.initClient()
+		lock.Unlock()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := closeClient(); err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestFetcher_CloseWithoutClient(t *testing.T) {
+	if err := closeClient(); err != nil {
+		t.Fatal(err)
+	}
+
+	t.Run("restored task", func(t *testing.T) {
+		_, build := new(FetcherManager).Restore()
+		f := build(&fetcher.FetcherMeta{
+			Req:  &base.Request{URL: "./testdata/test.torrent"},
+			Opts: &base.Options{},
+		}, &fetcherData{})
+		f.Setup(controller.NewController())
+		if err := f.Close(); err != nil {
+			t.Errorf("Close() got = %v, want nil", err)
+		}
+	})
+
+	t.Run("failed resolve", func(t *testing.T) {
+		f := buildFetcher(t)
+		if err := f.Resolve(&base.Request{URL: "./testdata/not-exist.torrent"}); err == nil {
+			t.Fatal("Resolve() got = nil, want error")
+		}
+		if err := f.Close(); err != nil {
+			t.Errorf("Close() got = %v, want nil", err)
+		}
+	})
+}
+
+func TestFetcher_ResolveWhileClientCloses(t *testing.T) {
+	if err := closeClient(); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { closeClient() })
+
+	var wg sync.WaitGroup
+	for _, u := range []string{"./testdata/test.torrent", "./testdata/ubuntu-22.04-live-server-amd64.iso.torrent"} {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 10; i++ {
+				f := buildFetcher(t)
+				// The other goroutine may close the client after the torrent was added.
+				if err := f.Resolve(&base.Request{URL: u}); err != nil && err.Error() != "torrent closed during resolution" {
+					t.Errorf("Resolve(%s) got = %v", u, err)
+					return
+				}
+				f.Close()
+				if err := closeClient(); err != nil {
+					t.Errorf("closeClient() got = %v", err)
+					return
+				}
+			}
+		}()
+	}
+	wg.Wait()
+}
+
 func doResolve(t *testing.T, fetcher fetcher.Fetcher) {
 	t.Run("Resolve Single File", func(t *testing.T) {
 		err := fetcher.Resolve(&base.Request{