From 4d1b53e6fb935fb93962cc4689da3ac9320ec87d Mon Sep 17 00:00:00 2001 From: Jhye Date: Wed, 22 Jul 2026 04:28:19 +1000 Subject: [PATCH] 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 | 20 ++++--- server/download_test.go | 113 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 6 deletions(-) create mode 100644 server/download_test.go diff --git a/server/download.go b/server/download.go index 0019fa13a..bdc03338b 100644 --- a/server/download.go +++ b/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 diff --git a/server/download_test.go b/server/download_test.go new file mode 100644 index 000000000..1fd8e3b80 --- /dev/null +++ b/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) + } +}