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

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