fix(http): wait for the prefetch goroutine to exit before reusing its file

internal/protocol/http/fetcher.go

@@ -313,6 +313,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
@@ -553,6 +554,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()
 	}
@@ -570,6 +572,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
@@ -648,23 +653,38 @@ func (f *Fetcher) asyncPrefetch() {
 	}
 }
 
-// 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)
-		}
+// 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
 	}
 
-	// Wait for prefetch to finish (with timeout)
-	for i := 0; f.prefetchStopCh != nil && i < 1000 && !f.prefetchDone.Load(); i++ {
-		time.Sleep(10 * time.Millisecond)
+	// 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 {
+	f.stopPrefetch()
 
 	prefetched := f.prefetchSize.Load()
 	if prefetched == 0 {
@@ -1620,11 +1640,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) {
@@ -2428,14 +2447,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 {
@@ -2445,11 +2457,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

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)

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")
+	}
+}