fix: Abort stalled plugin downloads quickly and retry them

managedplugin/download.go

@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/avast/retry-go/v5"
@@ -363,18 +364,42 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop
 			return err
 		}
 
+		// A server that stops sending — headers or body — otherwise holds the
+		// attempt until some middlebox kills the connection minutes later. The
+		// watchdog cancels the attempt after downloadStallTimeout without
+		// progress, and stallErr turns that cancellation into a retryable error
+		// instead of the terminal context.Canceled the caller's own cancel gets.
+		attemptCtx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		var stalled atomic.Bool
+		watchdog := time.AfterFunc(downloadStallTimeout, func() {
+			stalled.Store(true)
+			cancel()
+		})
+		defer watchdog.Stop()
+		stallErr := func(err error) error {
+			if !stalled.Load() {
+				return nil
+			}
+			return fmt.Errorf("%w: no data from %s for %s: %v", errDownloadStalled, urlForLog, downloadStallTimeout, redactURLError(err))
+		}
+
 		// Get the data
-		req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
+		req, err := http.NewRequestWithContext(attemptCtx, http.MethodGet, downloadURL, nil)
 		if err != nil {
 			return fmt.Errorf("failed create request %s: %w", urlForLog, redactURLError(err))
 		}
 
 		// Do http request
 		resp, err := http.DefaultClient.Do(req)
 		if err != nil {
+			if serr := stallErr(err); serr != nil {
+				return serr
+			}
 			return fmt.Errorf("failed to get url %s: %w", urlForLog, redactURLError(err))
 		}
 		defer resp.Body.Close()
+		watchdog.Reset(downloadStallTimeout)
 		// Check server response
 		if resp.StatusCode == http.StatusNotFound {
 			return errNotFound
@@ -398,8 +423,11 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop
 		}
 
 		// Write the body to file
-		written, err := io.Copy(io.MultiWriter(writers...), resp.Body)
+		written, err := io.Copy(io.MultiWriter(writers...), &stallResetReader{r: resp.Body, watchdog: watchdog, timeout: downloadStallTimeout})
 		if err != nil {
+			if serr := stallErr(err); serr != nil {
+				return serr
+			}
 			return fmt.Errorf("failed to copy body to file %s: %w", out.Name(), err)
 		}
 		if resp.ContentLength >= 0 && written != resp.ContentLength {
@@ -417,6 +445,23 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop
 	return checksum, nil
 }
 
+// stallResetReader defers the stall watchdog every time bytes arrive, so it
+// only fires when the peer stops sending entirely — a slow but moving
+// download never trips it.
+type stallResetReader struct {
+	r        io.Reader
+	watchdog *time.Timer
+	timeout  time.Duration
+}
+
+func (s *stallResetReader) Read(p []byte) (n int, err error) {
+	n, err = s.r.Read(p)
+	if n > 0 {
+		s.watchdog.Reset(s.timeout)
+	}
+	return n, err
+}
+
 func truncateFile(f *os.File) error {
 	if err := f.Truncate(0); err != nil {
 		return fmt.Errorf("failed to truncate file %s: %w", f.Name(), err)

managedplugin/download_retry.go

@@ -10,18 +10,21 @@ import (
 	"net/url"
 	"strings"
 	"syscall"
+	"time"
 )
 
 var (
-	errNotFound  = errors.New("not found")
-	errShortRead = errors.New("truncated response body")
+	errNotFound        = errors.New("not found")
+	errShortRead       = errors.New("truncated response body")
+	errDownloadStalled = errors.New("download stalled")
 )
 
 // Overridable so tests do not pay the real backoff.
 var (
 	downloadRetryAttempts = uint(RetryAttempts)
 	downloadRetryDelay    = RetryWaitTime
 	downloadRetryMaxDelay = MaxRetryWaitTime
+	downloadStallTimeout  = 30 * time.Second
 )
 
 type httpStatusError struct {
@@ -81,7 +84,8 @@ func isRetryableDownloadError(err error) bool {
 	}
 
 	switch {
-	case errors.Is(err, errShortRead),
+	case errors.Is(err, errDownloadStalled),
+		errors.Is(err, errShortRead),
 		errors.Is(err, io.ErrUnexpectedEOF),
 		errors.Is(err, io.EOF),
 		errors.Is(err, syscall.ECONNRESET),

managedplugin/download_retry_test.go

@@ -11,6 +11,7 @@ import (
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 	"syscall"
 	"testing"
 	"time"
@@ -58,6 +59,7 @@ func TestIsRetryableDownloadError(t *testing.T) {
 		{name: "bad gateway 502", err: &httpStatusError{statusCode: http.StatusBadGateway}, want: true},
 		{name: "service unavailable 503", err: &httpStatusError{statusCode: http.StatusServiceUnavailable}, want: true},
 
+		{name: "stalled download", err: fmt.Errorf("%w: no data from host for 30s", errDownloadStalled), want: true},
 		{name: "context canceled", err: fmt.Errorf("get url: %w", context.Canceled), want: false},
 		{name: "context deadline exceeded", err: fmt.Errorf("get url: %w", context.DeadlineExceeded), want: false},
 		{name: "checksum mismatch is permanent", err: errors.New("checksum mismatch: expected abc, got def"), want: false},
@@ -208,6 +210,115 @@ func fastRetries(t *testing.T) {
 	})
 }
 
+// fastStall shrinks the stall watchdog so the stall tests do not wait 30s per
+// attempt.
+func fastStall(t *testing.T, d time.Duration) {
+	t.Helper()
+
+	prev := downloadStallTimeout
+	downloadStallTimeout = d
+	t.Cleanup(func() {
+		downloadStallTimeout = prev
+	})
+}
+
+// TestDownloadFileRetriesStalledBody reproduces the assets.cloudquery.io outage
+// mode: the server sends part of the body and then goes silent without closing
+// the connection. Without the watchdog the attempt hangs until a middlebox
+// kills the stream minutes later; with it the attempt dies quickly and the
+// retry succeeds.
+func TestDownloadFileRetriesStalledBody(t *testing.T) {
+	fastRetries(t)
+	fastStall(t, 150*time.Millisecond)
+
+	body := []byte("cloudquery-plugin-binary-payload")
+	release := make(chan struct{})
+
+	var attempts atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		if attempts.Add(1) == 1 {
+			w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body)))
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write(body[:5])
+			w.(http.Flusher).Flush()
+			<-release
+			return
+		}
+		_, _ = w.Write(body)
+	}))
+	t.Cleanup(server.Close)
+	t.Cleanup(func() { close(release) })
+
+	localPath := filepath.Join(t.TempDir(), "plugin.zip")
+	checksum, err := downloadFile(context.Background(), localPath, server.URL, DownloaderOptions{NoProgress: true})
+	require.NoError(t, err)
+	require.EqualValues(t, 2, attempts.Load())
+
+	written, err := os.ReadFile(localPath)
+	require.NoError(t, err)
+	require.Equal(t, body, written)
+	require.Equal(t, sha256Hex(body), checksum)
+}
+
+// TestDownloadFileRetriesStalledHeaders covers a stall before any response
+// arrives — the watchdog must cut the header wait too, not just the body copy.
+func TestDownloadFileRetriesStalledHeaders(t *testing.T) {
+	fastRetries(t)
+	fastStall(t, 150*time.Millisecond)
+
+	body := []byte("payload")
+	release := make(chan struct{})
+
+	var attempts atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		if attempts.Add(1) == 1 {
+			<-release
+			return
+		}
+		_, _ = w.Write(body)
+	}))
+	t.Cleanup(server.Close)
+	t.Cleanup(func() { close(release) })
+
+	localPath := filepath.Join(t.TempDir(), "plugin.zip")
+	checksum, err := downloadFile(context.Background(), localPath, server.URL, DownloaderOptions{NoProgress: true})
+	require.NoError(t, err)
+	require.EqualValues(t, 2, attempts.Load())
+	require.Equal(t, sha256Hex(body), checksum)
+}
+
+// TestDownloadFileCallerCancelIsNotRetried pins the boundary between the
+// watchdog's own cancellation (retryable) and the caller's (terminal).
+func TestDownloadFileCallerCancelIsNotRetried(t *testing.T) {
+	fastRetries(t)
+
+	release := make(chan struct{})
+
+	var attempts atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		attempts.Add(1)
+		w.Header().Set("Content-Length", "32")
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("early"))
+		w.(http.Flusher).Flush()
+		<-release
+	}))
+	t.Cleanup(server.Close)
+	t.Cleanup(func() { close(release) })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		cancel()
+	}()
+
+	localPath := filepath.Join(t.TempDir(), "plugin.zip")
+	_, err := downloadFile(ctx, localPath, server.URL, DownloaderOptions{NoProgress: true})
+	require.Error(t, err)
+	require.EqualValues(t, 1, attempts.Load(), "the caller's own cancellation must not be retried")
+	require.NotErrorIs(t, err, errDownloadStalled)
+}
+
 // TestDownloadFileRedactsSignedTokenFromTransportErrors covers the leak that
 // url.Error reopens: it prints its URL verbatim, so wrapping one puts the signed
 // token back into the message once per attempt.