diff --git a/internal/protocol/http/fetcher.go b/internal/protocol/http/fetcher.go index 786e7b32..cfc4f9f8 100644 --- a/internal/protocol/http/fetcher.go +++ b/internal/protocol/http/fetcher.go @@ -314,6 +314,7 @@ type Fetcher struct { prefetchDone atomic.Bool // Prefetch completed or stopped prefetchErr error // Error during prefetch (if any) prefetchStopCh chan struct{} // Signal to stop prefetch + prefetchExitCh chan struct{} // Closed when the prefetch goroutine exits resolveFallback atomic.Bool // Resolve prefetch is still available until the first Range request takes over // Target file @@ -562,6 +563,7 @@ func (f *Fetcher) Resolve(req *base.Request) error { // For non-range resources, the response will be used directly in Start if res.Range && reuseResolveBody { f.prefetchStopCh = make(chan struct{}) + f.prefetchExitCh = make(chan struct{}) f.resolveFallback.Store(true) go f.asyncPrefetch() } @@ -579,6 +581,9 @@ func (f *Fetcher) Resolve(req *base.Request) error { func (f *Fetcher) asyncPrefetch() { defer func() { f.prefetchDone.Store(true) + if f.prefetchExitCh != nil { + close(f.prefetchExitCh) + } }() // Get the resolve response @@ -654,23 +659,38 @@ func (f *Fetcher) asyncPrefetch() { } } +// stopPrefetch signals the async prefetch to stop and waits until it exits. +// Closing the resolve body unblocks a pending Read, so a stalled server cannot +// keep the prefetch goroutine writing to the temporary file after we reuse it. +func (f *Fetcher) stopPrefetch() { + if f.prefetchStopCh == nil { + return + } + + // Pause and the first connection's takeover can stop concurrently, so the + // check-and-close must be under the lock to avoid closing the channel twice. + f.resolveRespLock.Lock() + select { + case <-f.prefetchStopCh: + // Already closed + default: + close(f.prefetchStopCh) + } + if f.resolveResp != nil { + f.resolveResp.Body.Close() + f.resolveResp = nil + } + f.resolveRespLock.Unlock() + + if f.prefetchExitCh != nil { + <-f.prefetchExitCh + } +} + // stopPrefetchAndGetData stops the async prefetch and returns prefetched bytes // It also copies prefetched data to the target file func (f *Fetcher) stopPrefetchAndCopyData() int64 { - // Signal prefetch to stop (safely) - if f.prefetchStopCh != nil { - select { - case <-f.prefetchStopCh: - // Already closed - default: - close(f.prefetchStopCh) - } - } - - // Wait for prefetch to finish (with timeout) - for i := 0; f.prefetchStopCh != nil && i < 1000 && !f.prefetchDone.Load(); i++ { - time.Sleep(10 * time.Millisecond) - } + f.stopPrefetch() prefetched := f.prefetchSize.Load() if prefetched == 0 { @@ -1631,11 +1651,10 @@ func (f *Fetcher) completeFromResolveFallback(conn *connection) bool { return false } - for !f.prefetchDone.Load() { - if conn.ctx.Err() != nil { - return false - } - time.Sleep(10 * time.Millisecond) + select { + case <-f.prefetchExitCh: + case <-conn.ctx.Done(): + return false } knownSize := f.hasKnownSize() if f.prefetchErr != nil || (knownSize && f.prefetchSize.Load() != f.meta.Res.Size) { @@ -2442,14 +2461,7 @@ func (f *Fetcher) Pause() error { } // Stop prefetch if running - if f.prefetchStopCh != nil { - select { - case <-f.prefetchStopCh: - // Already closed - default: - close(f.prefetchStopCh) - } - } + f.stopPrefetch() // Wait for downloadLoop to exit first (it will call wg.Wait internally) if f.downloadLoopDone != nil { @@ -2459,11 +2471,6 @@ func (f *Fetcher) Pause() error { // Wait for all connection goroutines to stop f.wg.Wait() - // Wait for prefetch to finish - for f.prefetchStopCh != nil && !f.prefetchDone.Load() { - time.Sleep(10 * time.Millisecond) - } - // If Pause wins before the first Range response is validated, materialize // the Resolve prefix before discarding its temporary file. The first // connection was allocated after the earlier prefetch snapshot, so advance diff --git a/internal/protocol/http/fetcher_test.go b/internal/protocol/http/fetcher_test.go index a64301c2..5b58c866 100644 --- a/internal/protocol/http/fetcher_test.go +++ b/internal/protocol/http/fetcher_test.go @@ -3261,16 +3261,17 @@ func downloadResume(listener net.Listener, connections int, t *testing.T) { } func downloadWithProxy(httpListener net.Listener, proxyListener net.Listener, t *testing.T) { - fetcher := downloadReady(httpListener, 4, t) - ctl := controller.NewController() - ctl.GetProxy = func(requestProxy *base.RequestProxy) func(*gohttp.Request) (*url.URL, error) { + // Setup must run before Resolve, as in the downloader, so the Resolve + // request also goes through the proxy. + f := buildFetcher() + f.ctl.GetProxy = func(requestProxy *base.RequestProxy) func(*gohttp.Request) (*url.URL, error) { return (&base.DownloaderProxyConfig{ Enable: true, Scheme: "socks5", Host: proxyListener.Addr().String(), }).ToHandler() } - fetcher.Setup(ctl) + fetcher := doDownloadReady(f, httpListener, 4, t) err := fetcher.Start() if err != nil { t.Fatal(err) diff --git a/internal/protocol/http/stop_prefetch_test.go b/internal/protocol/http/stop_prefetch_test.go new file mode 100644 index 00000000..baaef35b --- /dev/null +++ b/internal/protocol/http/stop_prefetch_test.go @@ -0,0 +1,56 @@ +package http + +import ( + "bytes" + "fmt" + gohttp "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/GopeedLab/gopeed/pkg/base" + "github.com/GopeedLab/gopeed/pkg/protocol/http" +) + +// A stalled Resolve body must not keep the prefetch goroutine alive after +// stopPrefetchAndCopyData returns, or it writes to a file we already removed. +func TestFetcher_StopPrefetchUnblocksStalledRead(t *testing.T) { + prefix := bytes.Repeat([]byte("p"), 64*1024) + stall := make(chan struct{}) + defer close(stall) + + server := httptest.NewServer(gohttp.HandlerFunc(func(w gohttp.ResponseWriter, r *gohttp.Request) { + w.Header().Set(base.HttpHeaderAcceptRanges, base.HttpHeaderBytes) + w.Header().Set(base.HttpHeaderContentLength, fmt.Sprintf("%d", len(prefix)*4)) + w.WriteHeader(gohttp.StatusOK) + _, _ = w.Write(prefix) + w.(gohttp.Flusher).Flush() + select { + case <-stall: + case <-r.Context().Done(): + } + })) + defer server.Close() + + f := buildFetcher() + f.Meta().Opts = &base.Options{ + Path: t.TempDir(), + Name: "stalled.data", + Extra: &http.OptsExtra{Connections: 1}, + } + if err := f.Resolve(&base.Request{URL: server.URL + "/stalled.data"}); err != nil { + t.Fatal(err) + } + for f.prefetchSize.Load() < int64(len(prefix)) { + time.Sleep(time.Millisecond) + } + + start := time.Now() + f.stopPrefetch() + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("stopPrefetch took %v, want it to interrupt the stalled read", elapsed) + } + if !f.prefetchDone.Load() { + t.Fatal("prefetch goroutine still running after stopPrefetch") + } +}