diff --git a/CLAUDE.md b/CLAUDE.md index 40a4321..0990221 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -209,6 +209,7 @@ See `config-appview.example.yaml` and `config-hold.example.yaml` for all options - **Hold stats are ATProto records in CAR store** — `io.atcr.hold.stats` records are stored via `repomgr.PutRecord()`, not in SQLite. Lost if CAR store is lost without backup. - **PLC auto-update on boot** — When using did:plc, `LoadOrCreateDID()` calls `EnsurePLCCurrent()` every startup. If local signing key or URL doesn't match plc.directory, it auto-updates (requires rotation key on disk). - **Hold CAR store is the source of truth** — Captain, crew, layer, stats, scan records, Bluesky posts, profiles are all ATProto records in the CAR store. SQLite holds only the records index and events. +- **Blob upload buffers are bounded process-wide** — `server.upload_buffer_budget_mb` (default 512) caps memory held in upload buffers; a push that would exceed it blocks in `Write` as backpressure on the client. A per-instance sweeper (not leased, 5-minute interval) cancels uploads idle past `server.upload_idle_timeout` (default 1h), aborting the hold-side multipart. Both live in `pkg/appview/storage/upload_budget.go`. - **Admin panel long-running loops MUST use background jobs** — Any admin handler that loops over many items doing per-item PDS writes/network calls (bulk crew tier remap, crew import, scan backfill) must NOT run synchronously on the request context. The reverse proxy 504s such a request around the 10s mark and cancels `r.Context()`, aborting the loop mid-flight (`getb tx, context canceled`). Use the `startJob` helper in `pkg/hold/admin/jobs.go`: do fast validation synchronously (return a 200 fragment on failure, never a 302 — the forms are htmx-driven), then kick off the loop via `ui.startJob(key, title, resultTemplate, timeout, fn)` and render `partials/job_progress.html`. The fragment polls `/admin/api/jobs/{key}/status` until done. `fn` runs under its own detached `context.Background()` timeout, so it survives the request ending. The GC subsystem (`gc.startBackground`) is a peer implementation in the `gc` package (which must not import `admin`). ## Common Tasks diff --git a/config-appview.example.yaml b/config-appview.example.yaml index a4707cd..b25eac1 100644 --- a/config-appview.example.yaml +++ b/config-appview.example.yaml @@ -36,6 +36,10 @@ server: # DIDs of holds this appview manages billing for (REQUIRED). The first entry is the default blob-storage hold. Tier updates are pushed to these holds. managed_holds: - did:web:172.28.0.3%3A8080 + # Process-wide ceiling, in MB, on memory held in blob upload buffers. Pushes block (backpressure) rather than exceeding it. Raised to 16 MB (one buffer) if set lower. + upload_buffer_budget_mb: 512 + # How long a blob upload may go without a write before it is treated as abandoned and cancelled. Measures inactivity, not age, so a slow push is never reaped. + upload_idle_timeout: 1h0m0s # Web UI settings. ui: # SQLite/libSQL database for OAuth sessions, stars, pull counts, and device approvals. diff --git a/deploy/upcloud/configs/appview.yaml.tmpl b/deploy/upcloud/configs/appview.yaml.tmpl index 0b67ebb..a2d4766 100644 --- a/deploy/upcloud/configs/appview.yaml.tmpl +++ b/deploy/upcloud/configs/appview.yaml.tmpl @@ -19,6 +19,8 @@ server: - "seamark.cr" managed_holds: - "{{.HoldDid}}" + upload_buffer_budget_mb: 512 + upload_idle_timeout: 1h0m0s ui: database_path: "{{.BasePath}}/ui.db" theme: seamark diff --git a/docs/appview.md b/docs/appview.md index fa6f22d..5e95310 100644 --- a/docs/appview.md +++ b/docs/appview.md @@ -137,7 +137,7 @@ jetstream.backfill_enabled → ATCR_JETSTREAM_BACKFILL_ENABLED | Section | Purpose | Notes | |---------|---------|-------| -| `server` | Listen address, public URL, managed holds, branding | Only `managed_holds` is required | +| `server` | Listen address, public URL, managed holds, branding, blob upload limits | Only `managed_holds` is required | | `ui` | Database path, theme, libSQL sync | All have defaults; auto-creates DB on first run | | `auth` | JWT signing key/cert paths | Auto-generated on first run | | `jetstream` | Real-time ATProto event streaming, backfill sync | Runs automatically; backfill enabled by default | @@ -145,6 +145,25 @@ jetstream.backfill_enabled → ATCR_JETSTREAM_BACKFILL_ENABLED | `log_shipper` | Remote log shipping (Victoria, OpenSearch, Loki) | Disabled by default | | `legal` | Terms/privacy page customization | Optional | +### Blob upload memory + +Each in-flight blob upload buffers up to 16MB in the AppView process, and Docker +pushes several layers at once per client, so concurrent pushes are bounded by two +`server` settings: + +| Field | Default | Purpose | +|-------|---------|---------| +| `upload_buffer_budget_mb` | `512` | Process-wide ceiling on memory held in upload buffers. A push that would exceed it blocks until another upload finishes, which is backpressure on the Docker client rather than an error. Raised to 16MB (one buffer) if configured lower, since a smaller budget could never satisfy a single upload. | +| `upload_idle_timeout` | `1h` | How long an upload may go without a write before it is treated as abandoned. | + +A background sweeper runs every 5 minutes on every instance (it is deliberately +not leased: the uploads it tracks are per-process). Anything idle past +`upload_idle_timeout` is cancelled: its buffer and budget are released, its +hold-side S3 multipart upload is aborted, and the client gets +`BLOB_UPLOAD_UNKNOWN` if it ever comes back, which makes Docker restart the +layer. Inactivity is the signal, not age, so a slow push that is still making +progress is never reaped. + ### Auto-generated files On first run (and each boot), AppView auto-generates these under `/var/lib/atcr/`: diff --git a/go.mod b/go.mod index a2af975..5f644b7 100644 --- a/go.mod +++ b/go.mod @@ -50,6 +50,7 @@ require ( golang.org/x/crypto v0.55.0 golang.org/x/image v0.45.0 golang.org/x/net v0.58.0 + golang.org/x/sync v0.22.0 golang.org/x/sys v0.47.0 golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da oras.land/oras-go/v2 v2.6.2 @@ -205,7 +206,6 @@ require ( go.uber.org/zap v1.28.0 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect golang.org/x/exp v0.0.0-20260824195058-e88cd73687aa // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/text v0.41.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/tools v0.49.0 // indirect diff --git a/pkg/appview/config.go b/pkg/appview/config.go index af0a069..29209e4 100644 --- a/pkg/appview/config.go +++ b/pkg/appview/config.go @@ -70,6 +70,15 @@ type ServerConfig struct { // DIDs of holds this appview manages billing for. The first entry is also // the default blob-storage hold (used when a user has no hold selected). ManagedHolds []string `yaml:"managed_holds" comment:"DIDs of holds this appview manages billing for (REQUIRED). The first entry is the default blob-storage hold. Tier updates are pushed to these holds."` + + // Process-wide ceiling on memory held in blob upload buffers. Each + // in-flight layer buffers up to 16MB, and Docker pushes five layers at a + // time per client, so without a ceiling concurrent pushes are unbounded. + // Writers block waiting for budget, which is backpressure on the client. + UploadBufferBudgetMB int `yaml:"upload_buffer_budget_mb" comment:"Process-wide ceiling, in MB, on memory held in blob upload buffers. Pushes block (backpressure) rather than exceeding it. Raised to 16 MB (one buffer) if set lower."` + + // How long a blob upload may sit idle before it is reaped. + UploadIdleTimeout time.Duration `yaml:"upload_idle_timeout" comment:"How long a blob upload may go without a write before it is treated as abandoned and cancelled. Measures inactivity, not age, so a slow push is never reaped."` } // PrimaryHoldDID returns the appview's default blob-storage hold, which is the @@ -245,6 +254,8 @@ func setDefaults(v *viper.Viper) { v.SetDefault("server.client_short_name", "ATCR") v.SetDefault("server.registry_domains", []string{}) v.SetDefault("server.managed_holds", []string{}) + v.SetDefault("server.upload_buffer_budget_mb", 512) + v.SetDefault("server.upload_idle_timeout", "1h") // UI defaults v.SetDefault("ui.database_path", "/var/lib/atcr/ui.db") diff --git a/pkg/appview/server.go b/pkg/appview/server.go index efee91d..8214cda 100644 --- a/pkg/appview/server.go +++ b/pkg/appview/server.go @@ -194,6 +194,13 @@ func NewAppViewServer(cfg *Config, branding *BrandingOverrides) (*AppViewServer, s.healthWorker.Start(workerCtx) slog.Info("Hold health worker started", "startup_delay", startupDelay, "refresh_interval", cfg.Health.CheckInterval, "cache_ttl", cfg.Health.CacheTTL) + // Blob upload buffers are bounded process-wide, and abandoned uploads are + // swept. Both are per-process concerns (the writers live in package state + // in storage, because distribution builds a new blob store per request), so + // the sweeper runs on every instance and is deliberately not leased. + storage.ConfigureUploads(int64(cfg.Server.UploadBufferBudgetMB)*1024*1024, cfg.Server.UploadIdleTimeout) + storage.StartUploadSweeper(workerCtx) + // Leader election for the singleton background workers. The health worker // above is deliberately not among them: it only refreshes a cache that each // instance needs locally, so running it everywhere is correct. diff --git a/pkg/appview/storage/proxy_blob_store.go b/pkg/appview/storage/proxy_blob_store.go index 6da2473..1b4e3fa 100644 --- a/pkg/appview/storage/proxy_blob_store.go +++ b/pkg/appview/storage/proxy_blob_store.go @@ -381,15 +381,18 @@ func (p *ProxyBlobStore) Create(ctx context.Context, options ...distribution.Blo // Generate unique writer ID writerID := fmt.Sprintf("upload-%d", time.Now().UnixNano()) + now := time.Now() writer := &ProxyBlobWriter{ - store: p, - options: opts, - parts: make([]CompletedPart, 0), - partNumber: 1, - buffer: &bytes.Buffer{}, - digester: digest.Canonical.Digester(), - id: writerID, - startedAt: time.Now(), + store: p, + options: opts, + parts: make([]CompletedPart, 0), + partNumber: 1, + buffer: &bytes.Buffer{}, + digester: digest.Canonical.Digester(), + id: writerID, + startedAt: now, + lastActivity: now, + requestCtx: ctx, } // Store in global uploads map for resume support @@ -408,9 +411,17 @@ func (p *ProxyBlobStore) Resume(ctx context.Context, id string) (distribution.Bl globalUploadsMu.RUnlock() if !ok { + // Also what a client sees after the sweeper reaped an abandoned upload: + // distribution turns this into BLOB_UPLOAD_UNKNOWN and the client + // starts the layer over. return nil, distribution.ErrBlobUploadUnknown } + // This request now owns the writer: hand it the context to block on, and + // count the resume as activity so a long push made of many PATCHes is never + // mistaken for an abandoned one. + writer.adopt(ctx) + // Just return the writer - parts are buffered and flushed on demand return writer, nil } @@ -674,6 +685,28 @@ type ProxyBlobWriter struct { closed bool id string // Distribution's upload ID (for state) startedAt time.Time + + // mu guards everything above against the abandoned-upload sweeper, which is + // the only thing that ever touches a writer concurrently with the request + // that owns it. Distribution hands a given upload to one request at a time, + // so Write, Commit and Cancel never contend with each other. + mu sync.Mutex + + // charged is the buffer budget this writer currently holds, in bytes. It is + // the single source of truth for release: every exit path releases exactly + // this and zeroes it, so nothing can be released twice. + charged int64 + + // lastActivity is when the writer last did something on the client's + // behalf. startedAt is the wrong signal for the sweeper: a slow but healthy + // push can run for hours, while an abandoned one goes quiet immediately. + lastActivity time.Time + + // requestCtx is the context of the request currently driving this writer, + // set by Create and by every Resume. Write has no context of its own (it is + // an io.Writer), and it can block waiting for budget, so it needs one: + // without it a client that has already hung up keeps a slot in the queue. + requestCtx context.Context } // ID returns the upload ID @@ -686,32 +719,179 @@ func (w *ProxyBlobWriter) StartedAt() time.Time { return w.startedAt } +// adopt records the request now driving this writer and counts as activity. +// Called on Create and on every Resume, so a client that keeps coming back for +// the next PATCH is never mistaken for an abandoned one. +func (w *ProxyBlobWriter) adopt(ctx context.Context) { + w.mu.Lock() + defer w.mu.Unlock() + w.requestCtx = ctx + w.lastActivity = time.Now() +} + +// waitContext is the context a blocking budget acquire should honour: the +// request currently driving the writer, or the background if there is none. +// Callers hold w.mu. +func (w *ProxyBlobWriter) waitContext() context.Context { + if w.requestCtx != nil { + return w.requestCtx + } + return context.Background() +} + +// projectedCap is the backing array size the buffer will hold once an n byte +// write has landed. It mirrors growBuffer plus bytes.Buffer's own doubling: the +// array either already fits the write, grows to the doubled size, or jumps to +// the threshold. Nothing is ever sized past the threshold: Write splits a +// larger slice into threshold-sized parts. +func (w *ProxyBlobWriter) projectedCap(n int) int64 { + have := int64(w.buffer.Cap()) + // Write fills the buffer to the threshold and flushes before taking more, + // so no write, however large, needs the buffer to grow past it. + need := min(int64(w.buffer.Len())+int64(n), maxBufferSize) + if need <= have { + return have + } + return min(max(need, 2*have), maxBufferSize) +} + +// budgetDelta is the extra budget the writer must acquire before an n byte +// write can land. +// +// What is charged is the buffer's backing array, not the bytes written: a +// writer holds exactly one buffer's worth at a time, whatever it has grown to. +// It is never released between parts, because bytes.Buffer.Reset keeps the +// array: releasing there would report memory as free while the writer still +// holds every byte of it. +// +// Callers hold w.mu. +func (w *ProxyBlobWriter) budgetDelta(n int) (int64, error) { + want := w.projectedCap(n) + if want <= w.charged { + return 0, nil + } + if want > uploadBudget.limit { + // Only reachable for a single write larger than the whole budget, which + // means a caller handing the writer an entire oversized blob in one + // call rather than streaming it. The semaphore would never grant this, + // so say so now instead of waiting out the acquire's deadline. + return 0, fmt.Errorf("upload buffer of %d bytes exceeds the process-wide budget of %d", want, uploadBudget.limit) + } + return want - w.charged, nil +} + +// releaseBudget returns everything this writer holds. Safe to call more than +// once: the second call has nothing to release. Callers hold w.mu. +func (w *ProxyBlobWriter) releaseBudget() { + uploadBudget.release(w.charged) + w.charged = 0 +} + +// reapIfIdle cancels the writer if it has been inactive for at least ttl, +// reporting how idle it was and whether it was reaped. The hold-side multipart +// is aborted on a detached context: there is no request left to borrow one from, +// and the request that opened this upload is long gone. +func (w *ProxyBlobWriter) reapIfIdle(now time.Time, ttl time.Duration) (time.Duration, bool) { + // A writer someone is actively inside is by definition not abandoned, and + // Write holds this lock across a part upload, which can take minutes. Skip + // it and look again on the next sweep rather than blocking every other + // writer's reap behind one live upload. + if !w.mu.TryLock() { + return 0, false + } + defer w.mu.Unlock() + + idle := now.Sub(w.lastActivity) + if w.closed || idle < ttl { + return 0, false + } + + w.closed = true + + ctx, cancel := context.WithTimeout(context.Background(), uploadAbortTimeout) + defer cancel() + w.abortIfStarted(ctx) + + w.releaseBudget() + return idle, true +} + // Write writes data to the upload. // Buffers data and flushes a part once the buffer reaches maxBufferSize. func (w *ProxyBlobWriter) Write(p []byte) (int, error) { + // Work out what this write costs, then wait for it without the lock held. + // Blocking here is the point: it is backpressure on the Docker client. But + // holding the writer's lock while blocked would make Cancel and the sweeper + // queue behind a write that is waiting on memory nobody has yet returned. + w.mu.Lock() if w.closed { + w.mu.Unlock() return 0, fmt.Errorf("writer closed") } + w.lastActivity = time.Now() + delta, err := w.budgetDelta(len(p)) + waitCtx := w.waitContext() + w.mu.Unlock() - w.growBuffer(len(p)) + if err != nil { + return 0, err + } + if err := uploadBudget.acquire(waitCtx, delta); err != nil { + return 0, err + } - n, err := w.buffer.Write(p) - w.size += int64(n) - if n > 0 { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + // Cancelled or reaped while this write waited for its budget. The delta + // was never folded into w.charged, so releasing it here cannot collide + // with the release that closing the writer already did. + uploadBudget.release(delta) + return 0, fmt.Errorf("writer closed") + } + // Only Write moves w.charged upward, and distribution drives a given upload + // from one request at a time, so nothing can have charged the buffer while + // this write was waiting. + w.charged += delta + + // Never let the buffer pass the threshold. Appending a whole chunk and + // checking afterwards let the last chunk before a flush land a few bytes + // past 16MB, which does not fit the 16MB backing array, so bytes.Buffer + // doubled it to 32MB and Reset kept that for the rest of the upload. Only + // chunk sizes that tile 16MB exactly avoided it, and the network read loop + // does not promise those. Filling to exactly the threshold, flushing, and + // continuing with the remainder pins capacity at 16MB for any chunk size, + // makes every part exactly one buffer, and means a single oversized Write + // streams through as parts instead of buffering whole. + written := 0 + for written < len(p) { + chunk := p[written:] + if room := maxBufferSize - w.buffer.Len(); len(chunk) > room { + chunk = chunk[:room] + } + + w.growBuffer(len(chunk)) + + // bytes.Buffer.Write only fails by panicking on allocation, never by + // returning an error, so n is always len(chunk). + n, _ := w.buffer.Write(chunk) + w.size += int64(n) + written += n // Hash as we go. Nothing else in this path ever looked at the bytes: // Commit took the digest in the client's final PUT on trust, which made // the content address of a blob whatever the client claimed it was. - w.digester.Hash().Write(p[:n]) - } + w.digester.Hash().Write(chunk[:n]) - // Flush if buffer reaches limit (S3 part size) - if w.buffer.Len() >= maxBufferSize { - if err := w.flushPart(); err != nil { - return n, err + // Flush once the buffer is full (S3 part size) + if w.buffer.Len() >= maxBufferSize { + if err := w.flushPart(); err != nil { + return written, err + } } } - return n, err + return written, nil } // growBuffer sizes the buffer's backing array ahead of an n byte write so that @@ -824,7 +1004,14 @@ func (w *ProxyBlobWriter) flushPart() error { // ReadFrom reads from a reader func (w *ProxyBlobWriter) ReadFrom(r io.Reader) (int64, error) { - if w.closed { + // The lock is taken for this check and released again: every byte below + // goes through Write, which takes it per chunk. Holding it across the whole + // copy would pin the writer for the length of an upload and starve the + // sweeper's TryLock for just as long. + w.mu.Lock() + closed := w.closed + w.mu.Unlock() + if closed { return 0, fmt.Errorf("writer closed") } @@ -854,6 +1041,8 @@ func (w *ProxyBlobWriter) ReadFrom(r io.Reader) (int64, error) { // Size returns the current size func (w *ProxyBlobWriter) Size() int64 { + w.mu.Lock() + defer w.mu.Unlock() return w.size } @@ -864,10 +1053,19 @@ func (w *ProxyBlobWriter) Size() int64 { // the final key if it is all still buffered, otherwise a final part plus the // hold's completeUpload. func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descriptor) (distribution.Descriptor, error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.closed { return distribution.Descriptor{}, fmt.Errorf("writer closed") } w.closed = true + w.lastActivity = time.Now() + + // Every path out of Commit is terminal, so the budget goes back here and + // nowhere else. The buffer is still needed below (the direct PUT reads + // straight out of it), which is why this is deferred rather than done now. + defer w.releaseBudget() // Remove from global uploads map globalUploadsMu.Lock() @@ -987,7 +1185,11 @@ func (w *ProxyBlobWriter) abortIfStarted(ctx context.Context) { // Cancel cancels the upload by aborting the multipart upload func (w *ProxyBlobWriter) Cancel(ctx context.Context) error { + w.mu.Lock() + defer w.mu.Unlock() + w.closed = true + defer w.releaseBudget() slog.Debug("Cancelling upload", "component", "proxy_blob_store/Cancel", "id", w.id) diff --git a/pkg/appview/storage/proxy_blob_store_test.go b/pkg/appview/storage/proxy_blob_store_test.go index 9514633..1313f53 100644 --- a/pkg/appview/storage/proxy_blob_store_test.go +++ b/pkg/appview/storage/proxy_blob_store_test.go @@ -2842,3 +2842,85 @@ func TestBufferGrowth_StaysWithinThreshold(t *testing.T) { } }) } + +// TestWrite_NeverOvershootsTheThreshold pins the flush boundary. Appending a +// whole chunk and checking afterwards let the last chunk before a flush land a +// few bytes past the threshold, which bytes.Buffer answered by doubling the +// backing array to 32MB, and Reset kept that for the rest of the upload. The +// budget then undercharged every large layer by a whole buffer. Chunk sizes +// that tile 16MB exactly hid it, so this streams in 24KB chunks, and also +// hands over one slice three buffers long to prove an oversized single Write +// streams through as exact parts instead of buffering whole. +func TestWrite_NeverOvershootsTheThreshold(t *testing.T) { + cases := []struct { + name string + chunk int + total int + }{ + {"24KB chunks across the boundary", 24 * 1024, maxBufferSize + 96*1024}, + {"one slice three buffers long", 3 * maxBufferSize, 3 * maxBufferSize}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, 4*maxBufferSize) + + store := createTestProxyBlobStore(t, holdServer.URL) + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + pbw := writer.(*ProxyBlobWriter) + + data := generateTestData(tc.total) + for off := 0; off < len(data); off += tc.chunk { + end := min(off+tc.chunk, len(data)) + if _, err := writer.Write(data[off:end]); err != nil { + t.Fatalf("Write() failed: %v", err) + } + } + + if c := pbw.buffer.Cap(); c > maxBufferSize { + t.Errorf("Buffer capacity overshot the threshold: cap %d > %d", c, maxBufferSize) + } + if pbw.buffer.Len() >= maxBufferSize { + t.Errorf("Buffer was left full without a flush: len %d", pbw.buffer.Len()) + } + if pbw.charged != maxBufferSize { + t.Errorf("Expected exactly one buffer charged, got %d", pbw.charged) + } + + s3Server.mu.Lock() + wantParts := tc.total / maxBufferSize + if len(s3Server.Parts) != wantParts { + t.Errorf("Expected %d flushed parts, got %d", wantParts, len(s3Server.Parts)) + } + for num, part := range s3Server.Parts { + if len(part) != maxBufferSize { + t.Errorf("Part %d is %d bytes, expected exactly %d", num, len(part), maxBufferSize) + } + } + s3Server.mu.Unlock() + + desc, err := writer.Commit(context.Background(), distribution.Descriptor{ + Digest: digest.FromBytes(data), + Size: int64(len(data)), + }) + if err != nil { + t.Fatalf("Commit() failed: %v", err) + } + if desc.Size != int64(tc.total) { + t.Errorf("Committed size %d, expected %d", desc.Size, tc.total) + } + if held := bytesHeld(); held != 0 { + t.Errorf("Expected budget returned after Commit, got %d held", held) + } + }) + } +} diff --git a/pkg/appview/storage/upload_budget.go b/pkg/appview/storage/upload_budget.go new file mode 100644 index 0000000..0cfac77 --- /dev/null +++ b/pkg/appview/storage/upload_budget.go @@ -0,0 +1,199 @@ +package storage + +import ( + "context" + "log/slog" + "sync/atomic" + "time" + + "golang.org/x/sync/semaphore" +) + +const ( + // defaultUploadBufferBudget is the process-wide ceiling on bytes held in + // blob upload buffers. Every in-flight push holds a buffer of up to + // maxBufferSize, Docker pushes up to five layers at once per client, and + // nothing used to stop N clients from multiplying that out until the + // AppView was killed by the OOM reaper. 512MB is 32 full buffers, which is + // far more concurrency than a single AppView instance sees in practice + // while still fitting comfortably in the smallest deployment. + defaultUploadBufferBudget = 512 * 1024 * 1024 // 512MB + + // defaultUploadIdleTimeout is how long a writer may go without a Write or + // a Commit before the sweeper treats it as abandoned. Generous on purpose: + // the signal is inactivity, not age, so a slow push that is still making + // progress is never reaped no matter how long it runs. + defaultUploadIdleTimeout = time.Hour + + // uploadSweepInterval is how often abandoned writers are looked for. Not + // configurable: the meaningful knob is the timeout, and sweeping a map of + // a few dozen entries every five minutes costs nothing. + uploadSweepInterval = 5 * time.Minute + + // uploadAbortTimeout bounds the hold-side abort issued for a reaped + // writer. It runs on a detached context, so it needs its own deadline. + uploadAbortTimeout = 30 * time.Second + + // uploadBudgetWait bounds how long a single Write waits for budget when + // the request it belongs to has no deadline of its own. Blocking is the + // point (it is backpressure on the Docker client), but blocking forever is + // not: a writer that cannot get budget in this long fails its write and + // the client retries. + uploadBudgetWait = 5 * time.Minute +) + +// bufferBudget is the process-wide allowance for bytes held in upload buffers. +// +// semaphore.Weighted does the waiting, and the held counter exists only so the +// sweeper and the tests can report what is outstanding: the semaphore itself +// does not expose its current weight. +type bufferBudget struct { + sem *semaphore.Weighted + limit int64 + held atomic.Int64 +} + +func newBufferBudget(limit int64) *bufferBudget { + return &bufferBudget{sem: semaphore.NewWeighted(limit), limit: limit} +} + +// acquire blocks until n bytes of budget are available, the context is done, or +// the wait cap expires. n is always at most maxBufferSize per step, and the +// budget is never smaller than maxBufferSize, so a single writer can always +// eventually be satisfied. +func (b *bufferBudget) acquire(ctx context.Context, n int64) error { + if n <= 0 { + return nil + } + waitCtx, cancel := context.WithTimeout(ctx, uploadBudgetWait) + defer cancel() + + if err := b.sem.Acquire(waitCtx, n); err != nil { + return err + } + b.held.Add(n) + return nil +} + +// release returns n bytes to the budget. Callers must release exactly what they +// acquired and no more; ProxyBlobWriter does that by tracking a single charged +// figure and zeroing it as it releases. +func (b *bufferBudget) release(n int64) { + if n <= 0 { + return + } + b.sem.Release(n) + b.held.Add(-n) +} + +// uploadBudget is package level for the same reason globalUploads is: a +// ProxyBlobStore is built fresh on every registry request, so anything that has +// to be shared across requests (and across the uploads that outlive a single +// request) cannot hang off the instance. Threading it through RegistryContext +// would have given each request its own view of a limit that is only meaningful +// process-wide, and would have made every writer's release depend on which +// request happened to construct it. +// +// ConfigureUploads replaces it once at startup, before the listener is up. +var ( + uploadBudget = newBufferBudget(defaultUploadBufferBudget) + uploadIdleTimeout = defaultUploadIdleTimeout +) + +// ConfigureUploads applies the configured upload buffer budget and idle +// timeout. Call it once during startup, before serving: it replaces the +// semaphore outright rather than resizing it, which is only safe while nothing +// holds budget. +// +// A budget below maxBufferSize is raised to it. Anything less could never be +// acquired by even a single writer, so it would not throttle uploads, it would +// stall every one of them until the wait cap expired. +func ConfigureUploads(budgetBytes int64, idleTimeout time.Duration) { + if budgetBytes < maxBufferSize { + slog.Warn("Upload buffer budget below one buffer, raising to the minimum", + "component", "proxy_blob_store", "configured", budgetBytes, "minimum", maxBufferSize) + budgetBytes = maxBufferSize + } + if idleTimeout <= 0 { + idleTimeout = defaultUploadIdleTimeout + } + + uploadBudget = newBufferBudget(budgetBytes) + uploadIdleTimeout = idleTimeout + + slog.Info("Upload buffer budget configured", + "component", "proxy_blob_store", "budget_bytes", budgetBytes, "idle_timeout", idleTimeout) +} + +// UploadStats reports what the in-flight uploads are currently holding: how +// many writers are tracked, and how much buffer budget they hold between them. +func UploadStats() (inFlight int, bytesHeld int64) { + globalUploadsMu.RLock() + inFlight = len(globalUploads) + globalUploadsMu.RUnlock() + return inFlight, uploadBudget.held.Load() +} + +// StartUploadSweeper runs the abandoned-upload sweep until ctx is cancelled. +// +// Not a leased worker. globalUploads is per-process memory, so every instance +// must sweep its own map; electing one instance to do it would leave the others +// leaking exactly as they do today. +func StartUploadSweeper(ctx context.Context) { + go func() { + ticker := time.NewTicker(uploadSweepInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + sweepAbandonedUploads(time.Now()) + } + } + }() +} + +// sweepAbandonedUploads cancels every writer that has been idle longer than the +// configured timeout, and returns how many it reaped. +// +// A Docker client that dies mid-push leaves its writer in globalUploads with +// nothing to remove it: only Commit and Cancel do that, and neither is ever +// called. The writer's buffer, its budget, and the hold-side S3 multipart +// session it may have opened all stay put for the life of the process. +func sweepAbandonedUploads(now time.Time) int { + // Snapshot under the map lock and reap outside it. Reaping takes the + // writer's own lock (and issues a network call), and Commit takes the map + // lock while holding the writer's lock, so holding both here in the other + // order would deadlock. + globalUploadsMu.RLock() + candidates := make([]*ProxyBlobWriter, 0, len(globalUploads)) + for _, w := range globalUploads { + candidates = append(candidates, w) + } + globalUploadsMu.RUnlock() + + reaped := 0 + for _, w := range candidates { + idle, ok := w.reapIfIdle(now, uploadIdleTimeout) + if !ok { + continue + } + reaped++ + + globalUploadsMu.Lock() + delete(globalUploads, w.id) + globalUploadsMu.Unlock() + + slog.Info("Reaped abandoned upload", + "component", "proxy_blob_store/sweep", "id", w.id, "idle", idle, "size", w.Size()) + } + + if reaped > 0 { + inFlight, held := UploadStats() + slog.Info("Abandoned upload sweep finished", + "component", "proxy_blob_store/sweep", "reaped", reaped, "in_flight", inFlight, "bytes_held", held) + } + return reaped +} diff --git a/pkg/appview/storage/upload_budget_test.go b/pkg/appview/storage/upload_budget_test.go new file mode 100644 index 0000000..da70f77 --- /dev/null +++ b/pkg/appview/storage/upload_budget_test.go @@ -0,0 +1,561 @@ +package storage + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/distribution/distribution/v3" + "github.com/opencontainers/go-digest" +) + +// withBudget swaps in a budget of the given size for the duration of the test. +// Sized directly rather than through ConfigureUploads so a test can use a +// budget smaller than one buffer, which ConfigureUploads deliberately refuses. +func withBudget(t *testing.T, limit int64) { + t.Helper() + prev := uploadBudget + uploadBudget = newBufferBudget(limit) + t.Cleanup(func() { uploadBudget = prev }) +} + +// withIdleTimeout swaps in a sweep timeout for the duration of the test. +func withIdleTimeout(t *testing.T, d time.Duration) { + t.Helper() + prev := uploadIdleTimeout + uploadIdleTimeout = d + t.Cleanup(func() { uploadIdleTimeout = prev }) +} + +// isolateUploads gives the test its own globalUploads map, so a sweep here only +// ever sees writers this test created. Other tests leave writers behind, and +// reaping those would fire aborts at hold servers that are already closed. +func isolateUploads(t *testing.T) { + t.Helper() + globalUploadsMu.Lock() + prev := globalUploads + globalUploads = make(map[string]*ProxyBlobWriter) + globalUploadsMu.Unlock() + + t.Cleanup(func() { + globalUploadsMu.Lock() + globalUploads = prev + globalUploadsMu.Unlock() + }) +} + +func bytesHeld() int64 { return uploadBudget.held.Load() } + +// TestUploadBudget_SecondWriterWaitsForTheFirst pins the whole point of the +// budget: when it is spent, the next write blocks instead of allocating, and it +// is released for the next writer the moment the first one finishes. +func TestUploadBudget_SecondWriterWaitsForTheFirst(t *testing.T) { + for _, tc := range []struct { + name string + release func(t *testing.T, w distribution.BlobWriter, data []byte) + }{ + { + name: "released by Commit", + release: func(t *testing.T, w distribution.BlobWriter, data []byte) { + if _, err := w.Commit(context.Background(), distribution.Descriptor{ + Digest: digest.FromBytes(data), + Size: int64(len(data)), + }); err != nil { + t.Errorf("Commit() failed: %v", err) + } + }, + }, + { + name: "released by Cancel", + release: func(t *testing.T, w distribution.BlobWriter, _ []byte) { + if err := w.Cancel(context.Background()); err != nil { + t.Errorf("Cancel() failed: %v", err) + } + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + + const budget = 1 << 20 // exactly one writer's worth, below + withBudget(t, budget) + + store := createTestProxyBlobStore(t, holdServer.URL) + + first, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + data := generateTestData(budget) + if _, err := first.Write(data); err != nil { + t.Fatalf("Write() failed: %v", err) + } + if held := bytesHeld(); held != budget { + t.Fatalf("Expected the first writer to hold the whole budget %d, got %d", budget, held) + } + + second, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer second.Cancel(context.Background()) + + blocked := make(chan error, 1) + go func() { + _, werr := second.Write([]byte("x")) + blocked <- werr + }() + + select { + case err := <-blocked: + t.Fatalf("Expected the second write to block on the budget, it returned %v", err) + case <-time.After(100 * time.Millisecond): + } + + tc.release(t, first, data) + + select { + case err := <-blocked: + if err != nil { + t.Fatalf("Second write failed after the budget was released: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("Second write never got the budget the first one released") + } + }) + } +} + +// TestUploadBudget_BlockedWriteHonoursRequestCancellation pins that a client +// that hangs up while waiting for budget does not keep waiting. Write has no +// context of its own, so the writer carries the one from the request that +// created or resumed it. +func TestUploadBudget_BlockedWriteHonoursRequestCancellation(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + + const budget = 1 << 20 + withBudget(t, budget) + + store := createTestProxyBlobStore(t, holdServer.URL) + + first, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer first.Cancel(context.Background()) + if _, err := first.Write(generateTestData(budget)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + reqCtx, cancelReq := context.WithCancel(context.Background()) + second, err := store.Create(reqCtx) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer second.Cancel(context.Background()) + + blocked := make(chan error, 1) + go func() { + _, werr := second.Write([]byte("x")) + blocked <- werr + }() + + select { + case err := <-blocked: + t.Fatalf("Expected the second write to block, it returned %v", err) + case <-time.After(100 * time.Millisecond): + } + + cancelReq() + + select { + case err := <-blocked: + if !errors.Is(err, context.Canceled) { + t.Fatalf("Expected context.Canceled from the blocked write, got %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("Blocked write ignored its request being cancelled") + } + + if held := bytesHeld(); held != budget { + t.Errorf("A write that never got budget must hold none of it: held %d, want %d", held, budget) + } +} + +// TestConfigureUploads_ClampsBudgetToOneBuffer pins the clamp. A budget below +// one buffer is not a tighter limit, it is a deadlock: no writer could ever +// acquire enough to buffer anything. +func TestConfigureUploads_ClampsBudgetToOneBuffer(t *testing.T) { + prevBudget, prevTimeout := uploadBudget, uploadIdleTimeout + t.Cleanup(func() { uploadBudget, uploadIdleTimeout = prevBudget, prevTimeout }) + + ConfigureUploads(4*1024*1024, 30*time.Minute) + if uploadBudget.limit != maxBufferSize { + t.Errorf("Expected a 4MB budget to be raised to %d, got %d", maxBufferSize, uploadBudget.limit) + } + if uploadIdleTimeout != 30*time.Minute { + t.Errorf("Expected the configured idle timeout to be kept, got %v", uploadIdleTimeout) + } + + ConfigureUploads(64*1024*1024, 0) + if uploadBudget.limit != 64*1024*1024 { + t.Errorf("Expected a budget above one buffer to be kept, got %d", uploadBudget.limit) + } + if uploadIdleTimeout != defaultUploadIdleTimeout { + t.Errorf("Expected a zero idle timeout to fall back to %v, got %v", defaultUploadIdleTimeout, uploadIdleTimeout) + } +} + +// TestUploadBudget_HeldAcrossPartsAndReturnedOnCommit pins what is charged: one +// buffer per writer, not one per part. bytes.Buffer.Reset keeps the backing +// array, so a flush frees no memory and must not free any budget either. +func TestUploadBudget_HeldAcrossPartsAndReturnedOnCommit(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, defaultUploadBufferBudget) + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + // Two and a half buffers, streamed the way ReadFrom feeds the writer, so + // two parts are flushed before Commit. + data := generateTestData(maxBufferSize*2 + maxBufferSize/2) + for off := 0; off < len(data); off += 64 * 1024 { + end := min(off+64*1024, len(data)) + if _, err := writer.Write(data[off:end]); err != nil { + t.Fatalf("Write() failed: %v", err) + } + } + + if held := bytesHeld(); held != maxBufferSize { + t.Errorf("Expected exactly one buffer charged across all parts, got %d bytes held", held) + } + + if _, err := writer.Commit(context.Background(), distribution.Descriptor{ + Digest: digest.FromBytes(data), + Size: int64(len(data)), + }); err != nil { + t.Fatalf("Commit() failed: %v", err) + } + + if held := bytesHeld(); held != 0 { + t.Errorf("Expected the budget back after Commit, got %d bytes still held", held) + } +} + +// TestUploadBudget_WriteLargerThanTheWholeBudgetFailsFast pins the one charge +// that can never be satisfied. Write never grows the buffer past +// maxBufferSize, and ConfigureUploads clamps the budget up to at least that, +// so in production this guard is unreachable. The test sets the budget below a +// buffer directly to prove that if the invariant is ever broken the writer +// fails fast instead of waiting out the acquire deadline on a semaphore that +// can never grant it. +func TestUploadBudget_WriteLargerThanTheWholeBudgetFailsFast(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, 1<<20) + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer writer.Cancel(context.Background()) + + done := make(chan error, 1) + go func() { + _, werr := writer.Write(generateTestData(2 << 20)) + done <- werr + }() + + select { + case err := <-done: + if err == nil { + t.Fatal("Expected an oversized write to be refused") + } + case <-time.After(2 * time.Second): + t.Fatal("Oversized write blocked instead of failing fast") + } + + if held := bytesHeld(); held != 0 { + t.Errorf("Expected a refused write to hold no budget, got %d", held) + } +} + +// TestSweep_ReapsIdleWritersAndLeavesLiveOnes pins the sweep: an abandoned +// upload loses its buffer, its budget, its hold-side multipart session and its +// slot in the map, while an upload that is merely slow is untouched. +func TestSweep_ReapsIdleWritersAndLeavesLiveOnes(t *testing.T) { + t.Run("idle writer with a multipart session is aborted", func(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, defaultUploadBufferBudget) + withIdleTimeout(t, time.Hour) + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + pbw := writer.(*ProxyBlobWriter) + + // Enough to flush a part, which is what opens the multipart session. + if _, err := writer.Write(generateTestData(maxBufferSize + 1024)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + if pbw.uploadID == "" { + t.Fatal("Expected a multipart session after a flush") + } + if held := bytesHeld(); held == 0 { + t.Fatal("Expected the writer to hold budget before the sweep") + } + + pbw.lastActivity = time.Now().Add(-2 * time.Hour) + + if reaped := sweepAbandonedUploads(time.Now()); reaped != 1 { + t.Fatalf("Expected 1 reaped upload, got %d", reaped) + } + + holdServer.mu.Lock() + aborts := len(holdServer.AbortCalls) + holdServer.mu.Unlock() + if aborts != 1 { + t.Errorf("Expected the hold-side multipart to be aborted once, got %d aborts", aborts) + } + + if held := bytesHeld(); held != 0 { + t.Errorf("Expected the budget back after the sweep, got %d bytes still held", held) + } + + globalUploadsMu.RLock() + _, stillTracked := globalUploads[pbw.id] + globalUploadsMu.RUnlock() + if stillTracked { + t.Error("Expected the reaped writer to be gone from the uploads map") + } + + if _, err := store.Resume(context.Background(), pbw.id); err != distribution.ErrBlobUploadUnknown { + t.Errorf("Expected ErrBlobUploadUnknown resuming a reaped upload, got %v", err) + } + }) + + t.Run("idle writer that never flushed sends no abort", func(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, defaultUploadBufferBudget) + withIdleTimeout(t, time.Hour) + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + pbw := writer.(*ProxyBlobWriter) + if _, err := writer.Write(generateTestData(2048)); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + pbw.lastActivity = time.Now().Add(-2 * time.Hour) + + if reaped := sweepAbandonedUploads(time.Now()); reaped != 1 { + t.Fatalf("Expected 1 reaped upload, got %d", reaped) + } + + holdServer.mu.Lock() + total := holdServer.TotalCalls + holdServer.mu.Unlock() + if total != 0 { + t.Errorf("Expected no hold call for a writer with no multipart session, got %d", total) + } + if held := bytesHeld(); held != 0 { + t.Errorf("Expected the budget back after the sweep, got %d bytes still held", held) + } + }) + + t.Run("active writer is left alone", func(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, defaultUploadBufferBudget) + withIdleTimeout(t, time.Hour) + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer writer.Cancel(context.Background()) + pbw := writer.(*ProxyBlobWriter) + + // Started long ago but still writing: StartedAt is stale, activity is not. + pbw.startedAt = time.Now().Add(-24 * time.Hour) + data := generateTestData(2048) + if _, err := writer.Write(data); err != nil { + t.Fatalf("Write() failed: %v", err) + } + + if reaped := sweepAbandonedUploads(time.Now()); reaped != 0 { + t.Fatalf("Expected an active writer to survive the sweep, %d were reaped", reaped) + } + + resumed, err := store.Resume(context.Background(), pbw.id) + if err != nil { + t.Fatalf("Resume() after the sweep failed: %v", err) + } + if resumed.Size() != int64(len(data)) { + t.Errorf("Expected the surviving writer to keep its %d bytes, got %d", len(data), resumed.Size()) + } + }) + + t.Run("resume counts as activity", func(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, defaultUploadBufferBudget) + withIdleTimeout(t, time.Hour) + + store := createTestProxyBlobStore(t, holdServer.URL) + + writer, err := store.Create(context.Background()) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + defer writer.Cancel(context.Background()) + pbw := writer.(*ProxyBlobWriter) + + pbw.lastActivity = time.Now().Add(-2 * time.Hour) + if _, err := store.Resume(context.Background(), pbw.id); err != nil { + t.Fatalf("Resume() failed: %v", err) + } + + if reaped := sweepAbandonedUploads(time.Now()); reaped != 0 { + t.Fatalf("Expected a just-resumed writer to survive the sweep, %d were reaped", reaped) + } + }) +} + +// TestSweep_ConcurrentWithLiveWriter runs the sweep against writers that are +// being written to and committed at the same time. Its job is to fail under +// -race if the sweeper and a live request ever touch the writer unguarded. +func TestSweep_ConcurrentWithLiveWriter(t *testing.T) { + s3Server := newMockS3Server(t, true) + defer s3Server.Close() + + holdServer := newMockHoldServer(t, s3Server.URL) + defer holdServer.Close() + + isolateUploads(t) + withBudget(t, defaultUploadBufferBudget) + // Aggressive on purpose: the sweeper is trying to reap the very writers the + // other goroutines are still using. + withIdleTimeout(t, time.Nanosecond) + + store := createTestProxyBlobStore(t, holdServer.URL) + + stop := make(chan struct{}) + var sweeper sync.WaitGroup + sweeper.Add(1) + go func() { + defer sweeper.Done() + for { + select { + case <-stop: + return + default: + sweepAbandonedUploads(time.Now()) + } + } + }() + + var writers sync.WaitGroup + for range 8 { + writers.Add(1) + go func() { + defer writers.Done() + + writer, err := store.Create(context.Background()) + if err != nil { + t.Errorf("Create() failed: %v", err) + return + } + + data := generateTestData(4096) + // A reaped writer refuses further work; that is a legitimate + // outcome of this race, not a failure. What must not happen is a + // data race or a lost budget charge. + if _, err := writer.Write(data); err != nil { + return + } + _ = writer.Size() + if _, err := writer.Commit(context.Background(), distribution.Descriptor{ + Digest: digest.FromBytes(data), + Size: int64(len(data)), + }); err != nil { + writer.Cancel(context.Background()) + } + }() + } + + writers.Wait() + close(stop) + sweeper.Wait() + + // Whichever side won each race, every writer is finished and the budget is + // whole again. + if held := bytesHeld(); held != 0 { + t.Errorf("Expected no budget held once every writer finished, got %d", held) + } + if inFlight, _ := UploadStats(); inFlight != 0 { + t.Errorf("Expected no uploads left in flight, got %d", inFlight) + } +}