server: detect download stalls before the first byte (#17259)

* server: detect download stalls before the first byte

* server: keep stall timeout out of download API

server/download.go

@@ -102,6 +102,8 @@ const (
 	maxDownloadPartSize int64 = 1000 * format.MegaByte
 )
 
+var downloadStallTimeout = 30 * time.Second
+
 func (p *blobDownloadPart) Name() string {
 	return strings.Join([]string{
 		p.blobDownload.Name, "partial", strconv.Itoa(p.N),
@@ -330,7 +332,11 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis
 
 func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w io.Writer, part *blobDownloadPart) error {
 	g, ctx := errgroup.WithContext(ctx)
+	attemptStarted := time.Now()
+	transferDone := make(chan struct{})
 	g.Go(func() error {
+		defer close(transferDone)
+
 		req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL.String(), nil)
 		if err != nil {
 			return err
@@ -359,19 +365,21 @@ func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w
 	})
 
 	g.Go(func() error {
-		ticker := time.NewTicker(time.Second)
+		ticker := time.NewTicker(min(time.Second, downloadStallTimeout/2))
+		defer ticker.Stop()
 		for {
 			select {
+			case <-transferDone:
+				return nil
 			case <-ticker.C:
-				if part.Completed.Load() >= part.Size {
-					return nil
-				}
-
 				part.lastUpdatedMu.Lock()
 				lastUpdated := part.lastUpdated
 				part.lastUpdatedMu.Unlock()
+				if lastUpdated.Before(attemptStarted) {
+					lastUpdated = attemptStarted
+				}
 
-				if !lastUpdated.IsZero() && time.Since(lastUpdated) > 30*time.Second {
+				if time.Since(lastUpdated) > downloadStallTimeout {
 					const msg = "%s part %d stalled; retrying. If this persists, press ctrl-c to exit, then 'ollama pull' to find a faster connection."
 					slog.Info(fmt.Sprintf(msg, b.Digest[7:19], part.N))
 					// reset last updated

server/download_test.go

@@ -0,0 +1,113 @@
+package server
+
+import (
+	"context"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+func BenchmarkDownloadChunkCompletion(b *testing.B) {
+	data := make([]byte, 1024*1024)
+	digest := fmt.Sprintf("sha256:%x", sha256.Sum256(data))
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Length", fmt.Sprint(len(data)))
+		w.WriteHeader(http.StatusPartialContent)
+		_, _ = w.Write(data)
+	}))
+	b.Cleanup(server.Close)
+
+	requestURL, err := url.Parse(server.URL)
+	if err != nil {
+		b.Fatal(err)
+	}
+	downloadPath := filepath.Join(b.TempDir(), "blob")
+
+	b.SetBytes(int64(len(data)))
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		download := &blobDownload{Name: downloadPath, Digest: digest}
+		part := &blobDownloadPart{Size: int64(len(data)), blobDownload: download}
+		if err := download.downloadChunk(b.Context(), requestURL, io.Discard, part); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func TestDownloadChunkReturnsWhenTransferCompletes(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Length", "1")
+		w.WriteHeader(http.StatusPartialContent)
+		_, _ = w.Write([]byte{0})
+	}))
+	t.Cleanup(server.Close)
+
+	requestURL, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	download := &blobDownload{
+		Name:   filepath.Join(t.TempDir(), "blob"),
+		Digest: "sha256:0000000000000000000000000000000000000000000000000000000000000000",
+	}
+	part := &blobDownloadPart{Size: 1, blobDownload: download}
+	ctx, cancel := context.WithTimeout(t.Context(), 250*time.Millisecond)
+	defer cancel()
+
+	if err := download.downloadChunk(ctx, requestURL, io.Discard, part); err != nil {
+		t.Fatalf("downloadChunk() error = %v, want nil", err)
+	}
+}
+
+func TestDownloadChunkDetectsStallBeforeFirstByte(t *testing.T) {
+	requestStarted := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		close(requestStarted)
+		w.Header().Set("Content-Length", "1")
+		w.WriteHeader(http.StatusPartialContent)
+		w.(http.Flusher).Flush()
+		<-r.Context().Done()
+	}))
+	t.Cleanup(server.Close)
+
+	requestURL, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	download := &blobDownload{Digest: "sha256:0000000000000000000000000000000000000000000000000000000000000000"}
+	part := &blobDownloadPart{Size: 1, blobDownload: download}
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
+	defer cancel()
+
+	originalStallTimeout := downloadStallTimeout
+	downloadStallTimeout = 50 * time.Millisecond
+	t.Cleanup(func() {
+		downloadStallTimeout = originalStallTimeout
+	})
+
+	started := time.Now()
+	err = download.downloadChunk(ctx, requestURL, io.Discard, part)
+	elapsed := time.Since(started)
+
+	select {
+	case <-requestStarted:
+	default:
+		t.Fatal("download request did not start")
+	}
+	if !errors.Is(err, errPartStalled) {
+		t.Fatalf("downloadChunk() error = %v after %v, want %v", err, elapsed, errPartStalled)
+	}
+	if elapsed >= 5*downloadStallTimeout {
+		t.Fatalf("downloadChunk() detected the stall after %v, want less than %v", elapsed, 5*downloadStallTimeout)
+	}
+}