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

package storage

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"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_ClampsBudgetToOneWriter pins the clamp. A budget below
// one writer's peak is not a tighter limit, it is a deadlock: no large upload
// could ever acquire enough to buffer anything. The floor is two buffers, not
// one, because a writer with a part in flight holds two.
func TestConfigureUploads_ClampsBudgetToOneWriter(t *testing.T) {
	prevBudget, prevTimeout := uploadBudget, uploadIdleTimeout
	t.Cleanup(func() { uploadBudget, uploadIdleTimeout = prevBudget, prevTimeout })

	ConfigureUploads(4*1024*1024, 30*time.Minute)
	if uploadBudget.limit != maxWriterFootprint {
		t.Errorf("Expected a 4MB budget to be raised to %d, got %d", maxWriterFootprint, 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 writer's peak 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: at
// most two buffers 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, and the two buffers are swapped rather than reallocated.
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)
		}
	}

	waitForParts(writer.(*ProxyBlobWriter))
	if held := bytesHeld(); held != maxWriterFootprint {
		t.Errorf("Expected exactly two buffers 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)
	}
}

// TestUploadBudget_ThreeWritersShareABudgetForTwo pins the invariant that keeps
// the budget from wedging: a writer never waits on the budget for anything past
// its first buffer.
//
// The wedge it regresses: every writer charged its second buffer as that buffer
// grew, so writers that were all mid-growth could hold the entire budget while
// none of them had reached its peak. Nothing is released before Commit, no
// writer could reach Commit without more budget, and there was none, so every
// Write sat out the five minute wait cap and failed. Three writers with room
// for two is exactly that shape: 48MB of first buffers out of 64MB, then all
// three carrying on at once.
func TestUploadBudget_ThreeWritersShareABudgetForTwo(t *testing.T) {
	s3Server := newMockS3Server(t, true)
	defer s3Server.Close()

	holdServer := newMockHoldServer(t, s3Server.URL)
	defer holdServer.Close()

	isolateUploads(t)

	const writers = 3
	withBudget(t, 2*maxWriterFootprint) // room for two writers at their peak

	store := createTestProxyBlobStore(t, holdServer.URL)

	// Two buffers each, so every writer flushes a part and then keeps writing,
	// which is the only way to hold a second buffer at all.
	data := generateTestData(2 * maxBufferSize)
	dgst := digest.FromBytes(data)

	// Nobody goes past its first buffer until all three have one. That is the
	// state the budget used to run out in, and it is the one a regression would
	// hang in.
	var firstBuffers sync.WaitGroup
	firstBuffers.Add(writers)
	done := make(chan error, writers)

	for i := range writers {
		go func(i int) {
			writer, err := store.Create(context.Background())
			if err != nil {
				firstBuffers.Done()
				done <- fmt.Errorf("writer %d: Create: %w", i, err)
				return
			}

			// Streamed rather than handed over whole, the way a PATCH body
			// arrives, so the writers interleave their growth.
			write := func(b []byte) error {
				const chunk = 4 << 20
				for off := 0; off < len(b); off += chunk {
					if _, err := writer.Write(b[off:min(off+chunk, len(b))]); err != nil {
						return fmt.Errorf("writer %d: Write: %w", i, err)
					}
				}
				return nil
			}

			err = write(data[:maxBufferSize])
			firstBuffers.Done()
			if err != nil {
				writer.Cancel(context.Background())
				done <- err
				return
			}
			firstBuffers.Wait()

			if err := write(data[maxBufferSize:]); err != nil {
				writer.Cancel(context.Background())
				done <- err
				return
			}

			if _, err := writer.Commit(context.Background(), distribution.Descriptor{
				Digest: dgst,
				Size:   int64(len(data)),
			}); err != nil {
				done <- fmt.Errorf("writer %d: Commit: %w", i, err)
				return
			}
			done <- nil
		}(i)
	}

	// Generous on purpose: the mock servers answer in microseconds, so the only
	// way to reach this is a writer stuck on the budget. It exists to fail
	// rather than sit through the acquire's own five minute cap.
	deadline := time.After(60 * time.Second)
	for range writers {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("A writer failed while three shared a budget for two: %v", err)
			}
		case <-deadline:
			t.Fatalf("Writers wedged on the budget: %d of %d bytes held with nothing able to release them",
				bytesHeld(), uploadBudget.limit)
		}
	}

	if held := bytesHeld(); held != 0 {
		t.Errorf("Expected the whole budget back once every writer committed, got %d bytes still held", held)
	}
}

// TestUploadBudget_PipelinesOnlyWhenTheBudgetIsFree pins both sides of the
// opportunistic second buffer: a writer that can have one keeps filling while
// its part goes up, and a writer that cannot still finishes the blob out of the
// one buffer it holds.
func TestUploadBudget_PipelinesOnlyWhenTheBudgetIsFree(t *testing.T) {
	t.Run("free budget buys a second buffer", func(t *testing.T) {
		s3Server := newMockS3Server(t, true)
		defer s3Server.Close()

		holdServer := newMockHoldServer(t, s3Server.URL)
		defer holdServer.Close()

		isolateUploads(t)
		withBudget(t, defaultUploadBufferBudget)

		// Hold the first part on its way to S3, so the only way the write below
		// can finish is in a second buffer.
		release := make(chan struct{})
		s3Server.setPartHook(func(int) error {
			<-release
			return nil
		})

		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(maxBufferSize + 4<<20)
		wrote := make(chan error, 1)
		go func() {
			_, werr := writer.Write(data)
			wrote <- werr
		}()

		select {
		case err := <-wrote:
			if err != nil {
				t.Fatalf("Write() failed: %v", err)
			}
		case <-time.After(30 * time.Second):
			close(release)
			t.Fatal("Write waited for the part it handed off instead of carrying on in a second buffer")
		}

		// The part is still going up, so both buffers are the writer's and both
		// are charged.
		if held := bytesHeld(); held != maxWriterFootprint {
			t.Errorf("Expected both buffers charged while a part is in flight, %d, got %d", maxWriterFootprint, held)
		}
		pbw.mu.Lock()
		pipelined := pbw.flight != nil
		pbw.mu.Unlock()
		if !pipelined {
			t.Error("Expected the part to still be in flight behind the second buffer")
		}

		close(release)
		s3Server.setPartHook(nil)

		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)
		}
	})

	t.Run("a budget of one buffer still finishes the blob", func(t *testing.T) {
		s3Server := newMockS3Server(t, true)
		defer s3Server.Close()

		holdServer := newMockHoldServer(t, s3Server.URL)
		defer holdServer.Close()

		isolateUploads(t)
		withBudget(t, maxBufferSize) // exactly one buffer: no room to pipeline

		store := createTestProxyBlobStore(t, holdServer.URL)
		writer, err := store.Create(context.Background())
		if err != nil {
			t.Fatalf("Create() failed: %v", err)
		}
		pbw := writer.(*ProxyBlobWriter)

		// Sample what the budget reports throughout: the writer must never hold
		// more than the single buffer it was able to charge.
		var peak atomic.Int64
		stop := make(chan struct{})
		var sampler sync.WaitGroup
		sampler.Add(1)
		go func() {
			defer sampler.Done()
			for {
				select {
				case <-stop:
					return
				default:
					if held := bytesHeld(); held > peak.Load() {
						peak.Store(held)
					}
				}
			}
		}()

		data := generateTestData(3 * maxBufferSize)
		for off := 0; off < len(data); off += 1 << 20 {
			if _, err := writer.Write(data[off:min(off+(1<<20), len(data))]); err != nil {
				close(stop)
				sampler.Wait()
				t.Fatalf("Write() failed with a budget of one buffer: %v", err)
			}
		}

		if other := pbw.otherBufferCap(); other != 0 {
			t.Errorf("Expected a writer under pressure to hold one buffer, found a second of %d bytes", other)
		}

		desc, err := writer.Commit(context.Background(), distribution.Descriptor{
			Digest: digest.FromBytes(data),
			Size:   int64(len(data)),
		})
		close(stop)
		sampler.Wait()
		if err != nil {
			t.Fatalf("Commit() failed with a budget of one buffer: %v", err)
		}
		if desc.Size != int64(len(data)) {
			t.Errorf("Committed size %d, expected %d", desc.Size, len(data))
		}

		if peak.Load() > maxBufferSize {
			t.Errorf("Expected never to hold more than one buffer, peaked at %d", peak.Load())
		}
		if parts := s3Server.finishedParts(); len(parts) != 3 {
			t.Errorf("Expected the blob to go up as 3 parts, got %v", parts)
		}
		if held := bytesHeld(); held != 0 {
			t.Errorf("Expected the budget back after Commit, got %d bytes still held", 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. The
		// part goes up in the background, so wait for it: a writer with a part
		// in flight is deliberately not reapable.
		if _, err := writer.Write(generateTestData(maxBufferSize + 1024)); err != nil {
			t.Fatalf("Write() failed: %v", err)
		}
		waitForParts(pbw)
		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)
	}
}