ollama source for Momentry Core verification
This commit is contained in:
Accusys
468
x/transfer/download.go
Normal file
468
x/transfer/download.go
Normal file
@@ -0,0 +1,468 @@
|
||||
package transfer
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"golang.org/x/sync/errgroup"
|
||||
"golang.org/x/sync/semaphore"
|
||||
)
|
||||
|
||||
var (
|
||||
errStalled = errors.New("download stalled")
|
||||
errSlow = errors.New("download too slow")
|
||||
)
|
||||
|
||||
type downloader struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
destDir string
|
||||
repository string // Repository path for blob URLs (e.g., "library/model")
|
||||
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 {
|
||||
if len(opts.Blobs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Calculate total from all blobs (for accurate progress reporting on resume)
|
||||
var total int64
|
||||
for _, b := range opts.Blobs {
|
||||
total += b.Size
|
||||
}
|
||||
|
||||
// Filter out already-downloaded blobs and track completed bytes
|
||||
var blobs []Blob
|
||||
var alreadyCompleted int64
|
||||
for _, b := range opts.Blobs {
|
||||
if fi, _ := os.Stat(filepath.Join(opts.DestDir, digestToPath(b.Digest))); fi != nil && fi.Size() == b.Size {
|
||||
if opts.Logger != nil {
|
||||
opts.Logger.Debug("blob already exists", "digest", b.Digest, "size", b.Size)
|
||||
}
|
||||
alreadyCompleted += b.Size
|
||||
continue
|
||||
}
|
||||
blobs = append(blobs, b)
|
||||
}
|
||||
if len(blobs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
progress := newProgressTracker(total, opts.Progress)
|
||||
progress.add(alreadyCompleted) // Report already-downloaded bytes upfront
|
||||
|
||||
d := &downloader{
|
||||
client: cmp.Or(opts.Client, defaultClient),
|
||||
baseURL: opts.BaseURL,
|
||||
destDir: opts.DestDir,
|
||||
repository: cmp.Or(opts.Repository, "library/_"),
|
||||
token: opts.Token,
|
||||
getToken: opts.GetToken,
|
||||
userAgent: cmp.Or(opts.UserAgent, defaultUserAgent),
|
||||
stallTimeout: cmp.Or(opts.StallTimeout, defaultStallTimeout),
|
||||
progress: progress,
|
||||
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 {
|
||||
if err := sem.Acquire(ctx, 1); err != nil {
|
||||
return err
|
||||
}
|
||||
defer sem.Release(1)
|
||||
return d.download(ctx, blob)
|
||||
})
|
||||
}
|
||||
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 {
|
||||
var lastErr error
|
||||
var slowRetries int
|
||||
attempt := 0
|
||||
|
||||
for attempt < maxRetries {
|
||||
if attempt > 0 {
|
||||
if err := backoff(ctx, attempt, time.Second<<uint(attempt-1)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
n, err := d.downloadOnce(ctx, blob)
|
||||
if err == nil {
|
||||
// Skip speed recording for tiny blobs — their transfer time is
|
||||
// dominated by HTTP overhead, not throughput, and would pollute
|
||||
// the median used for stall detection.
|
||||
if blob.Size >= smallBlobSpeedThreshold {
|
||||
if s := time.Since(start).Seconds(); s > 0 {
|
||||
d.speeds.record(float64(blob.Size) / s)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
d.progress.add(-n) // rollback
|
||||
|
||||
// Preserve partial .tmp files for large blobs to enable resume
|
||||
if blob.Size < resumeThreshold {
|
||||
dest := filepath.Join(d.destDir, digestToPath(blob.Digest))
|
||||
os.Remove(dest + ".tmp")
|
||||
}
|
||||
|
||||
switch {
|
||||
case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
|
||||
return err
|
||||
case errors.Is(err, errStalled):
|
||||
// Don't count stall retries against limit
|
||||
case errors.Is(err, errSlow):
|
||||
if slowRetries++; slowRetries >= 3 {
|
||||
attempt++ // Only count after 3 slow retries
|
||||
}
|
||||
default:
|
||||
attempt++
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
return fmt.Errorf("%w: %v", errMaxRetriesExceeded, lastErr)
|
||||
}
|
||||
|
||||
func (d *downloader) downloadOnce(ctx context.Context, blob Blob) (int64, error) {
|
||||
if d.logger != nil {
|
||||
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 {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// Check for existing partial .tmp file for resume
|
||||
dest := filepath.Join(d.destDir, digestToPath(blob.Digest))
|
||||
tmp := dest + ".tmp"
|
||||
var existingSize int64
|
||||
if blob.Size >= resumeThreshold {
|
||||
if fi, statErr := os.Stat(tmp); statErr == nil {
|
||||
if fi.Size() < blob.Size {
|
||||
existingSize = fi.Size()
|
||||
} else if fi.Size() > blob.Size {
|
||||
// .tmp larger than expected — discard
|
||||
os.Remove(tmp)
|
||||
}
|
||||
// fi.Size() == blob.Size handled in save (hash check + rename)
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
if t := d.authToken(); t != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+t)
|
||||
}
|
||||
}
|
||||
if existingSize > 0 {
|
||||
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", existingSize))
|
||||
}
|
||||
|
||||
resp, err := d.client.Do(req)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer func() { io.Copy(io.Discard, resp.Body); resp.Body.Close() }()
|
||||
|
||||
switch resp.StatusCode {
|
||||
case http.StatusOK:
|
||||
// Full response — reset any partial state
|
||||
existingSize = 0
|
||||
case http.StatusPartialContent:
|
||||
// Resume succeeded
|
||||
default:
|
||||
return 0, fmt.Errorf("status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
return d.save(ctx, blob, resp.Body, existingSize)
|
||||
}
|
||||
|
||||
func (d *downloader) save(ctx context.Context, blob Blob, r io.Reader, existingSize int64) (int64, error) {
|
||||
dest := filepath.Join(d.destDir, digestToPath(blob.Digest))
|
||||
tmp := dest + ".tmp"
|
||||
os.MkdirAll(filepath.Dir(dest), 0o755)
|
||||
|
||||
h := sha256.New()
|
||||
|
||||
var f *os.File
|
||||
var err error
|
||||
|
||||
if existingSize > 0 {
|
||||
// Resume — re-hash existing partial data, then append
|
||||
f, err = os.OpenFile(tmp, os.O_RDWR, 0o644)
|
||||
if err != nil {
|
||||
// Can't open partial file, start fresh
|
||||
existingSize = 0
|
||||
} else {
|
||||
// Hash the existing data
|
||||
if _, hashErr := io.CopyN(h, f, existingSize); hashErr != nil {
|
||||
f.Close()
|
||||
os.Remove(tmp)
|
||||
existingSize = 0
|
||||
} else {
|
||||
// Report resumed bytes as progress
|
||||
d.progress.add(existingSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if existingSize == 0 {
|
||||
f, err = os.Create(tmp)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
setSparse(f)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
n, err := d.copy(ctx, f, r, h)
|
||||
if err != nil {
|
||||
// Don't remove .tmp here — download() handles cleanup based on blob size
|
||||
return existingSize + n, err
|
||||
}
|
||||
f.Close()
|
||||
|
||||
if got := fmt.Sprintf("sha256:%x", h.Sum(nil)); got != blob.Digest {
|
||||
os.Remove(tmp)
|
||||
return existingSize + n, fmt.Errorf("digest mismatch")
|
||||
}
|
||||
totalWritten := existingSize + n
|
||||
if totalWritten != blob.Size {
|
||||
os.Remove(tmp)
|
||||
return totalWritten, fmt.Errorf("size mismatch")
|
||||
}
|
||||
return totalWritten, os.Rename(tmp, dest)
|
||||
}
|
||||
|
||||
func (d *downloader) copy(ctx context.Context, dst io.Writer, src io.Reader, h io.Writer) (int64, error) {
|
||||
var n int64
|
||||
var lastRead atomic.Int64
|
||||
lastRead.Store(time.Now().UnixNano())
|
||||
start := time.Now()
|
||||
|
||||
ctx, cancel := context.WithCancelCause(ctx)
|
||||
defer cancel(nil)
|
||||
|
||||
go func() {
|
||||
tick := time.NewTicker(time.Second)
|
||||
defer tick.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-tick.C:
|
||||
if time.Since(time.Unix(0, lastRead.Load())) > d.stallTimeout {
|
||||
cancel(errStalled)
|
||||
return
|
||||
}
|
||||
if e := time.Since(start); e > 5*time.Second {
|
||||
if m := d.speeds.median(); m > 0 && float64(n)/e.Seconds() < m*0.1 {
|
||||
cancel(errSlow)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
if err := ctx.Err(); err != nil {
|
||||
if c := context.Cause(ctx); c != nil {
|
||||
return n, c
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
nr, err := src.Read(buf)
|
||||
if nr > 0 {
|
||||
lastRead.Store(time.Now().UnixNano())
|
||||
dst.Write(buf[:nr])
|
||||
h.Write(buf[:nr])
|
||||
d.progress.add(int64(nr))
|
||||
n += int64(nr)
|
||||
}
|
||||
if err == io.EOF {
|
||||
return n, nil
|
||||
}
|
||||
if err != nil {
|
||||
return n, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// resolve follows redirects to find the final download URL.
|
||||
// Uses GET (not HEAD) because registries may return 200 for HEAD without
|
||||
// redirecting to CDN, while GET triggers the actual CDN redirect.
|
||||
func (d *downloader) resolve(ctx context.Context, rawURL string) (*url.URL, error) {
|
||||
u, _ := url.Parse(rawURL)
|
||||
for range 10 {
|
||||
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
req.Header.Set("User-Agent", d.userAgent)
|
||||
prev := d.authToken()
|
||||
if prev != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+prev)
|
||||
}
|
||||
|
||||
resp, err := d.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Drain body before close to enable HTTP connection reuse
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
switch resp.StatusCode {
|
||||
case http.StatusOK:
|
||||
return u, nil
|
||||
case http.StatusUnauthorized:
|
||||
if d.getToken == nil {
|
||||
return nil, fmt.Errorf("unauthorized")
|
||||
}
|
||||
ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate"))
|
||||
if err := d.refreshToken(ctx, ch, prev); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case http.StatusTemporaryRedirect, http.StatusFound, http.StatusMovedPermanently:
|
||||
loc, _ := resp.Location()
|
||||
if loc.Host != u.Host {
|
||||
return loc, nil
|
||||
}
|
||||
u = loc
|
||||
default:
|
||||
return nil, fmt.Errorf("status %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("too many redirects")
|
||||
}
|
||||
|
||||
type speedTracker struct {
|
||||
mu sync.Mutex
|
||||
speeds []float64
|
||||
}
|
||||
|
||||
func (s *speedTracker) record(v float64) {
|
||||
s.mu.Lock()
|
||||
s.speeds = append(s.speeds, v)
|
||||
if len(s.speeds) > 30 {
|
||||
s.speeds = s.speeds[1:]
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *speedTracker) median() float64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if len(s.speeds) < 5 {
|
||||
return 0
|
||||
}
|
||||
sorted := make([]float64, len(s.speeds))
|
||||
copy(sorted, s.speeds)
|
||||
slices.Sort(sorted)
|
||||
return sorted[len(sorted)/2]
|
||||
}
|
||||
|
||||
const defaultStallTimeout = 10 * time.Second
|
||||
12
x/transfer/sparse_other.go
Normal file
12
x/transfer/sparse_other.go
Normal file
@@ -0,0 +1,12 @@
|
||||
//go:build !windows
|
||||
|
||||
package transfer
|
||||
|
||||
import "os"
|
||||
|
||||
// setSparse is a no-op on non-Windows platforms.
|
||||
// On Windows, this sets the FSCTL_SET_SPARSE attribute which allows the OS
|
||||
// to not allocate disk blocks for zero-filled regions. This is useful for
|
||||
// partial downloads where not all data has been written yet. On Unix-like
|
||||
// systems, filesystems typically handle this automatically (sparse by default).
|
||||
func setSparse(_ *os.File) {}
|
||||
31
x/transfer/sparse_windows.go
Normal file
31
x/transfer/sparse_windows.go
Normal file
@@ -0,0 +1,31 @@
|
||||
//go:build windows
|
||||
|
||||
package transfer
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// setSparse sets the FSCTL_SET_SPARSE attribute on Windows files.
|
||||
// This allows the OS to not allocate disk blocks for zero-filled regions,
|
||||
// which is useful for large files that may not be fully written (e.g., partial
|
||||
// downloads). Without this, Windows may pre-allocate disk space for the full
|
||||
// file size even if most of it is zeros.
|
||||
//
|
||||
// Note: Errors are intentionally ignored because:
|
||||
// 1. The file will still work correctly without sparse support
|
||||
// 2. Not all Windows filesystems support sparse files (e.g., FAT32)
|
||||
// 3. This is an optimization, not a requirement
|
||||
func setSparse(file *os.File) {
|
||||
var bytesReturned uint32
|
||||
_ = windows.DeviceIoControl(
|
||||
windows.Handle(file.Fd()),
|
||||
windows.FSCTL_SET_SPARSE,
|
||||
nil, 0,
|
||||
nil, 0,
|
||||
&bytesReturned,
|
||||
nil,
|
||||
)
|
||||
}
|
||||
236
x/transfer/transfer.go
Normal file
236
x/transfer/transfer.go
Normal file
@@ -0,0 +1,236 @@
|
||||
// Package transfer provides minimal, fast blob transfer for tensor-based models.
|
||||
//
|
||||
// This package is in x/ because the tensor model storage format is under development.
|
||||
// It provides optimized transfer for models with many small blobs (tensor models)
|
||||
// rather than few large blobs (typical LLMs).
|
||||
//
|
||||
// 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:
|
||||
//
|
||||
// - 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 resumable downloads with JSON state files, async hashing from OS page
|
||||
// cache, and per-part speed tracking with rolling median.
|
||||
package transfer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"math/rand/v2"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Blob represents a content-addressed blob to transfer.
|
||||
type Blob struct {
|
||||
Digest string // sha256:...
|
||||
Size int64
|
||||
|
||||
// From enables cross-repository blob mounting (upload only).
|
||||
// When set, the upload will first attempt to mount the blob from this source
|
||||
// repository instead of uploading the data. This is a Docker Registry v2 API
|
||||
// feature that avoids re-uploading blobs that already exist elsewhere.
|
||||
//
|
||||
// Example: From="library/source-model" will add ?mount=<digest>&from=library/source-model
|
||||
// to the POST /blobs/uploads/ request. If the registry returns 201 Created,
|
||||
// the blob was mounted successfully and no upload is needed.
|
||||
//
|
||||
// See: https://distribution.github.io/distribution/spec/api/#cross-repository-blob-mount
|
||||
From string
|
||||
}
|
||||
|
||||
// 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 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 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
|
||||
ManifestRef string // Tag or digest for the manifest (e.g., "latest", "sha256:...")
|
||||
Repository string // Repository path for manifest URL (e.g., "library/model")
|
||||
}
|
||||
|
||||
// AuthChallenge represents a parsed WWW-Authenticate challenge.
|
||||
type AuthChallenge struct {
|
||||
Realm string
|
||||
Service string
|
||||
Scope string
|
||||
}
|
||||
|
||||
// Default concurrency limits and settings
|
||||
const (
|
||||
DefaultDownloadConcurrency = 64
|
||||
DefaultUploadConcurrency = 64
|
||||
maxRetries = 6
|
||||
defaultUserAgent = "ollama-transfer/1.0"
|
||||
|
||||
// resumeThreshold is the minimum blob size for resume support.
|
||||
// Only blobs above this size keep partial .tmp files on failure.
|
||||
resumeThreshold = 64 << 20 // 64 MB
|
||||
|
||||
// smallBlobSpeedThreshold is the size below which speed samples are skipped,
|
||||
// since their transfer time is dominated by HTTP overhead, not throughput.
|
||||
smallBlobSpeedThreshold = 100 << 10 // 100 KB
|
||||
)
|
||||
|
||||
var errMaxRetriesExceeded = errors.New("max retries exceeded")
|
||||
|
||||
// defaultClient is a shared HTTP client with connection pooling.
|
||||
var defaultClient = &http.Client{
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConns: 100,
|
||||
MaxIdleConnsPerHost: 100,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
},
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
|
||||
// progressTracker aggregates progress across concurrent operations.
|
||||
type progressTracker struct {
|
||||
completed atomic.Int64
|
||||
total int64
|
||||
callback func(completed, total int64)
|
||||
}
|
||||
|
||||
func newProgressTracker(total int64, callback func(completed, total int64)) *progressTracker {
|
||||
return &progressTracker{
|
||||
total: total,
|
||||
callback: callback,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *progressTracker) add(n int64) {
|
||||
if p == nil || p.callback == nil {
|
||||
return
|
||||
}
|
||||
completed := p.completed.Add(n)
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
slog.Debug("progress callback panic (likely closed channel)", "recovered", r)
|
||||
}
|
||||
}()
|
||||
p.callback(completed, p.total)
|
||||
}
|
||||
|
||||
// Download downloads blobs in parallel with streaming hash verification.
|
||||
func Download(ctx context.Context, opts DownloadOptions) error {
|
||||
return download(ctx, opts)
|
||||
}
|
||||
|
||||
// Upload uploads blobs in parallel.
|
||||
func Upload(ctx context.Context, opts UploadOptions) error {
|
||||
return upload(ctx, opts)
|
||||
}
|
||||
|
||||
// digestToPath converts sha256:abc123 to sha256-abc123
|
||||
func digestToPath(digest string) string {
|
||||
if len(digest) > 7 && digest[6] == ':' {
|
||||
return digest[:6] + "-" + digest[7:]
|
||||
}
|
||||
return digest
|
||||
}
|
||||
|
||||
// parseAuthChallenge parses a WWW-Authenticate header value.
|
||||
// Example: Bearer realm="https://auth.example.com",service="registry",scope="repository:foo:pull"
|
||||
func parseAuthChallenge(header string) AuthChallenge {
|
||||
header = strings.TrimPrefix(header, "Bearer ")
|
||||
|
||||
getValue := func(key string) string {
|
||||
startIdx := strings.Index(header, key+"=")
|
||||
if startIdx == -1 {
|
||||
return ""
|
||||
}
|
||||
startIdx += len(key) + 1
|
||||
if startIdx >= len(header) {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Handle quoted values
|
||||
if header[startIdx] == '"' {
|
||||
startIdx++
|
||||
endIdx := strings.Index(header[startIdx:], "\"")
|
||||
if endIdx == -1 {
|
||||
return header[startIdx:]
|
||||
}
|
||||
return header[startIdx : startIdx+endIdx]
|
||||
}
|
||||
|
||||
// Unquoted value - ends at comma or end of string
|
||||
endIdx := strings.Index(header[startIdx:], ",")
|
||||
if endIdx == -1 {
|
||||
return header[startIdx:]
|
||||
}
|
||||
return header[startIdx : startIdx+endIdx]
|
||||
}
|
||||
|
||||
return AuthChallenge{
|
||||
Realm: getValue("realm"),
|
||||
Service: getValue("service"),
|
||||
Scope: getValue("scope"),
|
||||
}
|
||||
}
|
||||
|
||||
// backoff returns a function that sleeps with exponential backoff.
|
||||
func backoff(ctx context.Context, attempt int, maxBackoff time.Duration) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
// n^2 backoff with jitter
|
||||
d := min(time.Duration(attempt*attempt)*10*time.Millisecond, maxBackoff)
|
||||
d = time.Duration(float64(d) * (rand.Float64() + 0.5))
|
||||
|
||||
t := time.NewTimer(d)
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-t.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
2909
x/transfer/transfer_test.go
Normal file
2909
x/transfer/transfer_test.go
Normal file
File diff suppressed because it is too large
Load Diff
813
x/transfer/upload.go
Normal file
813
x/transfer/upload.go
Normal file
@@ -0,0 +1,813 @@
|
||||
package transfer
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"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"
|
||||
|
||||
"golang.org/x/sync/errgroup"
|
||||
"golang.org/x/sync/semaphore"
|
||||
)
|
||||
|
||||
type uploader struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
srcDir string
|
||||
repository string // Repository path for blob URLs (e.g., "library/model")
|
||||
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 {
|
||||
if len(opts.Blobs) == 0 && len(opts.Manifest) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
u := &uploader{
|
||||
client: cmp.Or(opts.Client, defaultClient),
|
||||
baseURL: opts.BaseURL,
|
||||
srcDir: opts.SrcDir,
|
||||
repository: cmp.Or(opts.Repository, "library/_"),
|
||||
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 {
|
||||
// Discover which blobs the server already has so we can skip uploading
|
||||
needsUpload := make([]bool, len(opts.Blobs))
|
||||
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
|
||||
var toUpload []Blob
|
||||
var totalSize, alreadyExists int64
|
||||
for i, blob := range opts.Blobs {
|
||||
totalSize += blob.Size
|
||||
if needsUpload[i] {
|
||||
toUpload = append(toUpload, blob)
|
||||
} else {
|
||||
alreadyExists += blob.Size
|
||||
}
|
||||
}
|
||||
|
||||
// Progress includes all blobs — already-existing ones start as completed
|
||||
u.progress = newProgressTracker(totalSize, opts.Progress)
|
||||
u.progress.add(alreadyExists)
|
||||
|
||||
logutil.Trace("upload plan", "total_blobs", len(opts.Blobs), "need_upload", len(toUpload), "already_exist", len(opts.Blobs)-len(toUpload), "total_bytes", totalSize, "existing_bytes", alreadyExists)
|
||||
|
||||
if len(toUpload) == 0 {
|
||||
if u.logger != nil {
|
||||
u.logger.Debug("all blobs exist, nothing to upload")
|
||||
}
|
||||
} else {
|
||||
// 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 {
|
||||
if err := sem.Acquire(gctx, 1); err != nil {
|
||||
return err
|
||||
}
|
||||
defer sem.Release(1)
|
||||
return u.upload(gctx, blob)
|
||||
})
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(opts.Manifest) > 0 && opts.ManifestRef != "" && opts.Repository != "" {
|
||||
logutil.Trace("pushing manifest", "repo", opts.Repository, "ref", opts.ManifestRef, "size", len(opts.Manifest))
|
||||
if err := u.pushManifest(ctx, opts.Repository, opts.ManifestRef, opts.Manifest); err != nil {
|
||||
logutil.Trace("manifest push failed", "error", err)
|
||||
return err
|
||||
}
|
||||
logutil.Trace("manifest push succeeded", "repo", opts.Repository, "ref", opts.ManifestRef)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *uploader) upload(ctx context.Context, blob Blob) error {
|
||||
var lastErr error
|
||||
var n int64
|
||||
|
||||
for attempt := range maxRetries {
|
||||
if attempt > 0 {
|
||||
// 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
|
||||
}
|
||||
}
|
||||
|
||||
var err error
|
||||
n, err = u.uploadOnce(ctx, blob)
|
||||
if err == nil {
|
||||
logutil.Trace("blob upload complete", "digest", blob.Digest, "bytes", n, "attempt", attempt+1)
|
||||
return nil
|
||||
}
|
||||
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
logutil.Trace("blob upload cancelled", "digest", blob.Digest, "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
logutil.Trace("blob upload failed, retrying", "digest", blob.Digest, "attempt", attempt+1, "bytes", n, "error", err)
|
||||
u.progress.add(-n)
|
||||
lastErr = err
|
||||
}
|
||||
return fmt.Errorf("%w: %v", errMaxRetriesExceeded, lastErr)
|
||||
}
|
||||
|
||||
func (u *uploader) uploadOnce(ctx context.Context, blob Blob) (int64, error) {
|
||||
if u.logger != nil {
|
||||
u.logger.Debug("uploading blob", "digest", blob.Digest, "size", blob.Size)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
f, err := os.Open(filepath.Join(u.srcDir, digestToPath(blob.Digest)))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
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)
|
||||
prev := u.authToken()
|
||||
if prev != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+prev)
|
||||
}
|
||||
|
||||
resp, err := u.client.Do(req)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == http.StatusUnauthorized && u.getToken != nil {
|
||||
ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate"))
|
||||
if err := u.refreshToken(ctx, ch, prev); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return u.exists(ctx, blob)
|
||||
}
|
||||
|
||||
return resp.StatusCode == http.StatusOK, nil
|
||||
}
|
||||
|
||||
const maxInitRetries = 12
|
||||
|
||||
// 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)
|
||||
endpoint.RawQuery = q.Encode()
|
||||
|
||||
var lastErr error
|
||||
for attempt := range maxInitRetries {
|
||||
if attempt > 0 {
|
||||
if err := backoff(ctx, attempt, min(5*time.Second<<uint(attempt-1), 30*time.Second)); err != nil {
|
||||
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)
|
||||
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)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == http.StatusUnauthorized && u.getToken != nil {
|
||||
ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate"))
|
||||
if err := u.refreshToken(ctx, ch, prev); err != nil {
|
||||
return uploadEndpoint{}, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusCreated {
|
||||
// Server matched our ?digest= against existing storage —
|
||||
// nothing to upload.
|
||||
return uploadEndpoint{}, nil
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusAccepted {
|
||||
lastErr = fmt.Errorf("init upload: status %d", resp.StatusCode)
|
||||
continue
|
||||
}
|
||||
|
||||
loc := resp.Header.Get("Docker-Upload-Location")
|
||||
if loc == "" {
|
||||
loc = resp.Header.Get("Location")
|
||||
}
|
||||
if loc == "" {
|
||||
lastErr = fmt.Errorf("no upload location (server returned 202 without Location header)")
|
||||
continue
|
||||
}
|
||||
|
||||
sessionURL, _ := url.Parse(loc)
|
||||
if !sessionURL.IsAbs() {
|
||||
base, _ := url.Parse(u.baseURL)
|
||||
sessionURL = base.ResolveReference(sessionURL)
|
||||
}
|
||||
|
||||
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 uploadEndpoint{}, lastErr
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
splitParts := computeParts
|
||||
if u.makeParts != nil {
|
||||
splitParts = u.makeParts
|
||||
}
|
||||
parts := splitParts(blob.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)
|
||||
prev := u.authToken()
|
||||
if prev != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+prev)
|
||||
}
|
||||
|
||||
resp, err := u.client.Do(req)
|
||||
if err != nil {
|
||||
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() }()
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
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 nil, nil, cdnN, err
|
||||
}
|
||||
return next, cdnSum, cdnN, nil
|
||||
|
||||
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 nil, nil, pr.bytes(), fmt.Errorf("patch part %d: auth retry", part.n)
|
||||
|
||||
case resp.StatusCode >= http.StatusBadRequest:
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, nil, pr.bytes(), fmt.Errorf("patch part %d: status %d: %s", part.n, resp.StatusCode, body)
|
||||
}
|
||||
|
||||
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)
|
||||
prev := u.authToken()
|
||||
if prev != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+prev)
|
||||
}
|
||||
|
||||
resp, err := u.client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { io.Copy(io.Discard, resp.Body); resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode == http.StatusUnauthorized && u.getToken != nil {
|
||||
ch := parseAuthChallenge(resp.Header.Get("WWW-Authenticate"))
|
||||
if err := u.refreshToken(ctx, ch, prev); err != nil {
|
||||
return err
|
||||
}
|
||||
return u.pushManifest(ctx, repo, ref, manifest)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
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 atomic.Int64
|
||||
}
|
||||
|
||||
func (r *progressReader) Read(p []byte) (int, error) {
|
||||
n, err := r.reader.Read(p)
|
||||
if n > 0 {
|
||||
r.n.Add(int64(n))
|
||||
r.tracker.add(int64(n))
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *progressReader) bytes() int64 {
|
||||
return r.n.Load()
|
||||
}
|
||||
Reference in New Issue
Block a user