vendor/ollama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mlx: refined model push behavior (#15431)

Browse Source 
* 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
This commit is contained in:
Daniel Hiltgen
2026-05-08 14:25:30 -07:00
committed by GitHub
parent f866e7608f
commit 1e1b34dada
7 changed files with 1754 additions and 277 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
cmd/cmd.go
View File

@@ -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"],

53
envconfig/config.go
View File

@@ -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"},

38
server/images.go
View File

@@ -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(),
})
}

90
x/transfer/download.go
View File

@@ -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:

76
x/transfer/transfer.go
View File

@@ -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"

1145
x/transfer/transfer_test.go
View File

File diff suppressed because it is too large Load Diff

626
x/transfer/upload.go
View File

@@ -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<<uint(attempt-1)); err != nil {
return err
@@ -171,28 +232,41 @@ func (u *uploader) uploadOnce(ctx context.Context, blob Blob) (int64, error) {
u.logger.Debug("uploading blob", "digest", blob.Digest, "size", blob.Size)
}
// Init upload
uploadURL, err := u.initUpload(ctx, blob)
ep, err := u.initUpload(ctx, blob)
if err != nil {
return 0, err
}
if ep.sessionURL == "" {
// Server matched ?digest= against existing storage; nothing to
// upload. Credit the full size to progress so a retry-after-failure
// (where the prior attempt streamed bytes that were rolled back)
// still finishes at 100%.
u.progress.add(blob.Size)
return blob.Size, nil
}
// Open file
f, err := os.Open(filepath.Join(u.srcDir, digestToPath(blob.Digest)))
if err != nil {
return 0, err
}
defer f.Close()
// PUT blob
return u.put(ctx, uploadURL, f, blob.Size)
if ep.directUploadURL != "" {
// Body goes straight to the URL the server returned; the server
// only sees a tiny commit roundtrip.
return u.putDirect(ctx, ep, f, blob)
}
// Body goes to the server in parts via PATCH, followed by a finalize PUT.
return u.putChunked(ctx, ep.sessionURL, f, blob)
}
func (u *uploader) exists(ctx context.Context, blob Blob) (bool, error) {
req, _ := http.NewRequestWithContext(ctx, http.MethodHead, fmt.Sprintf("%s/v2/%s/blobs/%s", u.baseURL, u.repository, blob.Digest), nil)
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)
@@ -204,7 +278,7 @@ func (u *uploader) exists(ctx context.Context, blob Blob) (bool, error) {
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 false, err
}
return u.exists(ctx, blob)
@@ -215,7 +289,27 @@ func (u *uploader) exists(ctx context.Context, blob Blob) (bool, error) {
const maxInitRetries = 12
func (u *uploader) initUpload(ctx context.Context, blob Blob) (string, error) {
// uploadEndpoint describes where a blob's body should be uploaded after init.
//
// A zero-valued endpoint (sessionURL == "") means the server already has the
// blob and the caller should skip upload.
//
// When sessionURL is set but directUploadURL is empty, the body goes to the
// server in parts via PATCH, then commits with a finalize PUT.
//
// When directUploadURL is set, the body is PUT directly to that URL with any
// signedHeaders the server provided echoed back as request headers. A
// bodyless commit PUT to sessionURL?digest=... then records the blob.
type uploadEndpoint struct {
sessionURL string
directUploadURL string
signedHeaders http.Header // headers the server provided that the client must echo on the direct PUT
}
// initUpload announces the upload to the server and discovers which flow to
// use. The server may return a direct-upload URL alongside the session URL;
// the caller branches on whether one came back.
func (u *uploader) initUpload(ctx context.Context, blob Blob) (uploadEndpoint, error) {
endpoint, _ := url.Parse(fmt.Sprintf("%s/v2/%s/blobs/uploads/", u.baseURL, u.repository))
q := endpoint.Query()
q.Set("digest", blob.Digest)
@@ -224,18 +318,19 @@ func (u *uploader) initUpload(ctx context.Context, blob Blob) (string, error) {
var lastErr error
for attempt := range maxInitRetries {
if attempt > 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<<uint(attempt-1), 30*time.Second)); err != nil {
return "", err
return uploadEndpoint{}, err
}
logutil.Trace("retrying init upload", "digest", blob.Digest, "attempt", attempt+1, "error", lastErr)
}
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, endpoint.String(), nil)
req.Header.Set("User-Agent", u.userAgent)
if *u.token != "" {
req.Header.Set("Authorization", "Bearer "+*u.token)
req.Header.Set("X-Redirect-Uploads", "2")
req.Header.Set("X-Content-Length", fmt.Sprintf("%d", blob.Size))
prev := u.authToken()
if prev != "" {
req.Header.Set("Authorization", "Bearer "+prev)
}
resp, err := u.client.Do(req)
@@ -248,15 +343,16 @@ func (u *uploader) initUpload(ctx context.Context, blob Blob) (string, error) {
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 {
return "", err
if err := u.refreshToken(ctx, ch, prev); err != nil {
return uploadEndpoint{}, err
}
continue
}
if resp.StatusCode == http.StatusCreated {
// Blob was mounted or already exists — no upload needed
return "", nil
// Server matched our ?digest= against existing storage —
// nothing to upload.
return uploadEndpoint{}, nil
}
if resp.StatusCode != http.StatusAccepted {
@@ -269,101 +365,407 @@ func (u *uploader) initUpload(ctx context.Context, blob Blob) (string, error) {
loc = resp.Header.Get("Location")
}
if loc == "" {
// Server returned 202 but no Location — retry, the server may
// be under load and dropping headers.
lastErr = fmt.Errorf("no upload location (server returned 202 without Location header)")
continue
}
locURL, _ := url.Parse(loc)
if !locURL.IsAbs() {
sessionURL, _ := url.Parse(loc)
if !sessionURL.IsAbs() {
base, _ := url.Parse(u.baseURL)
locURL = base.ResolveReference(locURL)
sessionURL = base.ResolveReference(sessionURL)
}
q = locURL.Query()
q.Set("digest", blob.Digest)
locURL.RawQuery = q.Encode()
return locURL.String(), nil
ep := uploadEndpoint{sessionURL: sessionURL.String()}
// Opt-in direct-upload path: enabled only when the server returns an
// upload URL. Any X-Signed-Header-<name> response headers must be
// echoed back on the direct PUT under <name> — 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<<uint(try-1)); err != nil {
return err
}
}
req, _ := http.NewRequestWithContext(ctx, http.MethodPut, url, nil)
req.ContentLength = 0
req.Header.Set("Content-Type", "application/octet-stream")
req.Header.Set("User-Agent", u.userAgent)
prev := u.authToken()
if prev != "" {
req.Header.Set("Authorization", "Bearer "+prev)
}
resp, err := u.client.Do(req)
if err != nil {
lastErr = err
continue
}
switch {
case resp.StatusCode == http.StatusUnauthorized && u.getToken != nil:
io.Copy(io.Discard, resp.Body)
resp.Body.Close()
ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate"))
if err := u.refreshToken(ctx, ch, prev); err != nil {
return err
}
case resp.StatusCode == http.StatusCreated || resp.StatusCode == http.StatusOK:
io.Copy(io.Discard, resp.Body)
resp.Body.Close()
return nil
default:
// Capture body for the error message before closing.
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
resp.Body.Close()
lastErr = fmt.Errorf("%s: status %d: %s", op, resp.StatusCode, body)
}
}
return fmt.Errorf("%w: %v", errMaxRetriesExceeded, lastErr)
}
// putChunked is the fallback used when the server doesn't return a
// direct-upload URL. It splits the blob into parts and sends each via
// PATCH with a Content-Range, following any redirect on the response,
// then finalizes with a composite-MD5 etag PUT.
//
// On failure the function rolls back any progress it accumulated for this
// blob and returns 0 bytes written, so the outer per-blob retry can start
// from a clean state.
func (u *uploader) putChunked(ctx context.Context, uploadURL string, f *os.File, blob Blob) (int64, error) {
if uploadURL == "" {
return 0, nil
}
// Buffer reads for better throughput — 256KB reads instead of default 4KB
br := bufio.NewReaderSize(f, 256*1024)
pr := &progressReader{reader: br, tracker: u.progress}
splitParts := computeParts
if u.makeParts != nil {
splitParts = u.makeParts
}
parts := splitParts(blob.Size)
req, _ := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, pr)
req.ContentLength = size
current, err := url.Parse(uploadURL)
if err != nil {
return 0, fmt.Errorf("parse upload URL: %w", err)
}
composite := md5.New()
var written int64
for i := range parts {
part := &parts[i]
next, partHash, err := u.uploadOnePartWithRetry(ctx, current, part, f)
if err != nil {
u.progress.add(-written)
return 0, err
}
composite.Write(partHash)
written += part.size
if next != nil {
current = next
}
}
q := current.Query()
q.Set("digest", blob.Digest)
q.Set("etag", fmt.Sprintf("%x-%d", composite.Sum(nil), len(parts)))
current.RawQuery = q.Encode()
if err := u.bodylessRegistryPUT(ctx, current.String(), "finalize"); err != nil {
u.progress.add(-written)
return 0, err
}
return written, nil
}
// uploadOnePartWithRetry sends a single part with up to maxPartRetries
// attempts; rolls back per-attempt progress on transient failures so the
// progress tracker stays consistent.
func (u *uploader) uploadOnePartWithRetry(ctx context.Context, sessionURL *url.URL, part *uploadPart, f *os.File) (*url.URL, []byte, error) {
const maxPartRetries = 3
var lastErr error
for try := range maxPartRetries {
if try > 0 {
if err := backoff(ctx, try, 2*time.Second<<uint(try-1)); err != nil {
return nil, nil, err
}
}
next, partHash, n, err := u.uploadOnePart(ctx, sessionURL, part, f)
if err == nil {
return next, partHash, nil
}
// Roll back this attempt's progress so retries don't double-count.
u.progress.add(-n)
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return nil, nil, err
}
lastErr = err
}
return nil, nil, fmt.Errorf("part %d: %w", part.n, lastErr)
}
// uploadOnePart sends one PATCH for a single part and returns the next
// session URL, the part's MD5 sum, the bytes written, and any error. If the
// server replies 307, the body is re-uploaded to the redirect URL via a
// follow-up PUT; the next session URL still comes from the 307 response.
func (u *uploader) uploadOnePart(ctx context.Context, sessionURL *url.URL, part *uploadPart, f *os.File) (*url.URL, []byte, int64, error) {
// Hold the body slot across both the PATCH and any subsequent CDN PUT —
// both transfer body bytes and shouldn't double-count against the cap.
release, err := u.holdBody(ctx)
if err != nil {
return nil, nil, 0, err
}
defer release()
sr := io.NewSectionReader(f, part.offset, part.size)
br := bufio.NewReaderSize(sr, 256*1024)
partHash := md5.New()
pr := &progressReader{reader: br, tracker: u.progress}
body := io.TeeReader(pr, partHash)
req, _ := http.NewRequestWithContext(ctx, http.MethodPatch, sessionURL.String(), body)
req.ContentLength = part.size
req.Header.Set("Content-Type", "application/octet-stream")
req.Header.Set("Content-Range", fmt.Sprintf("%d-%d", part.offset, part.offset+part.size-1))
req.Header.Set("X-Redirect-Uploads", "1")
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)
if err != nil {
return pr.n, fmt.Errorf("put request: %w", err)
return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: %w", part.n, err)
}
defer func() { io.Copy(io.Discard, resp.Body); resp.Body.Close() }()
// Handle auth retry
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 {
return pr.n, err
// The server may return either an absolute or a relative URL in
// Location / Docker-Upload-Location; resolve relative ones against the
// request URL.
loc := resp.Header.Get("Docker-Upload-Location")
if loc == "" {
loc = resp.Header.Get("Location")
}
var next *url.URL
if loc != "" {
next, _ = url.Parse(loc)
if next != nil && !next.IsAbs() {
next = sessionURL.ResolveReference(next)
}
f.Seek(0, 0)
u.progress.add(-pr.n)
return u.put(ctx, uploadURL, f, size)
}
// Handle redirect to CDN
if resp.StatusCode == http.StatusTemporaryRedirect {
loc, _ := resp.Location()
f.Seek(0, 0)
u.progress.add(-pr.n)
br2 := bufio.NewReaderSize(f, 256*1024)
pr2 := &progressReader{reader: br2, tracker: u.progress}
req2, _ := http.NewRequestWithContext(ctx, http.MethodPut, loc.String(), pr2)
req2.ContentLength = size
req2.Header.Set("Content-Type", "application/octet-stream")
req2.Header.Set("User-Agent", u.userAgent)
resp2, err := u.client.Do(req2)
switch {
case resp.StatusCode == http.StatusTemporaryRedirect:
redirectURL, _ := resp.Location()
if redirectURL == nil {
return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: 307 without Location", part.n)
}
// The PATCH attempt's progress is wasted — we re-upload to CDN.
// We can't safely Reset partHash here: the http transport's
// writeLoop may still be feeding TeeReader bytes into it, so
// abandon it and let putPartToCDN compute a fresh hash from the
// bytes that actually land on the storage backend.
u.progress.add(-pr.bytes())
cdnSum, cdnN, err := u.putPartToCDN(ctx, redirectURL, part, f)
if err != nil {
return pr2.n, fmt.Errorf("cdn put request: %w", err)
return nil, nil, cdnN, err
}
defer func() { io.Copy(io.Discard, resp2.Body); resp2.Body.Close() }()
return next, cdnSum, cdnN, nil
if resp2.StatusCode != http.StatusCreated && resp2.StatusCode != http.StatusAccepted {
body, _ := io.ReadAll(resp2.Body)
return pr2.n, fmt.Errorf("status %d: %s", resp2.StatusCode, body)
case resp.StatusCode == http.StatusUnauthorized && u.getToken != nil:
ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate"))
if err := u.refreshToken(ctx, ch, prev); err != nil {
return nil, nil, pr.bytes(), err
}
return pr2.n, nil
}
return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: auth retry", part.n)
if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusAccepted {
case resp.StatusCode >= 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()
}