From 1e1b34dada7db6aa7c98db1d13f0dd9350b51f63 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Fri, 8 May 2026 14:25:30 -0700 Subject: [PATCH] mlx: refined model push behavior (#15431) * mlx: refined model push behavior Refine the algorithm for parallel push of safetensors based models to get better reliability and throughput. * review comments, hardening, and performance tuning for slow links * review comments --- cmd/cmd.go | 3 +- envconfig/config.go | 53 +- server/images.go | 38 +- x/transfer/download.go | 90 ++- x/transfer/transfer.go | 76 +-- x/transfer/transfer_test.go | 1145 ++++++++++++++++++++++++++++++++--- x/transfer/upload.go | 626 +++++++++++++++---- 7 files changed, 1754 insertions(+), 277 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 927f0a75..e238b5bd 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -2009,7 +2009,7 @@ func appendEnvDocs(cmd *cobra.Command, envs []envconfig.EnvVar) { Environment Variables: ` for _, e := range envs { - envUsage += fmt.Sprintf(" %-24s %s\n", e.Name, e.Description) + envUsage += fmt.Sprintf(" %-27s %s\n", e.Name, e.Description) } cmd.SetUsageTemplate(cmd.UsageTemplate() + envUsage) @@ -2473,6 +2473,7 @@ func NewCLI() *cobra.Command { envVars["OLLAMA_CONTEXT_LENGTH"], envVars["OLLAMA_KEEP_ALIVE"], envVars["OLLAMA_MAX_LOADED_MODELS"], + envVars["OLLAMA_MAX_TRANSFER_STREAMS"], envVars["OLLAMA_MAX_QUEUE"], envVars["OLLAMA_MODELS"], envVars["OLLAMA_NUM_PARALLEL"], diff --git a/envconfig/config.go b/envconfig/config.go index f4fe078a..f4bd4aa6 100644 --- a/envconfig/config.go +++ b/envconfig/config.go @@ -277,6 +277,12 @@ var ( MaxRunners = Uint("OLLAMA_MAX_LOADED_MODELS", 0) // MaxQueue sets the maximum number of queued requests. MaxQueue can be configured via the OLLAMA_MAX_QUEUE environment variable. MaxQueue = Uint("OLLAMA_MAX_QUEUE", 512) + // MaxTransferStreams caps the number of simultaneous body-bearing + // transfers during safetensors model pulls/pushes, keeping slower + // networks from being saturated. Tune higher for fast networks. Has + // no effect on GGUF transfers, which use the legacy upload/download + // paths. + MaxTransferStreams = Uint("OLLAMA_MAX_TRANSFER_STREAMS", 4) ) func Uint64(key string, defaultValue uint64) func() uint64 { @@ -304,29 +310,30 @@ type EnvVar struct { func AsMap() map[string]EnvVar { ret := map[string]EnvVar{ - "OLLAMA_DEBUG": {"OLLAMA_DEBUG", LogLevel(), "Show additional debug information (e.g. OLLAMA_DEBUG=1)"}, - "OLLAMA_DEBUG_LOG_REQUESTS": {"OLLAMA_DEBUG_LOG_REQUESTS", DebugLogRequests(), "Log inference request bodies and replay curl commands to a temp directory"}, - "OLLAMA_FLASH_ATTENTION": {"OLLAMA_FLASH_ATTENTION", FlashAttention(false), "Enabled flash attention"}, - "OLLAMA_KV_CACHE_TYPE": {"OLLAMA_KV_CACHE_TYPE", KvCacheType(), "Quantization type for the K/V cache (default: f16)"}, - "OLLAMA_GPU_OVERHEAD": {"OLLAMA_GPU_OVERHEAD", GpuOverhead(), "Reserve a portion of VRAM per GPU (bytes)"}, - "OLLAMA_HOST": {"OLLAMA_HOST", Host(), "IP Address for the ollama server (default 127.0.0.1:11434)"}, - "OLLAMA_KEEP_ALIVE": {"OLLAMA_KEEP_ALIVE", KeepAlive(), "The duration that models stay loaded in memory (default \"5m\")"}, - "OLLAMA_LLM_LIBRARY": {"OLLAMA_LLM_LIBRARY", LLMLibrary(), "Set LLM library to bypass autodetection"}, - "OLLAMA_LOAD_TIMEOUT": {"OLLAMA_LOAD_TIMEOUT", LoadTimeout(), "How long to allow model loads to stall before giving up (default \"5m\")"}, - "OLLAMA_MAX_LOADED_MODELS": {"OLLAMA_MAX_LOADED_MODELS", MaxRunners(), "Maximum number of loaded models per GPU"}, - "OLLAMA_MAX_QUEUE": {"OLLAMA_MAX_QUEUE", MaxQueue(), "Maximum number of queued requests"}, - "OLLAMA_MODELS": {"OLLAMA_MODELS", Models(), "The path to the models directory"}, - "OLLAMA_NO_CLOUD": {"OLLAMA_NO_CLOUD", NoCloud(), "Disable Ollama cloud features (remote inference and web search)"}, - "OLLAMA_NOHISTORY": {"OLLAMA_NOHISTORY", NoHistory(), "Do not preserve readline history"}, - "OLLAMA_NOPRUNE": {"OLLAMA_NOPRUNE", NoPrune(), "Do not prune model blobs on startup"}, - "OLLAMA_NUM_PARALLEL": {"OLLAMA_NUM_PARALLEL", NumParallel(), "Maximum number of parallel requests"}, - "OLLAMA_ORIGINS": {"OLLAMA_ORIGINS", AllowedOrigins(), "A comma separated list of allowed origins"}, - "OLLAMA_SCHED_SPREAD": {"OLLAMA_SCHED_SPREAD", SchedSpread(), "Always schedule model across all GPUs"}, - "OLLAMA_MULTIUSER_CACHE": {"OLLAMA_MULTIUSER_CACHE", MultiUserCache(), "Optimize prompt caching for multi-user scenarios"}, - "OLLAMA_CONTEXT_LENGTH": {"OLLAMA_CONTEXT_LENGTH", ContextLength(), "Context length to use unless otherwise specified (default: 4k/32k/256k based on VRAM)"}, - "OLLAMA_EDITOR": {"OLLAMA_EDITOR", Editor(), "Path to editor for interactive prompt editing (Ctrl+G)"}, - "OLLAMA_NEW_ENGINE": {"OLLAMA_NEW_ENGINE", NewEngine(), "Enable the new Ollama engine"}, - "OLLAMA_REMOTES": {"OLLAMA_REMOTES", Remotes(), "Allowed hosts for remote models (default \"ollama.com\")"}, + "OLLAMA_DEBUG": {"OLLAMA_DEBUG", LogLevel(), "Show additional debug information (e.g. OLLAMA_DEBUG=1)"}, + "OLLAMA_DEBUG_LOG_REQUESTS": {"OLLAMA_DEBUG_LOG_REQUESTS", DebugLogRequests(), "Log inference request bodies and replay curl commands to a temp directory"}, + "OLLAMA_FLASH_ATTENTION": {"OLLAMA_FLASH_ATTENTION", FlashAttention(false), "Enabled flash attention"}, + "OLLAMA_KV_CACHE_TYPE": {"OLLAMA_KV_CACHE_TYPE", KvCacheType(), "Quantization type for the K/V cache (default: f16)"}, + "OLLAMA_GPU_OVERHEAD": {"OLLAMA_GPU_OVERHEAD", GpuOverhead(), "Reserve a portion of VRAM per GPU (bytes)"}, + "OLLAMA_HOST": {"OLLAMA_HOST", Host(), "IP Address for the ollama server (default 127.0.0.1:11434)"}, + "OLLAMA_KEEP_ALIVE": {"OLLAMA_KEEP_ALIVE", KeepAlive(), "The duration that models stay loaded in memory (default \"5m\")"}, + "OLLAMA_LLM_LIBRARY": {"OLLAMA_LLM_LIBRARY", LLMLibrary(), "Set LLM library to bypass autodetection"}, + "OLLAMA_LOAD_TIMEOUT": {"OLLAMA_LOAD_TIMEOUT", LoadTimeout(), "How long to allow model loads to stall before giving up (default \"5m\")"}, + "OLLAMA_MAX_LOADED_MODELS": {"OLLAMA_MAX_LOADED_MODELS", MaxRunners(), "Maximum number of loaded models per GPU"}, + "OLLAMA_MAX_TRANSFER_STREAMS": {"OLLAMA_MAX_TRANSFER_STREAMS", MaxTransferStreams(), "Maximum parallel transfer streams for safetensors model pulls/pushes (default 4)"}, + "OLLAMA_MAX_QUEUE": {"OLLAMA_MAX_QUEUE", MaxQueue(), "Maximum number of queued requests"}, + "OLLAMA_MODELS": {"OLLAMA_MODELS", Models(), "The path to the models directory"}, + "OLLAMA_NO_CLOUD": {"OLLAMA_NO_CLOUD", NoCloud(), "Disable Ollama cloud features (remote inference and web search)"}, + "OLLAMA_NOHISTORY": {"OLLAMA_NOHISTORY", NoHistory(), "Do not preserve readline history"}, + "OLLAMA_NOPRUNE": {"OLLAMA_NOPRUNE", NoPrune(), "Do not prune model blobs on startup"}, + "OLLAMA_NUM_PARALLEL": {"OLLAMA_NUM_PARALLEL", NumParallel(), "Maximum number of parallel requests"}, + "OLLAMA_ORIGINS": {"OLLAMA_ORIGINS", AllowedOrigins(), "A comma separated list of allowed origins"}, + "OLLAMA_SCHED_SPREAD": {"OLLAMA_SCHED_SPREAD", SchedSpread(), "Always schedule model across all GPUs"}, + "OLLAMA_MULTIUSER_CACHE": {"OLLAMA_MULTIUSER_CACHE", MultiUserCache(), "Optimize prompt caching for multi-user scenarios"}, + "OLLAMA_CONTEXT_LENGTH": {"OLLAMA_CONTEXT_LENGTH", ContextLength(), "Context length to use unless otherwise specified (default: 4k/32k/256k based on VRAM)"}, + "OLLAMA_EDITOR": {"OLLAMA_EDITOR", Editor(), "Path to editor for interactive prompt editing (Ctrl+G)"}, + "OLLAMA_NEW_ENGINE": {"OLLAMA_NEW_ENGINE", NewEngine(), "Enable the new Ollama engine"}, + "OLLAMA_REMOTES": {"OLLAMA_REMOTES", Remotes(), "Allowed hosts for remote models (default \"ollama.com\")"}, // Informational "HTTP_PROXY": {"HTTP_PROXY", String("HTTP_PROXY")(), "HTTP proxy"}, diff --git a/server/images.go b/server/images.go index 0458c24a..df2dc7aa 100644 --- a/server/images.go +++ b/server/images.go @@ -761,14 +761,15 @@ func pullWithTransfer(ctx context.Context, n model.Name, layers []manifest.Layer } if err := transfer.Download(ctx, transfer.DownloadOptions{ - Blobs: blobs, - BaseURL: baseURL, - DestDir: destDir, - Repository: n.DisplayNamespaceModel(), - Progress: progress, - Token: regOpts.Token, - GetToken: getToken, - Logger: slog.Default(), + Blobs: blobs, + BaseURL: baseURL, + DestDir: destDir, + Repository: n.DisplayNamespaceModel(), + BodyConcurrency: max(1, int(envconfig.MaxTransferStreams())), + Progress: progress, + Token: regOpts.Token, + GetToken: getToken, + Logger: slog.Default(), }); err != nil { return err } @@ -837,16 +838,17 @@ func pushWithTransfer(ctx context.Context, n model.Name, layers []manifest.Layer } return transfer.Upload(ctx, transfer.UploadOptions{ - Blobs: blobs, - BaseURL: baseURL, - SrcDir: srcDir, - Progress: progress, - Token: regOpts.Token, - GetToken: getToken, - Logger: slog.Default(), - Manifest: manifestJSON, - ManifestRef: n.Tag, - Repository: n.DisplayNamespaceModel(), + Blobs: blobs, + BaseURL: baseURL, + SrcDir: srcDir, + BodyConcurrency: max(1, int(envconfig.MaxTransferStreams())), + Progress: progress, + Token: regOpts.Token, + GetToken: getToken, + Logger: slog.Default(), + Manifest: manifestJSON, + ManifestRef: n.Tag, + Repository: n.DisplayNamespaceModel(), }) } diff --git a/x/transfer/download.go b/x/transfer/download.go index 64a6f889..f2a21d04 100644 --- a/x/transfer/download.go +++ b/x/transfer/download.go @@ -31,13 +31,61 @@ type downloader struct { baseURL string destDir string repository string // Repository path for blob URLs (e.g., "library/model") - token *string + tokenMu sync.RWMutex + token string getToken func(context.Context, AuthChallenge) (string, error) userAgent string stallTimeout time.Duration progress *progressTracker speeds *speedTracker logger *slog.Logger + // bodySem caps the number of simultaneous body-bearing transfers so a + // modest home downlink isn't saturated. Always set by download(); nil + // only when tests build downloader directly (in which case holdBody is + // a no-op). + bodySem *semaphore.Weighted +} + +// authToken returns the current bearer token. Safe to call concurrently with +// refreshToken. +func (d *downloader) authToken() string { + d.tokenMu.RLock() + defer d.tokenMu.RUnlock() + return d.token +} + +// refreshToken coalesces token fetches so concurrent 401s don't all hit the +// auth server. prev is the token the caller used in the request that got +// rejected: if the stored token has already moved past prev, another +// goroutine has refreshed and we just observe its result; otherwise the +// caller holds the lock and performs the fetch. +func (d *downloader) refreshToken(ctx context.Context, ch AuthChallenge, prev string) error { + d.tokenMu.Lock() + defer d.tokenMu.Unlock() + if d.token != prev { + return nil + } + if d.getToken == nil { + return errors.New("no token refresh callback") + } + t, err := d.getToken(ctx, ch) + if err != nil { + return err + } + d.token = t + return nil +} + +// holdBody acquires a body-transfer slot. The returned release must be +// called exactly once after the body-bearing request completes (defer is fine). +func (d *downloader) holdBody(ctx context.Context) (func(), error) { + if d.bodySem == nil { + return func() {}, nil + } + if err := d.bodySem.Acquire(ctx, 1); err != nil { + return nil, err + } + return func() { d.bodySem.Release(1) }, nil } func download(ctx context.Context, opts DownloadOptions) error { @@ -68,7 +116,6 @@ func download(ctx context.Context, opts DownloadOptions) error { return nil } - token := opts.Token progress := newProgressTracker(total, opts.Progress) progress.add(alreadyCompleted) // Report already-downloaded bytes upfront @@ -77,7 +124,7 @@ func download(ctx context.Context, opts DownloadOptions) error { baseURL: opts.BaseURL, destDir: opts.DestDir, repository: cmp.Or(opts.Repository, "library/_"), - token: &token, + token: opts.Token, getToken: opts.GetToken, userAgent: cmp.Or(opts.UserAgent, defaultUserAgent), stallTimeout: cmp.Or(opts.StallTimeout, defaultStallTimeout), @@ -85,10 +132,13 @@ func download(ctx context.Context, opts DownloadOptions) error { speeds: &speedTracker{}, logger: opts.Logger, } + // 0 or negative serializes; never unbounded. + d.bodySem = semaphore.NewWeighted(int64(max(1, opts.BodyConcurrency))) concurrency := cmp.Or(opts.Concurrency, DefaultDownloadConcurrency) sem := semaphore.NewWeighted(int64(concurrency)) + start := time.Now() g, ctx := errgroup.WithContext(ctx) for _, blob := range blobs { g.Go(func() error { @@ -99,7 +149,18 @@ func download(ctx context.Context, opts DownloadOptions) error { return d.download(ctx, blob) }) } - return g.Wait() + err := g.Wait() + elapsed := time.Since(start) + done := d.progress.completed.Load() - alreadyCompleted + mbps := float64(done) / 1e6 / max(0.001, elapsed.Seconds()) + slog.Debug("download summary", + "blobs", len(blobs), + "bytes", done, + "duration", elapsed.Round(time.Millisecond), + "mb_per_sec", fmt.Sprintf("%.1f", mbps), + "max_transfers", max(1, opts.BodyConcurrency), + ) + return err } func (d *downloader) download(ctx context.Context, blob Blob) error { @@ -158,6 +219,14 @@ func (d *downloader) downloadOnce(ctx context.Context, blob Blob) (int64, error) d.logger.Debug("downloading blob", "digest", blob.Digest, "size", blob.Size) } + // Hold a body slot for the duration of the GET — released when the body + // has been read and the response closed. + release, err := d.holdBody(ctx) + if err != nil { + return 0, err + } + defer release() + baseURL, _ := url.Parse(d.baseURL) u, err := d.resolve(ctx, fmt.Sprintf("%s/v2/%s/blobs/%s", d.baseURL, d.repository, blob.Digest)) if err != nil { @@ -183,8 +252,10 @@ func (d *downloader) downloadOnce(ctx context.Context, blob Blob) (int64, error) req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) req.Header.Set("User-Agent", d.userAgent) // Add auth only for same-host (not CDN) - if u.Host == baseURL.Host && *d.token != "" { - req.Header.Set("Authorization", "Bearer "+*d.token) + if u.Host == baseURL.Host { + if t := d.authToken(); t != "" { + req.Header.Set("Authorization", "Bearer "+t) + } } if existingSize > 0 { req.Header.Set("Range", fmt.Sprintf("bytes=%d-", existingSize)) @@ -331,8 +402,9 @@ func (d *downloader) resolve(ctx context.Context, rawURL string) (*url.URL, erro for range 10 { req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) req.Header.Set("User-Agent", d.userAgent) - if *d.token != "" { - req.Header.Set("Authorization", "Bearer "+*d.token) + prev := d.authToken() + if prev != "" { + req.Header.Set("Authorization", "Bearer "+prev) } resp, err := d.client.Do(req) @@ -351,7 +423,7 @@ func (d *downloader) resolve(ctx context.Context, rawURL string) (*url.URL, erro return nil, fmt.Errorf("unauthorized") } ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate")) - if *d.token, err = d.getToken(ctx, ch); err != nil { + if err := d.refreshToken(ctx, ch, prev); err != nil { return nil, err } case http.StatusTemporaryRedirect, http.StatusFound, http.StatusMovedPermanently: diff --git a/x/transfer/transfer.go b/x/transfer/transfer.go index 33e6d589..4c987acc 100644 --- a/x/transfer/transfer.go +++ b/x/transfer/transfer.go @@ -7,20 +7,24 @@ // TODO (jmorganca): Integrate into server/download.go and server/upload.go when stable. // // Design Philosophy: -// This package is intentionally simpler than the main server's download/upload code. -// Key simplifications for many-small-blob workloads: +// This package is intentionally simpler than the main server's download/upload +// code. Key simplifications for many-small-blob workloads: // -// - Whole-blob transfers: No part-based chunking. Each blob downloads/uploads as one unit. -// - Resume for large blobs: Blobs >= 64MB preserve partial .tmp files on failure -// and use HTTP Range requests on retry. Small blobs restart from scratch. -// - Inline hashing: SHA256 computed during streaming, not asynchronously after parts complete. -// - Stall and speed detection: Cancels on no data (stall) or speed below 10% of median. +// - Whole-blob downloads: Each blob downloads as one unit, with HTTP Range +// resume for blobs >= 64MB on retry; small blobs restart from scratch. +// - Whole-blob uploads by default: A single PUT per blob. When the server +// returns a direct-upload URL the body goes straight to the storage +// backend; otherwise the body goes to the registry in one shot. +// - Multi-part upload fallback: If the server requires it, blobs are split +// into parts and sent via PATCH with a finalize PUT carrying a composite +// etag. This is a server-side compatibility path, not the fast path. +// - Inline hashing: digests computed during streaming. +// - Stall and speed detection (downloads): cancels on no data (stall) or +// speed below 10% of median. // -// For large models (multi-GB), use the server's download/upload code which has: -// - Part-based transfers with 64MB chunks -// - Resumable downloads with JSON state files -// - Async streamHasher that hashes from OS page cache as parts complete -// - Speed tracking with rolling median to detect and restart slow parts +// For large models (multi-GB), use the server's download/upload code which +// has resumable downloads with JSON state files, async hashing from OS page +// cache, and per-part speed tracking with rolling median. package transfer import ( @@ -54,32 +58,34 @@ type Blob struct { // DownloadOptions configures a parallel download operation. type DownloadOptions struct { - Blobs []Blob // Blobs to download - BaseURL string // Registry base URL - DestDir string // Destination directory for blobs - Repository string // Repository path for blob URLs (e.g., "library/model") - Concurrency int // Max parallel downloads (default 64) - Progress func(completed, total int64) // Progress callback (optional) - Client *http.Client // HTTP client (optional, uses default) - Token string // Auth token (optional) - GetToken func(ctx context.Context, challenge AuthChallenge) (string, error) // Token refresh callback - Logger *slog.Logger // Optional structured logger - UserAgent string // User-Agent header (optional, has default) - StallTimeout time.Duration // Timeout for stall detection (default 10s) + Blobs []Blob // Blobs to download + BaseURL string // Registry base URL + DestDir string // Destination directory for blobs + Repository string // Repository path for blob URLs (e.g., "library/model") + Concurrency int // Max parallel downloads (default DefaultDownloadConcurrency) + BodyConcurrency int // Max simultaneous body-bearing transfers; 0 or negative serializes (capacity 1) + Progress func(completed, total int64) // Progress callback (optional) + Client *http.Client // HTTP client (optional, uses default) + Token string // Auth token (optional) + GetToken func(ctx context.Context, challenge AuthChallenge) (string, error) // Token refresh callback + Logger *slog.Logger // Optional structured logger + UserAgent string // User-Agent header (optional, has default) + StallTimeout time.Duration // Timeout for stall detection (default 10s) } // UploadOptions configures a parallel upload operation. type UploadOptions struct { - Blobs []Blob // Blobs to upload - BaseURL string // Registry base URL - SrcDir string // Source directory containing blobs - Concurrency int // Max parallel uploads (default 32) - Progress func(completed, total int64) // Progress callback (optional) - Client *http.Client // HTTP client (optional, uses default) - Token string // Auth token (optional) - GetToken func(ctx context.Context, challenge AuthChallenge) (string, error) // Token refresh callback - Logger *slog.Logger // Optional structured logger - UserAgent string // User-Agent header (optional, has default) + Blobs []Blob // Blobs to upload + BaseURL string // Registry base URL + SrcDir string // Source directory containing blobs + Concurrency int // Max parallel uploads (default DefaultUploadConcurrency) + BodyConcurrency int // Max simultaneous body-bearing transfers; 0 or negative serializes (capacity 1) + Progress func(completed, total int64) // Progress callback (optional) + Client *http.Client // HTTP client (optional, uses default) + Token string // Auth token (optional) + GetToken func(ctx context.Context, challenge AuthChallenge) (string, error) // Token refresh callback + Logger *slog.Logger // Optional structured logger + UserAgent string // User-Agent header (optional, has default) // Manifest fields (optional) - if set, manifest is pushed after all blobs complete Manifest []byte // Raw manifest JSON to push @@ -97,7 +103,7 @@ type AuthChallenge struct { // Default concurrency limits and settings const ( DefaultDownloadConcurrency = 64 - DefaultUploadConcurrency = 32 + DefaultUploadConcurrency = 64 maxRetries = 6 defaultUserAgent = "ollama-transfer/1.0" diff --git a/x/transfer/transfer_test.go b/x/transfer/transfer_test.go index f28ecd15..206f144b 100644 --- a/x/transfer/transfer_test.go +++ b/x/transfer/transfer_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "crypto/sha256" + "encoding/base64" "errors" "fmt" "io" @@ -18,6 +19,45 @@ import ( "time" ) +// chunkedSession tracks accumulated PATCH body bytes for an upload session. +// Tests that mock the registry use it to handle the GGUF-style POST → PATCH → +// PUT-finalize flow without each test reimplementing the bookkeeping. +type chunkedSession struct { + mu sync.Mutex + sessions map[string]*bytes.Buffer +} + +func newChunkedSession() *chunkedSession { + return &chunkedSession{sessions: make(map[string]*bytes.Buffer)} +} + +// recordPatch reads the request body into the session buffer and writes a +// 202 Accepted response with Docker-Upload-Location pointing at the same +// session URL. Use this from a mock handler's PATCH branch. +func (c *chunkedSession) recordPatch(w http.ResponseWriter, r *http.Request) { + c.mu.Lock() + buf, ok := c.sessions[r.URL.Path] + if !ok { + buf = &bytes.Buffer{} + c.sessions[r.URL.Path] = buf + } + c.mu.Unlock() + io.Copy(buf, r.Body) + w.Header().Set("Docker-Upload-Location", r.URL.Path) + w.WriteHeader(http.StatusAccepted) +} + +// finalize returns the bytes accumulated for the given session URL path. +// Use it from the PUT-finalize branch of a mock handler. +func (c *chunkedSession) finalize(sessionPath string) []byte { + c.mu.Lock() + defer c.mu.Unlock() + if buf, ok := c.sessions[sessionPath]; ok { + return buf.Bytes() + } + return nil +} + // createTestBlob creates a blob with deterministic content and returns its digest func createTestBlob(t *testing.T, dir string, size int) (Blob, []byte) { t.Helper() @@ -388,26 +428,26 @@ func TestUpload(t *testing.T) { blob2, _ := createTestBlob(t, clientDir, 2048) var uploadedBlobs sync.Map - uploadID := 0 + var uploadID atomic.Int32 + session := newChunkedSession() var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case r.Method == http.MethodHead: - // Blob doesn't exist http.NotFound(w, r) case r.Method == http.MethodPost && r.URL.Path == "/v2/library/_/blobs/uploads/": - // Initiate upload - uploadID++ - w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, uploadID)) + id := uploadID.Add(1) + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, id)) w.WriteHeader(http.StatusAccepted) + case r.Method == http.MethodPatch: + session.recordPatch(w, r) + case r.Method == http.MethodPut: - // Complete upload digest := r.URL.Query().Get("digest") - data, _ := io.ReadAll(r.Body) - uploadedBlobs.Store(digest, data) + uploadedBlobs.Store(digest, session.finalize(r.URL.Path)) w.WriteHeader(http.StatusCreated) default: @@ -451,35 +491,49 @@ func TestUploadWithRedirect(t *testing.T) { var uploadedBlobs sync.Map var cdnCalled atomic.Bool - // CDN server (redirect target) + // CDN server (PATCH redirect target). PATCH is redirected to a PUT here, + // matching production: server issues 307 + Location to a presigned CDN URL, + // and the client re-uploads the part body via PUT. cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { cdnCalled.Store(true) if r.Method == http.MethodPut { - digest := r.URL.Query().Get("digest") data, _ := io.ReadAll(r.Body) - uploadedBlobs.Store(digest, data) + // Stash the body keyed by the path so the main server can pick it + // up at finalize time. + uploadedBlobs.Store(r.URL.Path, data) w.WriteHeader(http.StatusCreated) } })) defer cdn.Close() var serverURL string - uploadID := 0 + var uploadID atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch { - case r.Method == http.MethodHead: + switch r.Method { + case http.MethodHead: http.NotFound(w, r) - case r.Method == http.MethodPost && r.URL.Path == "/v2/library/_/blobs/uploads/": - uploadID++ - w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, uploadID)) + case http.MethodPost: + id := uploadID.Add(1) + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, id)) w.WriteHeader(http.StatusAccepted) - case r.Method == http.MethodPut: - // Redirect to CDN - cdnURL := cdn.URL + r.URL.Path + "?" + r.URL.RawQuery + case http.MethodPatch: + // Redirect PATCH body to CDN, mirroring server behavior + cdnURL := cdn.URL + r.URL.Path + w.Header().Set("Docker-Upload-Location", r.URL.Path) http.Redirect(w, r, cdnURL, http.StatusTemporaryRedirect) + case http.MethodPut: + // Finalize: copy body the CDN received under this session path + // to the uploadedBlobs map keyed by digest. + digest := r.URL.Query().Get("digest") + if v, ok := uploadedBlobs.Load(r.URL.Path); ok { + uploadedBlobs.Store(digest, v) + uploadedBlobs.Delete(r.URL.Path) + } + w.WriteHeader(http.StatusCreated) + default: http.NotFound(w, r) } @@ -511,7 +565,8 @@ func TestUploadWithAuth(t *testing.T) { var uploadedBlobs sync.Map var authCalled atomic.Bool - uploadID := 0 + var uploadID atomic.Int32 + session := newChunkedSession() var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -523,19 +578,21 @@ func TestUploadWithAuth(t *testing.T) { return } - switch { - case r.Method == http.MethodHead: + switch r.Method { + case http.MethodHead: http.NotFound(w, r) - case r.Method == http.MethodPost && r.URL.Path == "/v2/library/_/blobs/uploads/": - uploadID++ - w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, uploadID)) + case http.MethodPost: + id := uploadID.Add(1) + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, id)) w.WriteHeader(http.StatusAccepted) - case r.Method == http.MethodPut: + case http.MethodPatch: + session.recordPatch(w, r) + + case http.MethodPut: digest := r.URL.Query().Get("digest") - data, _ := io.ReadAll(r.Body) - uploadedBlobs.Store(digest, data) + uploadedBlobs.Store(digest, session.finalize(r.URL.Path)) w.WriteHeader(http.StatusCreated) default: @@ -620,23 +677,27 @@ func TestUploadWithCustomRepository(t *testing.T) { var headPath, postPath string var mu sync.Mutex + session := newChunkedSession() var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - mu.Lock() switch r.Method { case http.MethodHead: + mu.Lock() headPath = r.URL.Path + mu.Unlock() w.WriteHeader(http.StatusNotFound) // Blob doesn't exist case http.MethodPost: + mu.Lock() postPath = r.URL.Path + mu.Unlock() w.Header().Set("Location", fmt.Sprintf("%s/v2/myorg/mymodel/blobs/uploads/1", serverURL)) w.WriteHeader(http.StatusAccepted) + case http.MethodPatch: + session.recordPatch(w, r) case http.MethodPut: - io.Copy(io.Discard, r.Body) w.WriteHeader(http.StatusCreated) } - mu.Unlock() })) defer server.Close() serverURL = server.URL @@ -798,6 +859,82 @@ func verifyBlob(t *testing.T, dir string, blob Blob, expected []byte) { // ==================== Parallelism Tests ==================== +func TestDownloadParallelism(t *testing.T) { + // Create many blobs to test parallelism + serverDir := t.TempDir() + numBlobs := 10 + blobs := make([]Blob, numBlobs) + blobData := make([][]byte, numBlobs) + + for i := range numBlobs { + blobs[i], blobData[i] = createTestBlob(t, serverDir, 1024+i*100) + } + + var activeRequests atomic.Int32 + var maxConcurrent atomic.Int32 + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + current := activeRequests.Add(1) + defer activeRequests.Add(-1) + + // Track max concurrent requests + for { + old := maxConcurrent.Load() + if current <= old || maxConcurrent.CompareAndSwap(old, current) { + break + } + } + + // Simulate network latency to ensure parallelism is visible + time.Sleep(50 * time.Millisecond) + + digest := filepath.Base(r.URL.Path) + path := filepath.Join(serverDir, digestToPath(digest)) + data, err := os.ReadFile(path) + if err != nil { + http.NotFound(w, r) + return + } + w.WriteHeader(http.StatusOK) + w.Write(data) + })) + defer server.Close() + + clientDir := t.TempDir() + + start := time.Now() + err := Download(context.Background(), DownloadOptions{ + Blobs: blobs, + BaseURL: server.URL, + DestDir: clientDir, + Concurrency: 4, + BodyConcurrency: 4, + }) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Download failed: %v", err) + } + + // Verify all blobs downloaded + for i, blob := range blobs { + verifyBlob(t, clientDir, blob, blobData[i]) + } + + // Verify parallelism was used + if maxConcurrent.Load() < 2 { + t.Errorf("Max concurrent requests was %d, expected at least 2 for parallelism", maxConcurrent.Load()) + } + + // With 10 blobs at 50ms each, sequential would take ~500ms + // Parallel with 4 workers should take ~150ms (relax to 1s for CI variance) + if elapsed > time.Second { + t.Errorf("Downloads took %v, expected faster with parallelism", elapsed) + } + + t.Logf("Downloaded %d blobs in %v with max %d concurrent requests", numBlobs, elapsed, maxConcurrent.Load()) +} + func TestUploadParallelism(t *testing.T) { clientDir := t.TempDir() numBlobs := 10 @@ -811,6 +948,7 @@ func TestUploadParallelism(t *testing.T) { var maxConcurrent atomic.Int32 var uploadedBlobs sync.Map var uploadID atomic.Int32 + session := newChunkedSession() var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -825,20 +963,22 @@ func TestUploadParallelism(t *testing.T) { } } - switch { - case r.Method == http.MethodHead: + switch r.Method { + case http.MethodHead: http.NotFound(w, r) - case r.Method == http.MethodPost: + case http.MethodPost: id := uploadID.Add(1) w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", serverURL, id)) w.WriteHeader(http.StatusAccepted) - case r.Method == http.MethodPut: - time.Sleep(50 * time.Millisecond) // Simulate upload time + case http.MethodPatch: + time.Sleep(50 * time.Millisecond) // simulate upload time on body chunk + session.recordPatch(w, r) + + case http.MethodPut: digest := r.URL.Query().Get("digest") - data, _ := io.ReadAll(r.Body) - uploadedBlobs.Store(digest, data) + uploadedBlobs.Store(digest, session.finalize(r.URL.Path)) w.WriteHeader(http.StatusCreated) default: @@ -1020,16 +1160,16 @@ func TestUploadCancellation(t *testing.T) { var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch { - case r.Method == http.MethodHead: + switch r.Method { + case http.MethodHead: http.NotFound(w, r) - case r.Method == http.MethodPost: + case http.MethodPost: w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) w.WriteHeader(http.StatusAccepted) - case r.Method == http.MethodPut: - // Read slowly + case http.MethodPatch: + // Read slowly so the cancellation has time to interrupt the body upload. buf := make([]byte, 1024) for { _, err := r.Body.Read(buf) @@ -1038,6 +1178,10 @@ func TestUploadCancellation(t *testing.T) { } time.Sleep(5 * time.Millisecond) } + w.Header().Set("Docker-Upload-Location", r.URL.Path) + w.WriteHeader(http.StatusAccepted) + + case http.MethodPut: w.WriteHeader(http.StatusCreated) } })) @@ -1173,29 +1317,33 @@ func TestUploadRetryOnFailure(t *testing.T) { clientDir := t.TempDir() blob, _ := createTestBlob(t, clientDir, 1024) - var putCount atomic.Int32 + var patchCount atomic.Int32 var uploadedBlobs sync.Map + session := newChunkedSession() var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch { - case r.Method == http.MethodHead: + switch r.Method { + case http.MethodHead: http.NotFound(w, r) - case r.Method == http.MethodPost: + case http.MethodPost: w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) w.WriteHeader(http.StatusAccepted) - case r.Method == http.MethodPut: - count := putCount.Add(1) + case http.MethodPatch: + count := patchCount.Add(1) if count < 3 { - // Fail first 2 attempts + // Fail first 2 PATCH attempts to exercise the retry path + io.Copy(io.Discard, r.Body) http.Error(w, "server error", http.StatusInternalServerError) return } + session.recordPatch(w, r) + + case http.MethodPut: digest := r.URL.Query().Get("digest") - data, _ := io.ReadAll(r.Body) - uploadedBlobs.Store(digest, data) + uploadedBlobs.Store(digest, session.finalize(r.URL.Path)) w.WriteHeader(http.StatusCreated) } })) @@ -1215,8 +1363,8 @@ func TestUploadRetryOnFailure(t *testing.T) { t.Error("Blob not uploaded after retry") } - if putCount.Load() < 3 { - t.Errorf("Expected at least 3 PUT attempts, got %d", putCount.Load()) + if patchCount.Load() < 3 { + t.Errorf("Expected at least 3 PATCH attempts, got %d", patchCount.Load()) } } @@ -1235,29 +1383,34 @@ func TestProgressRollback(t *testing.T) { t.Fatal(err) } - var putCount atomic.Int32 + var patchCount atomic.Int32 var progressValues []int64 var mu sync.Mutex + session := newChunkedSession() var serverURL string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch { - case r.Method == http.MethodHead: + switch r.Method { + case http.MethodHead: http.NotFound(w, r) - case r.Method == http.MethodPost: + case http.MethodPost: w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) w.WriteHeader(http.StatusAccepted) - case r.Method == http.MethodPut: - // Read some data before failing - io.CopyN(io.Discard, r.Body, 10) - count := putCount.Add(1) + case http.MethodPatch: + // Read some bytes (so the client reports progress) before failing, + // to exercise the rollback-on-retry path. + count := patchCount.Add(1) if count < 3 { + io.CopyN(io.Discard, r.Body, 10) + io.Copy(io.Discard, r.Body) http.Error(w, "server error", http.StatusInternalServerError) return } - io.Copy(io.Discard, r.Body) + session.recordPatch(w, r) + + case http.MethodPut: w.WriteHeader(http.StatusCreated) } })) @@ -1388,36 +1541,32 @@ func TestManifestPush(t *testing.T) { var manifestContentType string var serverURL string + session := newChunkedSession() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Handle blob check (HEAD) - if r.Method == http.MethodHead { + switch { + case r.Method == http.MethodHead: http.NotFound(w, r) - return - } - // Handle blob upload initiate (POST) - if r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/blobs/uploads") { + case r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/blobs/uploads"): w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) w.WriteHeader(http.StatusAccepted) - return - } - // Handle blob upload (PUT to blobs) - if r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/blobs/") { + case r.Method == http.MethodPatch: + session.recordPatch(w, r) + + case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/blobs/"): + // Finalize the chunked blob upload w.WriteHeader(http.StatusCreated) - return - } - // Handle manifest push (PUT to manifests) - if r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/manifests/") { + case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/manifests/"): manifestPath = r.URL.Path manifestContentType = r.Header.Get("Content-Type") manifestReceived, _ = io.ReadAll(r.Body) w.WriteHeader(http.StatusCreated) - return - } - http.NotFound(w, r) + default: + http.NotFound(w, r) + } })) defer server.Close() serverURL = server.URL @@ -1923,8 +2072,838 @@ func TestResumePartialFileExactSize(t *testing.T) { if err != nil { t.Fatalf("Failed to read final file: %v", err) } - finalHash := sha256.Sum256(finalData) - if fmt.Sprintf("sha256:%x", finalHash) != digest { + resumeHash := sha256.Sum256(finalData) + if fmt.Sprintf("sha256:%x", resumeHash) != digest { t.Error("Final file hash mismatch") } } + +// ==================== Chunked Upload Tests ==================== + +// chunkedUploadServer creates a test server that implements the OCI chunked +// upload protocol: POST → PATCH* (with Content-Range) → PUT (finalize). +type chunkedUploadServer struct { + t *testing.T + mu sync.Mutex + parts map[int][]byte // part offset -> received data + patchCount int + finalized bool + finalDigest string + finalEtag string + patchHandler func(w http.ResponseWriter, r *http.Request) // optional override + uploadCounter int + serverURL *string +} + +func newChunkedUploadServer(t *testing.T) *chunkedUploadServer { + return &chunkedUploadServer{ + t: t, + parts: make(map[int][]byte), + } +} + +func (s *chunkedUploadServer) handler() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodHead: + http.NotFound(w, r) + + case r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/uploads"): + s.mu.Lock() + s.uploadCounter++ + id := s.uploadCounter + s.mu.Unlock() + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/%d", *s.serverURL, id)) + w.WriteHeader(http.StatusAccepted) + + case r.Method == http.MethodPatch: + if s.patchHandler != nil { + s.patchHandler(w, r) + return + } + s.defaultPatchHandler(w, r) + + case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/uploads"): + s.mu.Lock() + s.finalized = true + s.finalDigest = r.URL.Query().Get("digest") + s.finalEtag = r.URL.Query().Get("etag") + s.mu.Unlock() + w.WriteHeader(http.StatusCreated) + + default: + http.NotFound(w, r) + } + } +} + +func (s *chunkedUploadServer) defaultPatchHandler(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + s.patchCount++ + patchNum := s.patchCount + s.mu.Unlock() + + cr := r.Header.Get("Content-Range") + if cr == "" { + http.Error(w, "missing Content-Range", http.StatusBadRequest) + return + } + + var start, end int64 + fmt.Sscanf(cr, "%d-%d", &start, &end) + + data, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + s.mu.Lock() + s.parts[int(start)] = data + s.mu.Unlock() + + w.Header().Set("Docker-Upload-Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/session-%d", *s.serverURL, patchNum+1)) + w.WriteHeader(http.StatusAccepted) +} + +func (s *chunkedUploadServer) reassemble(totalSize int) []byte { + s.mu.Lock() + defer s.mu.Unlock() + result := make([]byte, totalSize) + for offset, data := range s.parts { + copy(result[offset:], data) + } + return result +} + +func TestComputeParts(t *testing.T) { + tests := []struct { + name string + totalSize int64 + wantParts int + wantFirst int64 + }{ + { + name: "1GB blob — clamped to min part size", + totalSize: 1 << 30, + wantParts: int((1<<30 + minUploadPartSize - 1) / minUploadPartSize), + wantFirst: minUploadPartSize, + }, + { + name: "5GB blob — 16 parts", + totalSize: 5 << 30, + wantParts: 16, + wantFirst: 5 << 30 / 16, + }, + { + name: "20GB blob — clamped to max part size", + totalSize: 20 << 30, + wantParts: int((20<<30 + maxUploadPartSize - 1) / maxUploadPartSize), + wantFirst: maxUploadPartSize, + }, + { + name: "exactly min part size", + totalSize: minUploadPartSize, + wantParts: 1, + wantFirst: minUploadPartSize, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + parts := computeParts(tt.totalSize) + if len(parts) != tt.wantParts { + t.Errorf("computeParts(%d) = %d parts, want %d", tt.totalSize, len(parts), tt.wantParts) + } + if len(parts) > 0 && parts[0].size != tt.wantFirst { + t.Errorf("first part size = %d, want %d", parts[0].size, tt.wantFirst) + } + + // Verify parts cover entire blob with no gaps + var total int64 + for i, p := range parts { + if p.offset != total { + t.Errorf("part %d offset = %d, want %d", i, p.offset, total) + } + if p.n != i { + t.Errorf("part %d n = %d, want %d", i, p.n, i) + } + total += p.size + } + if total != tt.totalSize { + t.Errorf("total part sizes = %d, want %d", total, tt.totalSize) + } + }) + } +} + +func TestChunkedUploadBasic(t *testing.T) { + blobSize := resumeThreshold + 1024 + data := make([]byte, blobSize) + for i := range data { + data[i] = byte((i * 7) % 256) + } + h := sha256.Sum256(data) + digest := fmt.Sprintf("sha256:%x", h) + blob := Blob{Digest: digest, Size: int64(blobSize)} + + clientDir := t.TempDir() + path := filepath.Join(clientDir, digestToPath(digest)) + os.MkdirAll(filepath.Dir(path), 0o755) + os.WriteFile(path, data, 0o644) + + srv := newChunkedUploadServer(t) + var serverURL string + srv.serverURL = &serverURL + server := httptest.NewServer(srv.handler()) + defer server.Close() + serverURL = server.URL + + var progressCalls atomic.Int32 + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + Progress: func(completed, total int64) { + progressCalls.Add(1) + }, + }) + if err != nil { + t.Fatalf("Chunked upload failed: %v", err) + } + + reassembled := srv.reassemble(blobSize) + reassembledHash := sha256.Sum256(reassembled) + if fmt.Sprintf("sha256:%x", reassembledHash) != digest { + t.Error("Reassembled data hash mismatch") + } + + srv.mu.Lock() + if !srv.finalized { + t.Error("Finalize PUT was not called") + } + if srv.finalDigest != digest { + t.Errorf("Finalize digest = %s, want %s", srv.finalDigest, digest) + } + if srv.finalEtag == "" { + t.Error("Finalize etag is empty") + } + if srv.patchCount == 0 { + t.Error("No PATCH requests were sent") + } + srv.mu.Unlock() + + if progressCalls.Load() == 0 { + t.Error("Progress callback never called") + } +} + +// TestSmallBlobUsesChunkedFlow verifies that even small blobs go through the +// PATCH+finalize chunked flow. Server-side redirect logic is gated on PATCH, +// so a single-PUT path could never trigger CDN redirection — every blob must +// use PATCH so the server has the chance to redirect. +func TestChunkedUploadCDNRedirect(t *testing.T) { + blobSize := resumeThreshold + 1024 + data := make([]byte, blobSize) + for i := range data { + data[i] = byte((i * 7) % 256) + } + h := sha256.Sum256(data) + digest := fmt.Sprintf("sha256:%x", h) + blob := Blob{Digest: digest, Size: int64(blobSize)} + + clientDir := t.TempDir() + path := filepath.Join(clientDir, digestToPath(digest)) + os.MkdirAll(filepath.Dir(path), 0o755) + os.WriteFile(path, data, 0o644) + + cdnParts := make(map[string][]byte) + var cdnMu sync.Mutex + var cdnGotAuth atomic.Bool + + cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + cdnGotAuth.Store(true) + } + cdnData, _ := io.ReadAll(r.Body) + cdnMu.Lock() + cdnParts[r.URL.Path] = cdnData + cdnMu.Unlock() + w.WriteHeader(http.StatusCreated) + })) + defer cdn.Close() + + var serverURL string + var patchCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodHead: + http.NotFound(w, r) + case r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/uploads"): + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) + w.WriteHeader(http.StatusAccepted) + case r.Method == http.MethodPatch: + n := patchCount.Add(1) + w.Header().Set("Docker-Upload-Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/session-%d", serverURL, n+1)) + cdnPath := fmt.Sprintf("/cdn/part-%d", n) + http.Redirect(w, r, cdn.URL+cdnPath, http.StatusTemporaryRedirect) + case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/uploads"): + w.WriteHeader(http.StatusCreated) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + }) + if err != nil { + t.Fatalf("Upload with CDN redirect failed: %v", err) + } + + cdnMu.Lock() + totalCDNBytes := 0 + for _, d := range cdnParts { + totalCDNBytes += len(d) + } + cdnMu.Unlock() + + if totalCDNBytes != blobSize { + t.Errorf("CDN received %d bytes, want %d", totalCDNBytes, blobSize) + } + + if cdnGotAuth.Load() { + t.Error("CDN received Authorization header — should not be sent to CDN") + } +} + +func TestChunkedUploadPartRetry(t *testing.T) { + blobSize := resumeThreshold + 1024 + data := make([]byte, blobSize) + for i := range data { + data[i] = byte(i % 256) + } + h := sha256.Sum256(data) + digest := fmt.Sprintf("sha256:%x", h) + blob := Blob{Digest: digest, Size: int64(blobSize)} + + clientDir := t.TempDir() + path := filepath.Join(clientDir, digestToPath(digest)) + os.MkdirAll(filepath.Dir(path), 0o755) + os.WriteFile(path, data, 0o644) + + var patchAttempts atomic.Int32 + + srv := newChunkedUploadServer(t) + srv.patchHandler = func(w http.ResponseWriter, r *http.Request) { + attempt := patchAttempts.Add(1) + if attempt == 1 { + http.Error(w, "server error", http.StatusInternalServerError) + return + } + srv.defaultPatchHandler(w, r) + } + var serverURL string + srv.serverURL = &serverURL + server := httptest.NewServer(srv.handler()) + defer server.Close() + serverURL = server.URL + + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + }) + if err != nil { + t.Fatalf("Upload with retry failed: %v", err) + } + + if patchAttempts.Load() < 2 { + t.Errorf("Expected at least 2 PATCH attempts, got %d", patchAttempts.Load()) + } + + reassembled := srv.reassemble(blobSize) + reassembledHash := sha256.Sum256(reassembled) + if fmt.Sprintf("sha256:%x", reassembledHash) != digest { + t.Error("Data integrity failed after retry") + } +} + +// multiPartTestHelper creates an uploader with small part sizes so multi-part +// behavior can be tested with small blobs. Returns the uploader and blob data. +func multiPartTestHelper(t *testing.T, blobSize int, partSize int64, serverURL string) (*uploader, Blob, []byte) { + t.Helper() + data := make([]byte, blobSize) + for i := range data { + data[i] = byte((i*7 + 13) % 256) + } + h := sha256.Sum256(data) + digest := fmt.Sprintf("sha256:%x", h) + blob := Blob{Digest: digest, Size: int64(blobSize)} + + clientDir := t.TempDir() + path := filepath.Join(clientDir, digestToPath(digest)) + os.MkdirAll(filepath.Dir(path), 0o755) + os.WriteFile(path, data, 0o644) + + u := &uploader{ + client: defaultClient, + baseURL: serverURL, + srcDir: clientDir, + userAgent: defaultUserAgent, + progress: newProgressTracker(int64(blobSize), nil), + makeParts: func(totalSize int64) []uploadPart { + return computePartsWithLimits(totalSize, 16, partSize, partSize*10) + }, + } + return u, blob, data +} + +func TestChunkedUploadMultiPartSessionURLChain(t *testing.T) { + // Use 10KB blobs with 2KB parts → 5 parts, exercising the URL chain + blobSize := 10240 + partSize := int64(2048) + + var patchURLs []string + var mu sync.Mutex + + srv := newChunkedUploadServer(t) + srv.patchHandler = func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + patchURLs = append(patchURLs, r.URL.Path) + mu.Unlock() + srv.defaultPatchHandler(w, r) + } + var serverURL string + srv.serverURL = &serverURL + server := httptest.NewServer(srv.handler()) + defer server.Close() + serverURL = server.URL + + u, blob, data := multiPartTestHelper(t, blobSize, partSize, server.URL) + + f, err := os.Open(filepath.Join(u.srcDir, digestToPath(blob.Digest))) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + initURL := fmt.Sprintf("%s/v2/library/_/blobs/uploads/init-1", server.URL) + n, err := u.putChunked(context.Background(), initURL, f, blob) + if err != nil { + t.Fatalf("putChunked failed: %v", err) + } + if n != int64(blobSize) { + t.Errorf("bytes written = %d, want %d", n, blobSize) + } + + // Verify data integrity + reassembled := srv.reassemble(blobSize) + if !bytes.Equal(reassembled, data) { + t.Error("Reassembled data mismatch") + } + + mu.Lock() + defer mu.Unlock() + + // Should have 5 parts with distinct URLs + if len(patchURLs) != 5 { + t.Fatalf("Expected 5 PATCH requests, got %d", len(patchURLs)) + } + + // First PATCH uses the init URL + if !strings.Contains(patchURLs[0], "init-1") { + t.Errorf("First PATCH URL should contain init-1, got %s", patchURLs[0]) + } + + // Subsequent PATCHes should use session URLs from Docker-Upload-Location + for i := 1; i < len(patchURLs); i++ { + if patchURLs[i] == patchURLs[i-1] { + t.Errorf("PATCH %d used same URL as PATCH %d — chain broken", i, i-1) + } + if !strings.Contains(patchURLs[i], "session-") { + t.Errorf("PATCH %d URL should contain session-, got %s", i, patchURLs[i]) + } + } +} + +func TestChunkedUploadMultiPartDataIntegrity(t *testing.T) { + // Non-evenly-divisible: 10001 bytes with 3000-byte parts → 4 parts (3000+3000+3000+1001) + blobSize := 10001 + partSize := int64(3000) + + srv := newChunkedUploadServer(t) + var serverURL string + srv.serverURL = &serverURL + server := httptest.NewServer(srv.handler()) + defer server.Close() + serverURL = server.URL + + u, blob, data := multiPartTestHelper(t, blobSize, partSize, server.URL) + + f, err := os.Open(filepath.Join(u.srcDir, digestToPath(blob.Digest))) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + initURL := fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", server.URL) + _, err = u.putChunked(context.Background(), initURL, f, blob) + if err != nil { + t.Fatalf("putChunked failed: %v", err) + } + + reassembled := srv.reassemble(blobSize) + if !bytes.Equal(reassembled, data) { + t.Error("Reassembled data mismatch with non-evenly-divisible parts") + } + + srv.mu.Lock() + if srv.patchCount != 4 { + t.Errorf("Expected 4 PATCH requests, got %d", srv.patchCount) + } + srv.mu.Unlock() +} + +func TestChunkedUploadMultiPartProgressRollback(t *testing.T) { + blobSize := 6000 + partSize := int64(2000) // 3 parts + + var patchAttempts atomic.Int32 + + srv := newChunkedUploadServer(t) + srv.patchHandler = func(w http.ResponseWriter, r *http.Request) { + attempt := patchAttempts.Add(1) + // Fail the second PATCH attempt (part 1, first try). Drain the body + // before erroring so the server sends 100 Continue (under Expect: + // 100-continue) and the client uploads the body — that's what makes + // the progress rollback observable. + if attempt == 2 { + io.Copy(io.Discard, r.Body) + http.Error(w, "server error", http.StatusInternalServerError) + return + } + srv.defaultPatchHandler(w, r) + } + var serverURL string + srv.serverURL = &serverURL + server := httptest.NewServer(srv.handler()) + defer server.Close() + serverURL = server.URL + + u, blob, data := multiPartTestHelper(t, blobSize, partSize, server.URL) + // Track progress + var progressValues []int64 + var mu sync.Mutex + u.progress = newProgressTracker(int64(blobSize), func(completed, total int64) { + mu.Lock() + progressValues = append(progressValues, completed) + mu.Unlock() + }) + + f, err := os.Open(filepath.Join(u.srcDir, digestToPath(blob.Digest))) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + initURL := fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", server.URL) + _, err = u.putChunked(context.Background(), initURL, f, blob) + if err != nil { + t.Fatalf("putChunked failed: %v", err) + } + + // Verify data integrity despite retry + reassembled := srv.reassemble(blobSize) + if !bytes.Equal(reassembled, data) { + t.Error("Data mismatch after retry") + } + + // Verify progress had a rollback (decrease) then recovered + mu.Lock() + defer mu.Unlock() + + hadDecrease := false + for i := 1; i < len(progressValues); i++ { + if progressValues[i] < progressValues[i-1] { + hadDecrease = true + break + } + } + if !hadDecrease { + t.Error("Expected progress to decrease (rollback) during retry, but it was monotonic") + } + + // Final should equal blob size + if len(progressValues) > 0 && progressValues[len(progressValues)-1] != int64(blobSize) { + t.Errorf("Final progress = %d, want %d", progressValues[len(progressValues)-1], blobSize) + } +} + +// ==================== v2 direct-upload extension tests ==================== + +// TestV2InitRequestShape verifies the init POST advertises the v2 capability +// with the expected query parameter and headers, and that the request body +// is empty. +func TestV2InitRequestShape(t *testing.T) { + clientDir := t.TempDir() + blob, _ := createTestBlob(t, clientDir, 4096) + + var sawDigestQuery, sawCapHeader, sawSizeHeader string + var bodyLen int + session := newChunkedSession() + + var serverURL string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodHead: + http.NotFound(w, r) + case http.MethodPost: + sawDigestQuery = r.URL.Query().Get("digest") + sawCapHeader = r.Header.Get("X-Redirect-Uploads") + sawSizeHeader = r.Header.Get("X-Content-Length") + body, _ := io.ReadAll(r.Body) + bodyLen = len(body) + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) + w.WriteHeader(http.StatusAccepted) + case http.MethodPatch: + session.recordPatch(w, r) + case http.MethodPut: + w.WriteHeader(http.StatusCreated) + } + })) + defer server.Close() + serverURL = server.URL + + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + }) + if err != nil { + t.Fatalf("Upload failed: %v", err) + } + + if sawDigestQuery != blob.Digest { + t.Errorf("init POST ?digest= = %q, want %q", sawDigestQuery, blob.Digest) + } + if sawCapHeader != "2" { + t.Errorf("init POST X-Redirect-Uploads = %q, want %q", sawCapHeader, "2") + } + if sawSizeHeader != fmt.Sprintf("%d", blob.Size) { + t.Errorf("init POST X-Content-Length = %q, want %q", sawSizeHeader, fmt.Sprintf("%d", blob.Size)) + } + if bodyLen != 0 { + t.Errorf("init POST body length = %d, want 0", bodyLen) + } +} + +// TestV2DirectUpload verifies the v2 happy path: server returns +// X-Direct-Upload-URL + X-Signed-Header-X-Amz-Checksum-Sha256, the client +// PUTs body to the direct URL with the forwarded checksum header, then +// commits via a bodyless PUT to the session URL. +func TestV2DirectUpload(t *testing.T) { + clientDir := t.TempDir() + blob, data := createTestBlob(t, clientDir, 8192) + + var ( + cdnReceived []byte + cdnChecksumHeader string + cdnAuthHeader string + cdnHits atomic.Int32 + ) + + cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cdnHits.Add(1) + if r.Method != http.MethodPut { + http.Error(w, "want PUT", http.StatusMethodNotAllowed) + return + } + cdnChecksumHeader = r.Header.Get("X-Amz-Checksum-Sha256") + cdnAuthHeader = r.Header.Get("Authorization") + cdnReceived, _ = io.ReadAll(r.Body) + w.WriteHeader(http.StatusOK) + })) + defer cdn.Close() + + // Compute the checksum the registry would sign — base64 of the SHA-256 + // binary digest. The mock just hands the value back to the client; the + // client forwards it to the CDN. + hexDigest := strings.TrimPrefix(blob.Digest, "sha256:") + digestBytes := make([]byte, len(hexDigest)/2) + for i := range digestBytes { + fmt.Sscanf(hexDigest[i*2:i*2+2], "%02x", &digestBytes[i]) + } + expectedChecksum := base64.StdEncoding.EncodeToString(digestBytes) + + var ( + commitDigest atomic.Value + commitBody atomic.Int32 + ) + + var serverURL string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodHead: + http.NotFound(w, r) + case http.MethodPost: + directURL := cdn.URL + "/upload/" + blob.Digest + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) + w.Header().Set("X-Direct-Upload-URL", directURL) + w.Header().Set("X-Signed-Header-X-Amz-Checksum-Sha256", expectedChecksum) + w.WriteHeader(http.StatusAccepted) + case http.MethodPut: + commitDigest.Store(r.URL.Query().Get("digest")) + body, _ := io.ReadAll(r.Body) + commitBody.Store(int32(len(body))) + w.WriteHeader(http.StatusCreated) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + }) + if err != nil { + t.Fatalf("Upload failed: %v", err) + } + + if got := cdnHits.Load(); got != 1 { + t.Errorf("CDN hits = %d, want 1", got) + } + if !bytes.Equal(cdnReceived, data) { + t.Errorf("CDN body length = %d, want %d", len(cdnReceived), len(data)) + } + if cdnChecksumHeader != expectedChecksum { + t.Errorf("CDN x-amz-checksum-sha256 = %q, want %q", cdnChecksumHeader, expectedChecksum) + } + if cdnAuthHeader != "" { + t.Errorf("CDN Authorization = %q, want empty (presigned URL shouldn't carry auth)", cdnAuthHeader) + } + if got, _ := commitDigest.Load().(string); got != blob.Digest { + t.Errorf("commit ?digest= = %q, want %q", got, blob.Digest) + } + if commitBody.Load() != 0 { + t.Errorf("commit body length = %d, want 0", commitBody.Load()) + } +} + +// TestV2FallbackToChunked verifies that when the server returns a standard +// 202 without v2 extension headers, the client falls back to the chunked +// PATCH path. This exercises the vanilla Docker Registry compatibility. +func TestV2FallbackToChunked(t *testing.T) { + clientDir := t.TempDir() + blob, data := createTestBlob(t, clientDir, 8192) + + var ( + uploadedBlobs sync.Map + patchHit atomic.Int32 + commitHit atomic.Int32 + ) + session := newChunkedSession() + + var serverURL string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodHead: + http.NotFound(w, r) + case http.MethodPost: + // vanilla: standard Location only, no v2 extension headers + w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL)) + w.WriteHeader(http.StatusAccepted) + case http.MethodPatch: + patchHit.Add(1) + session.recordPatch(w, r) + case http.MethodPut: + commitHit.Add(1) + digest := r.URL.Query().Get("digest") + uploadedBlobs.Store(digest, session.finalize(r.URL.Path)) + w.WriteHeader(http.StatusCreated) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + }) + if err != nil { + t.Fatalf("Upload failed: %v", err) + } + + if patchHit.Load() < 1 { + t.Error("expected at least one PATCH (chunked fallback), got none") + } + if commitHit.Load() != 1 { + t.Errorf("commit PUT hits = %d, want 1", commitHit.Load()) + } + if got, ok := uploadedBlobs.Load(blob.Digest); !ok { + t.Error("blob not uploaded") + } else if !bytes.Equal(got.([]byte), data) { + t.Errorf("uploaded body length = %d, want %d", len(got.([]byte)), len(data)) + } +} + +// TestV2BlobAlreadyExists verifies that a 201 Created response from the init +// POST short-circuits the upload — the server has matched our ?digest= +// against existing storage and there's nothing to upload. +func TestV2BlobAlreadyExists(t *testing.T) { + clientDir := t.TempDir() + blob, _ := createTestBlob(t, clientDir, 1024) + + var ( + postHits atomic.Int32 + patchHits atomic.Int32 + putHits atomic.Int32 + ) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodHead: + http.NotFound(w, r) + case http.MethodPost: + postHits.Add(1) + // Server matched ?digest= against existing storage. + w.WriteHeader(http.StatusCreated) + case http.MethodPatch: + patchHits.Add(1) + w.WriteHeader(http.StatusAccepted) + case http.MethodPut: + putHits.Add(1) + w.WriteHeader(http.StatusCreated) + } + })) + defer server.Close() + + err := Upload(context.Background(), UploadOptions{ + Blobs: []Blob{blob}, + BaseURL: server.URL, + SrcDir: clientDir, + }) + if err != nil { + t.Fatalf("Upload failed: %v", err) + } + + if postHits.Load() != 1 { + t.Errorf("init POST hits = %d, want 1", postHits.Load()) + } + if patchHits.Load() != 0 { + t.Errorf("PATCH hits = %d, want 0 (blob existed; nothing to upload)", patchHits.Load()) + } + if putHits.Load() != 0 { + t.Errorf("PUT hits = %d, want 0 (blob existed; nothing to upload)", putHits.Load()) + } +} diff --git a/x/transfer/upload.go b/x/transfer/upload.go index c75ee3e6..fa7a7962 100644 --- a/x/transfer/upload.go +++ b/x/transfer/upload.go @@ -5,14 +5,19 @@ import ( "bytes" "cmp" "context" + "crypto/md5" "errors" "fmt" "io" "log/slog" + "maps" "net/http" "net/url" "os" "path/filepath" + "strings" + "sync" + "sync/atomic" "time" "github.com/ollama/ollama/logutil" @@ -26,11 +31,59 @@ type uploader struct { baseURL string srcDir string repository string // Repository path for blob URLs (e.g., "library/model") - token *string + tokenMu sync.RWMutex + token string getToken func(context.Context, AuthChallenge) (string, error) userAgent string progress *progressTracker logger *slog.Logger + // bodySem caps the number of simultaneous body-bearing transfers so a + // modest home uplink isn't saturated. Always set by upload(); nil only + // when tests build uploader directly (in which case holdBody is a no-op). + bodySem *semaphore.Weighted + makeParts func(int64) []uploadPart // controls how blobs are split for chunked upload +} + +// authToken returns the current bearer token. Safe to call concurrently with +// refreshToken. +func (u *uploader) authToken() string { + u.tokenMu.RLock() + defer u.tokenMu.RUnlock() + return u.token +} + +// refreshToken coalesces token fetches so concurrent 401s don't all hit the +// auth server. prev is the token the caller used in the request that got +// rejected: if the stored token has already moved past prev, another +// goroutine has refreshed and we just observe its result; otherwise the +// caller holds the lock and performs the fetch. +func (u *uploader) refreshToken(ctx context.Context, ch AuthChallenge, prev string) error { + u.tokenMu.Lock() + defer u.tokenMu.Unlock() + if u.token != prev { + return nil + } + if u.getToken == nil { + return errors.New("no token refresh callback") + } + t, err := u.getToken(ctx, ch) + if err != nil { + return err + } + u.token = t + return nil +} + +// holdBody acquires a body-transfer slot. The returned release must be called +// exactly once after the body-bearing request completes (defer is fine). +func (u *uploader) holdBody(ctx context.Context) (func(), error) { + if u.bodySem == nil { + return func() {}, nil + } + if err := u.bodySem.Acquire(ctx, 1); err != nil { + return nil, err + } + return func() { u.bodySem.Release(1) }, nil } func upload(ctx context.Context, opts UploadOptions) error { @@ -38,45 +91,40 @@ func upload(ctx context.Context, opts UploadOptions) error { return nil } - token := opts.Token u := &uploader{ client: cmp.Or(opts.Client, defaultClient), baseURL: opts.BaseURL, srcDir: opts.SrcDir, repository: cmp.Or(opts.Repository, "library/_"), - token: &token, + token: opts.Token, getToken: opts.GetToken, userAgent: cmp.Or(opts.UserAgent, defaultUserAgent), logger: opts.Logger, } + // 0 or negative serializes; never unbounded. + u.bodySem = semaphore.NewWeighted(int64(max(1, opts.BodyConcurrency))) if len(opts.Blobs) > 0 { - // Phase 1: Fast parallel HEAD checks to find which blobs need uploading + // Discover which blobs the server already has so we can skip uploading needsUpload := make([]bool, len(opts.Blobs)) - { - sem := semaphore.NewWeighted(128) // High concurrency for HEAD checks - g, gctx := errgroup.WithContext(ctx) - for i, blob := range opts.Blobs { - g.Go(func() error { - if err := sem.Acquire(gctx, 1); err != nil { - return err - } - defer sem.Release(1) - exists, err := u.exists(gctx, blob) - if err != nil { - return err - } - if !exists { - needsUpload[i] = true - } else if u.logger != nil { - u.logger.Debug("blob exists", "digest", blob.Digest) - } - return nil - }) - } - if err := g.Wait(); err != nil { - return err - } + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(128) + for i, blob := range opts.Blobs { + g.Go(func() error { + exists, err := u.exists(gctx, blob) + if err != nil { + return err + } + if !exists { + needsUpload[i] = true + } else if u.logger != nil { + u.logger.Debug("blob exists", "digest", blob.Digest) + } + return nil + }) + } + if err := g.Wait(); err != nil { + return err } // Filter to only blobs that need uploading, but track total across all blobs @@ -102,10 +150,12 @@ func upload(ctx context.Context, opts UploadOptions) error { u.logger.Debug("all blobs exist, nothing to upload") } } else { - // Phase 2: Upload blobs that don't exist + // Upload the blobs the server doesn't already have. Concurrency + // caps blob-level parallelism. concurrency := cmp.Or(opts.Concurrency, DefaultUploadConcurrency) sem := semaphore.NewWeighted(int64(concurrency)) + start := time.Now() g, gctx := errgroup.WithContext(ctx) for _, blob := range toUpload { g.Go(func() error { @@ -116,7 +166,18 @@ func upload(ctx context.Context, opts UploadOptions) error { return u.upload(gctx, blob) }) } - if err := g.Wait(); err != nil { + err := g.Wait() + elapsed := time.Since(start) + done := u.progress.completed.Load() - alreadyExists + mbps := float64(done) / 1e6 / max(0.001, elapsed.Seconds()) + slog.Debug("upload summary", + "blobs", len(toUpload), + "bytes", done, + "duration", elapsed.Round(time.Millisecond), + "mb_per_sec", fmt.Sprintf("%.1f", mbps), + "max_transfers", max(1, opts.BodyConcurrency), + ) + if err != nil { return err } } @@ -139,8 +200,8 @@ func (u *uploader) upload(ctx context.Context, blob Blob) error { for attempt := range maxRetries { if attempt > 0 { - // Use longer backoff for uploads — server-side rate limiting - // and S3 upload session creation need real breathing room. + // Longer backoff for uploads — server-side rate limiting and + // upload-session bookkeeping need real breathing room. // attempt 1: up to 2s, attempt 2: up to 4s, attempt 3: up to 8s, etc. if err := backoff(ctx, attempt, 2*time.Second< 0 { - // Start at 5s and cap at 30s — the server needs real breathing - // room when it's dropping Location headers under load. if err := backoff(ctx, attempt, min(5*time.Second< response headers must be + // echoed back on the direct PUT under — the client doesn't + // need to know which headers, just to forward whatever was signed. + if directURL := resp.Header.Get("X-Direct-Upload-URL"); directURL != "" { + // Validate it parses and is absolute, but keep the original + // string. url.Parse + String() round-trips with normalization + // (percent-encoding case, query ordering) which can change the + // canonical form a signed URL was computed over. + if d, err := url.Parse(directURL); err == nil && d.IsAbs() { + ep.directUploadURL = directURL + ep.signedHeaders = make(http.Header) + const signedPrefix = "X-Signed-Header-" + for k, vs := range resp.Header { + name, ok := strings.CutPrefix(k, signedPrefix) + if !ok { + continue + } + for _, v := range vs { + ep.signedHeaders.Add(name, v) + } + } + } + } + return ep, nil } - return "", lastErr + return uploadEndpoint{}, lastErr } -func (u *uploader) put(ctx context.Context, uploadURL string, f *os.File, size int64) (int64, error) { - // uploadURL is empty when initUpload determined the blob already exists (201 Created) +// putDirect PUTs the blob body to the URL the server returned, echoing any +// signed headers it provided. The follow-up commit PUT records the blob on +// the server side with no body. +func (u *uploader) putDirect(ctx context.Context, ep uploadEndpoint, f *os.File, blob Blob) (int64, error) { + pr, err := u.streamPutBody(ctx, ep, f, blob) + if err != nil { + return pr.bytes(), err + } + // Body slot is released; commit is bookkeeping (no body) and shouldn't + // hold the cap from other body uploads. + if err := u.commit(ctx, ep.sessionURL, blob.Digest); err != nil { + return pr.bytes(), err + } + return pr.bytes(), nil +} + +// streamPutBody PUTs the blob body to the server-supplied URL, holding a +// body-transfer slot only for the duration of the body PUT (not the +// follow-up commit). Returns the progressReader so the caller can report +// pr.n on commit failure. +func (u *uploader) streamPutBody(ctx context.Context, ep uploadEndpoint, f *os.File, blob Blob) (*progressReader, error) { + release, err := u.holdBody(ctx) + if err != nil { + return &progressReader{}, err + } + defer release() + + br := bufio.NewReaderSize(f, 256*1024) + pr := &progressReader{reader: br, tracker: u.progress} + + req, _ := http.NewRequestWithContext(ctx, http.MethodPut, ep.directUploadURL, pr) + req.ContentLength = blob.Size + req.Header.Set("Content-Type", "application/octet-stream") + req.Header.Set("User-Agent", u.userAgent) + // Echo signed headers — overwrite any defaults we set above so the + // signed value wins. Appending would leave duplicates that change the + // signature canonical form and the upload would be rejected. + maps.Copy(req.Header, ep.signedHeaders) + // No Authorization — the direct-upload URL carries its own credential. + + resp, err := u.client.Do(req) + if err != nil { + return pr, fmt.Errorf("direct put: %w", err) + } + defer func() { io.Copy(io.Discard, resp.Body); resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + body, _ := io.ReadAll(resp.Body) + return pr, fmt.Errorf("direct put: status %d: %s", resp.StatusCode, body) + } + return pr, nil +} + +// commit sends a bodyless PUT to the session URL with ?digest= so the server +// records a blob whose body was uploaded out-of-band. +func (u *uploader) commit(ctx context.Context, sessionURL, digest string) error { + finalURL, err := url.Parse(sessionURL) + if err != nil { + return fmt.Errorf("parse session URL: %w", err) + } + q := finalURL.Query() + q.Set("digest", digest) + finalURL.RawQuery = q.Encode() + + return u.bodylessRegistryPUT(ctx, finalURL.String(), "commit") +} + +// bodylessRegistryPUT sends a zero-body PUT to a registry URL, retrying with +// backoff on transport/server errors and once on auth challenge. op is used +// as the error prefix. +func (u *uploader) bodylessRegistryPUT(ctx context.Context, url string, op string) error { + var lastErr error + for try := range maxRetries { + if try > 0 { + if err := backoff(ctx, try, 2*time.Second< 0 { + if err := backoff(ctx, try, 2*time.Second<= http.StatusBadRequest: body, _ := io.ReadAll(resp.Body) - return pr.n, fmt.Errorf("status %d: %s", resp.StatusCode, body) + return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: status %d: %s", part.n, resp.StatusCode, body) } - return pr.n, nil + + if next == nil { + return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: no next URL in response", part.n) + } + return next, partHash.Sum(nil), pr.bytes(), nil +} + +// putPartToCDN re-uploads a part's data to a CDN redirect URL via PUT. +// Returns the md5 sum of bytes actually streamed to the CDN, the byte count, +// and any error. The hash is fed inline so the composite etag we eventually +// send to the registry reflects what the storage backend stored, not what the +// client tried to PATCH. +func (u *uploader) putPartToCDN(ctx context.Context, cdnURL *url.URL, part *uploadPart, f *os.File) ([]byte, int64, error) { + sr := io.NewSectionReader(f, part.offset, part.size) + br := bufio.NewReaderSize(sr, 256*1024) + pr := &progressReader{reader: br, tracker: u.progress} + partHash := md5.New() + body := io.TeeReader(pr, partHash) + + req, _ := http.NewRequestWithContext(ctx, http.MethodPut, cdnURL.String(), body) + req.ContentLength = part.size + req.Header.Set("Content-Type", "application/octet-stream") + req.Header.Set("User-Agent", u.userAgent) + // No Authorization — the redirect URL carries its own credential. + + resp, err := u.client.Do(req) + if err != nil { + return nil, pr.bytes(), fmt.Errorf("cdn put part %d: %w", part.n, err) + } + defer func() { io.Copy(io.Discard, resp.Body); resp.Body.Close() }() + + if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, pr.bytes(), fmt.Errorf("cdn put part %d: status %d: %s", part.n, resp.StatusCode, body) + } + return partHash.Sum(nil), pr.bytes(), nil +} + +// Chunked-upload sizing — when computeParts splits a blob into parts for the +// multipart fallback, parts are sized in [minUploadPartSize, maxUploadPartSize] +// with a target count of numUploadParts. Smaller blobs end up as a single +// sub-minimum part. +const ( + numUploadParts = 16 + minUploadPartSize int64 = 100 << 20 // 100 MB + maxUploadPartSize int64 = 1000 << 20 // ~1 GB +) + +// uploadPart represents a chunk of a blob for the multipart fallback. +type uploadPart struct { + n int + offset int64 + size int64 +} + +// computeParts divides a blob into upload parts using default limits. +func computeParts(totalSize int64) []uploadPart { + return computePartsWithLimits(totalSize, numUploadParts, minUploadPartSize, maxUploadPartSize) +} + +// computePartsWithLimits divides a blob into upload parts with configurable limits. +func computePartsWithLimits(totalSize int64, nParts int, minPart, maxPart int64) []uploadPart { + partSize := totalSize / int64(nParts) + partSize = max(partSize, minPart) + partSize = min(partSize, maxPart) + + var parts []uploadPart + var offset int64 + for offset < totalSize { + size := partSize + if offset+size > totalSize { + size = totalSize - offset + } + parts = append(parts, uploadPart{n: len(parts), offset: offset, size: size}) + offset += size + } + return parts } func (u *uploader) pushManifest(ctx context.Context, repo, ref string, manifest []byte) error { req, _ := http.NewRequestWithContext(ctx, http.MethodPut, fmt.Sprintf("%s/v2/%s/manifests/%s", u.baseURL, repo, ref), bytes.NewReader(manifest)) req.Header.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json") req.Header.Set("User-Agent", u.userAgent) - if *u.token != "" { - req.Header.Set("Authorization", "Bearer "+*u.token) + prev := u.authToken() + if prev != "" { + req.Header.Set("Authorization", "Bearer "+prev) } resp, err := u.client.Do(req) @@ -374,7 +776,7 @@ func (u *uploader) pushManifest(ctx context.Context, repo, ref string, manifest if resp.StatusCode == http.StatusUnauthorized && u.getToken != nil { ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate")) - if *u.token, err = u.getToken(ctx, ch); err != nil { + if err := u.refreshToken(ctx, ch, prev); err != nil { return err } return u.pushManifest(ctx, repo, ref, manifest) @@ -387,17 +789,25 @@ func (u *uploader) pushManifest(ctx context.Context, repo, ref string, manifest return nil } +// progressReader counts bytes streamed through Read. The byte counter is +// atomic because the HTTP transport's writeLoop runs concurrently with the +// goroutine that returns the count after a non-2xx response — the transport +// may still be calling Read while we're already returning from uploadOnePart. type progressReader struct { reader io.Reader tracker *progressTracker - n int64 + n atomic.Int64 } func (r *progressReader) Read(p []byte) (int, error) { n, err := r.reader.Read(p) if n > 0 { - r.n += int64(n) + r.n.Add(int64(n)) r.tracker.add(int64(n)) } return n, err } + +func (r *progressReader) bytes() int64 { + return r.n.Load() +}