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

package storage

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/distribution/distribution/v3"
	"github.com/opencontainers/go-digest"
)

// These tests are about the pipelining in ProxyBlobWriter: one part upload in
// flight while the next buffer fills. They observe the behaviour rather than
// only the result, so most of them hold a part PUT open on a channel and assert
// what the writer does (or refuses to do) while it is stuck there.

// blockGate is a one-shot release for a blocked part PUT, plus the record of
// which parts have gone through it.
type blockGate struct {
	release chan struct{}
	once    sync.Once
}

func newBlockGate() *blockGate { return &blockGate{release: make(chan struct{})} }

// hookFor blocks the named part until the gate is released and lets every other
// part through untouched.
func (g *blockGate) hookFor(part int) func(int) error {
	return func(n int) error {
		if n == part {
			<-g.release
		}
		return nil
	}
}

func (g *blockGate) open() { g.once.Do(func() { close(g.release) }) }

// waitUntil polls cond until it holds, failing the test if it never does. Used
// to wait for something a background goroutine does (a part PUT arriving at
// S3), where there is no channel to wait on from the test side.
func waitUntil(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

// newPipelineWriter wires a hold and an S3 in front of a fresh writer, with the
// uploads map and the budget isolated to this test.
func newPipelineWriter(t *testing.T) (*mockS3Server, *mockHoldServer, *ProxyBlobWriter) {
	t.Helper()

	s3Server := newMockS3Server(t, true)
	t.Cleanup(s3Server.Close)

	holdServer := newMockHoldServer(t, s3Server.URL)
	t.Cleanup(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)
	}

	return s3Server, holdServer, writer.(*ProxyBlobWriter)
}

// mustWrite writes the whole slice and fails the test if it does not land.
func mustWrite(t *testing.T, w *ProxyBlobWriter, p []byte) {
	t.Helper()
	n, err := w.Write(p)
	if err != nil {
		t.Fatalf("Write() of %d bytes failed after %d: %v", len(p), n, err)
	}
	if n != len(p) {
		t.Fatalf("Write() took %d of %d bytes", n, len(p))
	}
}

// TestWrite_OverlapsPartUploadWithTheNextBuffer is the point of the whole
// change. The upload used to be strictly serial: while a part was being PUT to
// S3 the PATCH body was not being drained, so a layer cost receive time plus
// send time. Here the first part is held open at S3 and the writer must keep
// accepting a whole second buffer's worth of bytes while it is stuck.
//
// The other half of the contract is the ceiling: only one part in flight. So
// the write that fills the second buffer must block, because handing it off
// would make two.
func TestWrite_OverlapsPartUploadWithTheNextBuffer(t *testing.T) {
	s3Server, holdServer, pbw := newPipelineWriter(t)

	gate := newBlockGate()
	s3Server.setPartHook(gate.hookFor(1))
	defer gate.open()

	// Two full buffers plus a tail, so the commit has a final part to flush.
	const tail = 1024
	data := generateTestData(2*maxBufferSize + tail)

	// First buffer: fills, hands part 1 off, returns while the PUT is blocked.
	mustWrite(t, pbw, data[:maxBufferSize])
	waitUntil(t, "part 1's PUT to reach S3", func() bool {
		return len(s3Server.startedParts()) == 1
	})

	// The second buffer, all but the byte that would fill it. Every one of
	// these writes has to complete with part 1 still stuck in S3.
	mustWrite(t, pbw, data[maxBufferSize:2*maxBufferSize-1])

	if got := s3Server.finishedParts(); len(got) != 0 {
		t.Fatalf("Part 1 was supposed to still be in flight, S3 has finished %v", got)
	}

	// The byte that fills the second buffer, and the tail behind it. This one
	// must block: handing off now would put two parts in flight.
	blocked := make(chan error, 1)
	go func() {
		_, err := pbw.Write(data[2*maxBufferSize-1:])
		blocked <- err
	}()

	select {
	case err := <-blocked:
		t.Fatalf("The write that fills the second buffer returned (%v) while part 1 was still in flight: two parts at once", err)
	case <-time.After(250 * time.Millisecond):
		// Still waiting for the in-flight slot, which is what we want.
	}

	gate.open()

	select {
	case err := <-blocked:
		if err != nil {
			t.Fatalf("Write() failed once part 1 was released: %v", err)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("Write() never returned after part 1 was released")
	}

	dgst := digest.FromBytes(data)
	desc, err := pbw.Commit(context.Background(), distribution.Descriptor{
		Digest: dgst,
		Size:   int64(len(data)),
	})
	if err != nil {
		t.Fatalf("Commit() failed: %v", err)
	}
	if desc.Digest != dgst {
		t.Errorf("Committed digest %s, want %s", desc.Digest, dgst)
	}
	if desc.Size != int64(len(data)) {
		t.Errorf("Committed size %d, want %d", desc.Size, len(data))
	}

	// Three parts: two full buffers and the tail, in order, each with the bytes
	// that were written into it.
	if got, want := s3Server.finishedParts(), []int{1, 2, 3}; fmt.Sprint(got) != fmt.Sprint(want) {
		t.Fatalf("S3 recorded parts %v, want %v", got, want)
	}
	for i, want := range [][]byte{
		data[:maxBufferSize],
		data[maxBufferSize : 2*maxBufferSize],
		data[2*maxBufferSize:],
	} {
		got := s3Server.part(i + 1)
		if len(got) != len(want) {
			t.Errorf("Part %d is %d bytes, want %d", i+1, len(got), len(want))
			continue
		}
		if string(got) != string(want) {
			t.Errorf("Part %d carries the wrong bytes", i+1)
		}
	}

	// The writer named them to the hold in order, with the ETags S3 handed back.
	holdServer.mu.Lock()
	completes := append([]mockCompleteCall(nil), holdServer.CompleteCalls...)
	holdServer.mu.Unlock()

	if len(completes) != 1 {
		t.Fatalf("Expected 1 completeUpload call, got %d", len(completes))
	}
	if len(completes[0].Parts) != 3 {
		t.Fatalf("Expected 3 parts named at completion, got %d", len(completes[0].Parts))
	}
	for i, part := range completes[0].Parts {
		wantNum := float64(i + 1)
		if part["part_number"] != wantNum {
			t.Errorf("Part at index %d is numbered %v, want %v: parts must stay ordered", i, part["part_number"], wantNum)
		}
		wantETag := fmt.Sprintf(`"etag-part-%d"`, i+1)
		if part["etag"] != wantETag {
			t.Errorf("Part %d has ETag %v, want %s", i+1, part["etag"], wantETag)
		}
	}
}

// TestWrite_PartFailureSurfacesOnePartLate pins how a background part reports a
// failure. Nothing is watching the goroutine, so the error cannot be returned by
// the write that handed the buffer over: it lands on whatever touches the
// writer next. The session is aborted from the goroutine itself, so it is gone
// whether or not anyone comes back for it.
func TestWrite_PartFailureSurfacesOnePartLate(t *testing.T) {
	t.Run("reported by a later write", func(t *testing.T) {
		s3Server, holdServer, pbw := newPipelineWriter(t)

		s3Server.setPartHook(func(n int) error {
			if n == 1 {
				return fmt.Errorf("S3 refused part one")
			}
			return nil
		})

		// The write that fills the buffer only hands it off, so it succeeds.
		mustWrite(t, pbw, generateTestData(maxBufferSize))
		waitForParts(pbw)

		_, err := pbw.Write(generateTestData(1024))
		if err == nil {
			t.Fatal("Expected the failed part to be reported by the next write")
		}
		if !strings.Contains(err.Error(), "part 1 upload failed") {
			t.Errorf("Expected the error to name the part that failed, got: %v", err)
		}
		if !strings.Contains(err.Error(), "S3 refused part one") {
			t.Errorf("Expected the underlying cause to survive the wrapping, got: %v", err)
		}

		// The writer is closed, and stays closed with the same explanation.
		if _, err := pbw.Write(generateTestData(16)); err == nil {
			t.Error("Expected further writes to fail after a part failure")
		} else if !strings.Contains(err.Error(), "S3 refused part one") {
			t.Errorf("Expected the cause to stay visible on later writes, got: %v", err)
		}

		holdServer.mu.Lock()
		aborts := len(holdServer.AbortCalls)
		holdServer.mu.Unlock()
		if aborts != 1 {
			t.Errorf("Expected the multipart session to be aborted exactly once, got %d aborts", aborts)
		}

		if _, err := pbw.Commit(context.Background(), distribution.Descriptor{}); err == nil {
			t.Error("Expected Commit to fail on a writer closed by a failed part")
		}
	})

	t.Run("reported by Commit when nothing else touches the writer", func(t *testing.T) {
		s3Server, holdServer, pbw := newPipelineWriter(t)

		s3Server.setPartHook(func(n int) error {
			if n == 1 {
				return fmt.Errorf("S3 refused part one")
			}
			return nil
		})

		data := generateTestData(maxBufferSize)
		mustWrite(t, pbw, data)

		// No waitForParts: Commit has to do the waiting itself, and report what
		// it finds.
		_, err := pbw.Commit(context.Background(), distribution.Descriptor{
			Digest: digest.FromBytes(data),
			Size:   int64(len(data)),
		})
		if err == nil {
			t.Fatal("Expected Commit to fail when the in-flight part failed")
		}
		if !strings.Contains(err.Error(), "S3 refused part one") {
			t.Errorf("Expected the cause in Commit's error, got: %v", err)
		}

		holdServer.mu.Lock()
		aborts := len(holdServer.AbortCalls)
		completes := len(holdServer.CompleteCalls)
		holdServer.mu.Unlock()
		if aborts != 1 {
			t.Errorf("Expected exactly 1 abort, got %d", aborts)
		}
		if completes != 0 {
			t.Errorf("Expected no completeUpload after a failed part, got %d", completes)
		}
	})
}

// TestCommit_WaitsForTheInFlightPart pins that Commit does not run ahead of a
// part that is still going up. Its ETag has to be in the list completeUpload
// names, so committing without waiting would silently drop a part.
func TestCommit_WaitsForTheInFlightPart(t *testing.T) {
	s3Server, holdServer, pbw := newPipelineWriter(t)

	gate := newBlockGate()
	s3Server.setPartHook(gate.hookFor(1))
	defer gate.open()

	const tail = 4096
	data := generateTestData(maxBufferSize + tail)

	mustWrite(t, pbw, data[:maxBufferSize])
	waitUntil(t, "part 1's PUT to reach S3", func() bool {
		return len(s3Server.startedParts()) == 1
	})
	mustWrite(t, pbw, data[maxBufferSize:])

	type result struct {
		desc distribution.Descriptor
		err  error
	}
	done := make(chan result, 1)
	go func() {
		desc, err := pbw.Commit(context.Background(), distribution.Descriptor{
			Digest: digest.FromBytes(data),
			Size:   int64(len(data)),
		})
		done <- result{desc, err}
	}()

	select {
	case r := <-done:
		t.Fatalf("Commit returned (%v, %v) while part 1 was still in flight", r.desc, r.err)
	case <-time.After(250 * time.Millisecond):
		// Still waiting, which is the point.
	}

	// completeUpload must not have gone out either.
	holdServer.mu.Lock()
	early := len(holdServer.CompleteCalls)
	holdServer.mu.Unlock()
	if early != 0 {
		t.Fatalf("completeUpload went out before the in-flight part landed (%d calls)", early)
	}

	gate.open()

	select {
	case r := <-done:
		if r.err != nil {
			t.Fatalf("Commit() failed once part 1 landed: %v", r.err)
		}
		if r.desc.Size != int64(len(data)) {
			t.Errorf("Committed size %d, want %d", r.desc.Size, len(data))
		}
	case <-time.After(10 * time.Second):
		t.Fatal("Commit never returned after part 1 was released")
	}

	if got, want := s3Server.finishedParts(), []int{1, 2}; fmt.Sprint(got) != fmt.Sprint(want) {
		t.Errorf("S3 recorded parts %v, want %v", got, want)
	}
}

// TestCancel_WaitsForTheInFlightPartBeforeAborting pins the teardown order. An
// abort that overtakes a running part PUT can arrive before the hold has issued
// the upload ID at all, and then nothing left knows the S3 session exists. So
// Cancel waits for the goroutine first and aborts after it.
func TestCancel_WaitsForTheInFlightPartBeforeAborting(t *testing.T) {
	s3Server, holdServer, pbw := newPipelineWriter(t)

	var eventsMu sync.Mutex
	var events []string
	record := func(what string) {
		eventsMu.Lock()
		events = append(events, what)
		eventsMu.Unlock()
	}

	gate := newBlockGate()
	s3Server.setPartHook(func(n int) error {
		if n == 1 {
			<-gate.release
			record("part-put-finished")
		}
		return nil
	})
	defer gate.open()

	holdServer.mu.Lock()
	holdServer.AbortHook = func() { record("abort") }
	holdServer.mu.Unlock()

	mustWrite(t, pbw, generateTestData(maxBufferSize))
	waitUntil(t, "part 1's PUT to reach S3", func() bool {
		return len(s3Server.startedParts()) == 1
	})
	waitUntil(t, "the multipart session to be opened", func() bool {
		holdServer.mu.Lock()
		defer holdServer.mu.Unlock()
		return len(holdServer.InitiateCalls) == 1
	})

	cancelled := make(chan error, 1)
	go func() { cancelled <- pbw.Cancel(context.Background()) }()

	select {
	case err := <-cancelled:
		t.Fatalf("Cancel returned (%v) while the part was still in flight", err)
	case <-time.After(250 * time.Millisecond):
		// Waiting for the goroutine, as it should be.
	}

	gate.open()

	select {
	case err := <-cancelled:
		if err != nil {
			t.Fatalf("Cancel() failed: %v", err)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("Cancel never returned after the part was released")
	}

	eventsMu.Lock()
	got := append([]string(nil), events...)
	eventsMu.Unlock()
	want := []string{"part-put-finished", "abort"}
	if fmt.Sprint(got) != fmt.Sprint(want) {
		t.Errorf("Event order was %v, want %v: the abort must not race the part PUT", got, want)
	}

	holdServer.mu.Lock()
	aborts := len(holdServer.AbortCalls)
	holdServer.mu.Unlock()
	if aborts != 1 {
		t.Errorf("Expected exactly 1 abort, got %d", aborts)
	}

	if held := bytesHeld(); held != 0 {
		t.Errorf("Expected both buffers' budget back after Cancel, got %d bytes held", held)
	}

	globalUploadsMu.RLock()
	_, tracked := globalUploads[pbw.id]
	globalUploadsMu.RUnlock()
	if tracked {
		t.Error("Expected the cancelled writer to be gone from the uploads map")
	}
}

// TestUploadBudget_ChargesTwoBuffersAtPeak pins the memory accounting for the
// overlap. A writer with a part in flight holds two buffers, so a large upload
// costs 32MB rather than 16, and it must be charged for both. A blob that never
// fills a buffer never allocates the second one and must keep paying only for
// what it grew.
func TestUploadBudget_ChargesTwoBuffersAtPeak(t *testing.T) {
	t.Run("large upload charges exactly two buffers", func(t *testing.T) {
		s3Server, _, pbw := newPipelineWriter(t)

		gate := newBlockGate()
		s3Server.setPartHook(gate.hookFor(1))
		defer gate.open()

		data := generateTestData(2*maxBufferSize + 1024)

		mustWrite(t, pbw, data[:maxBufferSize])
		waitUntil(t, "part 1's PUT to reach S3", func() bool {
			return len(s3Server.startedParts()) == 1
		})

		// Fill the second buffer to one byte short of full (filling it exactly
		// would block on the in-flight slot). Two writes rather than one: the
		// first takes the backing array past halfway, which is what makes
		// growBuffer size it to the whole threshold, so the charge is a round
		// buffer rather than an awkward fraction of one.
		half := maxBufferSize + maxBufferSize/2 + 1
		mustWrite(t, pbw, data[maxBufferSize:half])
		mustWrite(t, pbw, data[half:2*maxBufferSize-1])

		// One buffer in flight, one being filled: this is the peak.
		if held := bytesHeld(); held != maxWriterFootprint {
			t.Errorf("Expected the peak to be two buffers (%d), got %d bytes held", maxWriterFootprint, held)
		}

		gate.open()
		mustWrite(t, pbw, data[2*maxBufferSize-1:])
		waitForParts(pbw)

		// Still two: the buffers are swapped, never reallocated, and nothing is
		// released between parts because Reset keeps the backing array.
		if held := bytesHeld(); held != maxWriterFootprint {
			t.Errorf("Expected two buffers still charged after the swap, got %d bytes held", held)
		}

		if _, err := pbw.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 whole charge back after Commit, got %d bytes held", held)
		}
	})

	t.Run("small blob charges only the buffer it grew", func(t *testing.T) {
		_, _, pbw := newPipelineWriter(t)

		data := generateTestData(4096)
		mustWrite(t, pbw, data)

		pbw.mu.Lock()
		charged, bufCap, spare, flight := pbw.charged, int64(pbw.buffer.Cap()), pbw.spare, pbw.flight
		pbw.mu.Unlock()

		if spare != nil || flight != nil {
			t.Error("A blob that never filled a buffer must not have allocated a second one")
		}
		if charged != bufCap {
			t.Errorf("Charged %d for a buffer of %d: a small blob pays for what it grew", charged, bufCap)
		}
		if charged >= maxBufferSize {
			t.Errorf("A 4KB blob was charged %d bytes, a whole buffer is %d", charged, maxBufferSize)
		}

		if _, err := pbw.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 charge back after Commit, got %d bytes held", held)
		}
	})
}

// TestSweep_LeavesAWriterWithAPartInFlight pins the sweeper's read of what busy
// means now that a part goes up off the writer's lock. lastActivity is only
// bumped when a part lands, so a writer mid-PUT can look arbitrarily idle by the
// clock; reaping it would abort the multipart session out from under a goroutine
// that is still feeding it.
func TestSweep_LeavesAWriterWithAPartInFlight(t *testing.T) {
	s3Server, holdServer, pbw := newPipelineWriter(t)
	withIdleTimeout(t, time.Hour)

	gate := newBlockGate()
	s3Server.setPartHook(gate.hookFor(1))
	defer gate.open()

	mustWrite(t, pbw, generateTestData(maxBufferSize))
	waitUntil(t, "part 1's PUT to reach S3", func() bool {
		return len(s3Server.startedParts()) == 1
	})

	// As idle as the clock can make it: the last Write was, as far as
	// lastActivity knows, two hours ago.
	pbw.mu.Lock()
	pbw.lastActivity = time.Now().Add(-2 * time.Hour)
	pbw.mu.Unlock()

	if _, reaped := pbw.reapIfIdle(time.Now(), time.Hour); reaped {
		t.Fatal("A writer with a part in flight was reaped mid-upload")
	}
	if reaped := sweepAbandonedUploads(time.Now()); reaped != 0 {
		t.Fatalf("The sweep reaped %d writers that had a part in flight", reaped)
	}

	holdServer.mu.Lock()
	aborts := len(holdServer.AbortCalls)
	holdServer.mu.Unlock()
	if aborts != 0 {
		t.Fatalf("The sweep aborted a live upload's session (%d aborts)", aborts)
	}

	gate.open()
	waitForParts(pbw)

	// The part landed, so the writer is idle again and reapable. Its own clock
	// was reset by the landing, hence the future now.
	if _, reaped := pbw.reapIfIdle(time.Now().Add(2*time.Hour), time.Hour); !reaped {
		t.Error("Expected the writer to be reapable once its part had landed")
	}

	holdServer.mu.Lock()
	aborts = len(holdServer.AbortCalls)
	holdServer.mu.Unlock()
	if aborts != 1 {
		t.Errorf("Expected the reap to abort the session once, got %d aborts", aborts)
	}
	if held := bytesHeld(); held != 0 {
		t.Errorf("Expected the budget back after the reap, got %d bytes held", held)
	}
}

// TestCancel_ConcurrentWithWrite exists for -race. Cancel now unlocks the writer
// while it waits for a part, and the part goroutine takes the lock to record its
// result, so there are three parties on w.mu instead of two. Whoever wins, the
// budget must come back whole and nothing may be left in flight.
func TestCancel_ConcurrentWithWrite(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)

	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
			}

			var inner sync.WaitGroup
			inner.Add(2)
			go func() {
				defer inner.Done()
				// Enough to hand a part off, so a Cancel can land while one is
				// in flight. A cancelled writer refuses writes, which is a
				// legitimate outcome of this race.
				chunk := generateTestData(maxBufferSize / 4)
				for range 6 {
					if _, err := writer.Write(chunk); err != nil {
						return
					}
				}
			}()
			go func() {
				defer inner.Done()
				writer.Cancel(context.Background())
			}()
			inner.Wait()

			// A second Cancel must be harmless, and must not double-abort.
			writer.Cancel(context.Background())
		}()
	}
	writers.Wait()

	if held := bytesHeld(); held != 0 {
		t.Errorf("Expected no budget held once every writer was cancelled, got %d", held)
	}
	if inFlight, _ := UploadStats(); inFlight != 0 {
		t.Errorf("Expected no uploads left in the map, got %d", inFlight)
	}
}