at-container-registry/pkg/appview/storage/upload_budget.go

package storage

import (
	"context"
	"log/slog"
	"sync/atomic"
	"time"

	"golang.org/x/sync/semaphore"
)

const (
	// maxWriterBuffers is how many buffers one writer can hold at its peak: the
	// one being filled plus the one being uploaded. A blob small enough never
	// to flush holds one and never allocates the second, and so does a large
	// one whose second buffer the budget could not afford.
	maxWriterBuffers = 2

	// maxWriterFootprint is a writer's peak buffer memory, 32MB: what a large
	// upload costs when it gets to upload a part in the background while the
	// next buffer fills. It is a ceiling, not a price of admission. Only the
	// first buffer is ever waited for; the second is taken only when the budget
	// can spare it, and a writer that cannot have one uploads its parts
	// synchronously out of the single buffer it already holds.
	maxWriterFootprint = maxWriterBuffers * maxBufferSize

	// defaultUploadBufferBudget is the process-wide ceiling on bytes held in
	// blob upload buffers. Every in-flight push holds up to maxWriterFootprint,
	// 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.
	//
	// What the number buys is progress for budget/maxBufferSize large uploads,
	// 32 of them at 512MB, because one full buffer is all a writer needs to
	// finish. Pipelining is what the rest of the memory is spent on when it is
	// there: up to 16 of those 32 can also hold a part in flight. That 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

	// uploadPartTimeout bounds a background part upload: one hold call for the
	// presigned URL (two on the first part, which also starts the session) and
	// the S3 PUT itself. Detached from the request that filled the buffer,
	// because that request is usually answered before the PUT finishes, so it
	// needs a deadline of its own. It is also what stops a wedged part from
	// hiding a writer from the sweeper forever.
	uploadPartTimeout = 5 * time.Minute

	// uploadFlightAbandonTimeout bounds how long Cancel waits for a part that
	// is still going up before aborting the session anyway. Waiting at all is
	// what keeps the abort from overtaking the PUT and leaking the S3 upload
	// ID; not waiting forever is what keeps a client's DELETE from hanging on
	// a PUT that has stopped making progress.
	uploadFlightAbandonTimeout = 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. Only a writer's first buffer is ever acquired this way,
// and that buffer never exceeds maxBufferSize, which the budget is never
// smaller than, 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
}

// tryAcquire takes n bytes of budget if they are free right now and keepFree
// bytes would still be free afterwards, and reports whether it got them. It
// never blocks and never queues, which is what makes it safe for a charge a
// writer must be able to do without: waiting for anything past the first buffer
// is what wedged the budget.
//
// keepFree is asked for as part of the same weight and handed straight back,
// because the semaphore has no way to ask whether that much more is free
// without taking it. Only n ends up held. Between the two calls the budget
// looks fuller than it is, which can make a concurrent tryAcquire decline; that
// costs a writer its pipelining for one part and nothing more.
//
// x/sync/semaphore's TryAcquire also fails outright while anyone is queued on
// Acquire, so an opportunistic second buffer can never jump ahead of a writer
// that is waiting for its first.
func (b *bufferBudget) tryAcquire(n, keepFree int64) bool {
	if n <= 0 {
		return true
	}
	if !b.sem.TryAcquire(n + keepFree) {
		return false
	}
	if keepFree > 0 {
		b.sem.Release(keepFree)
	}
	b.held.Add(n)
	return true
}

// 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 maxWriterFootprint is raised to it. One buffer is all a writer
// needs to finish, so the floor could be maxBufferSize and still be correct;
// two is the smallest figure that leaves the process room to do anything but
// serialise every push behind one buffer. At the floor itself that is two large
// uploads each filling a buffer rather than one going twice as fast: a second
// buffer is only ever taken while a buffer's worth of budget would still be
// free afterwards, which at 32MB it never is.
func ConfigureUploads(budgetBytes int64, idleTimeout time.Duration) {
	if budgetBytes < maxWriterFootprint {
		slog.Warn("Upload buffer budget below one writer's peak, raising to the minimum",
			"component", "proxy_blob_store", "configured", budgetBytes, "minimum", maxWriterFootprint)
		budgetBytes = maxWriterFootprint
	}
	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
}