mirror of
https://tangled.org/evan.jarrett.net/at-container-registry
synced 2026-09-20 01:04:15 +00:00
appview: keep one part upload in flight while the next buffer fills
A large layer went up strictly one step at a time: fill 16MB from the client, stop reading, fetch a part URL and PUT the part to S3, reset, resume reading. While the part was in flight Docker sat on a full TCP window; while the buffer filled S3 sat idle. Wall clock was receive time plus send time. The writer now hands a full buffer to a goroutine that does the hold call and the PUT, and keeps filling a second buffer from the client. When that one fills it waits for the previous part, takes its buffer back, and hands the new one off. At most one part is in flight, so part numbers and ETags stay ordered, and a blob that never fills a buffer never allocates the second one. No network runs under the writer lock on the happy path. Peak memory for a large upload is now two buffers, 32MB. Both are charged to the process budget through the existing accounting, the second as it grows, and the budget floor rises to match so a large upload can never be refused outright. The 512MB default holds sixteen. A failed part records a sticky error, closes the writer, and aborts the multipart from the goroutine that still holds the upload ID; the next Write, hand-off, or Commit reports the cause. Commit verifies the digest first, then waits for the flight, sends the final part, and completes. Cancel waits for the flight, bounded, before aborting so the abort cannot overtake a PUT that has not yet been issued its upload ID. The sweeper refuses to reap a writer with a part in flight, since last activity is only stamped when a part lands. Tests observe the overlap directly: the fake S3 blocks the first PUT and the second buffer's writes are asserted to return before it is released, while the third buffer's writes block. Also covered: the one-part-late error, Commit waiting, Cancel during flight, peak budget, the sweeper, and concurrent Cancel and Write under the race detector. The integration suite passed with a 72MB layer pushed through the pipeline by three clients. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018Yf1ZVA7sXYhQNb9tCo1m5
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
8b5f195f09
commit
56bde61555
+6
-4
@@ -147,13 +147,15 @@ jetstream.backfill_enabled → ATCR_JETSTREAM_BACKFILL_ENABLED
|
||||
|
||||
### 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:
|
||||
Each in-flight blob upload buffers up to 16MB in the AppView process, and a large
|
||||
upload holds two of those buffers at its peak (32MB), because one part is uploaded
|
||||
to S3 in the background while the next buffer fills; a blob small enough never to
|
||||
flush only ever holds one. 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_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 32MB (one writer's peak) if configured lower, since a smaller budget could never satisfy a single large 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
|
||||
|
||||
@@ -0,0 +1,658 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -668,28 +668,93 @@ type CompletedPart struct {
|
||||
ETag string `json:"etag"`
|
||||
}
|
||||
|
||||
// errWriterClosed is what every entry point reports once the writer is done
|
||||
// with, whether it was committed, cancelled, reaped, or closed by a part
|
||||
// upload that failed. Where there is a more specific cause (a failed part) the
|
||||
// writer reports that instead; see closedErr.
|
||||
var errWriterClosed = errors.New("writer closed")
|
||||
|
||||
// inFlightPart is the one part upload a writer may have running in the
|
||||
// background. The writer hands it a full buffer and carries straight on
|
||||
// filling a second one.
|
||||
//
|
||||
// Ownership of buf passes to the goroutine at hand-off and comes back to the
|
||||
// writer when done is closed. Nothing else may touch buf in between: the
|
||||
// goroutine is reading straight out of it, with no copy.
|
||||
type inFlightPart struct {
|
||||
// number is the S3 part number this upload was assigned. Assigned under
|
||||
// w.mu at hand-off, so part numbers follow the order the bytes arrived in.
|
||||
number int
|
||||
|
||||
// buf is the full buffer being uploaded.
|
||||
buf *bytes.Buffer
|
||||
|
||||
// done is closed once the goroutine has recorded its outcome under w.mu
|
||||
// and handed buf back. A reader that takes w.mu after observing this close
|
||||
// sees everything the goroutine wrote.
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// ProxyBlobWriter implements distribution.BlobWriter for proxy uploads.
|
||||
//
|
||||
// Small blobs (everything that still fits in buffer at Commit) are PUT once to
|
||||
// their final content-addressed key. Larger ones fall back to an S3 multipart
|
||||
// upload, started on the first flush.
|
||||
//
|
||||
// A large blob keeps one part in flight while the next buffer fills. The
|
||||
// upload used to be strictly serial: while a part was being PUT to S3 the
|
||||
// Docker PATCH body was not being drained, and while the buffer filled S3 sat
|
||||
// idle, so the wall clock for a layer was receive time plus send time. The
|
||||
// price is a second buffer, so a large upload's peak memory is
|
||||
// 2 * maxBufferSize (32MB), not one buffer's worth.
|
||||
type ProxyBlobWriter struct {
|
||||
store *ProxyBlobStore
|
||||
options distribution.CreateOptions
|
||||
uploadID string // S3 multipart upload ID; empty until the first flush starts one
|
||||
parts []CompletedPart // Track uploaded parts with ETags
|
||||
partNumber int // Current part number (starts at 1)
|
||||
buffer *bytes.Buffer // Buffer for current part
|
||||
partNumber int // Next part number to hand out (starts at 1)
|
||||
buffer *bytes.Buffer // Buffer currently being filled by Write
|
||||
digester digest.Digester // Hashes every byte written, for verification at Commit
|
||||
size int64 // Total bytes written
|
||||
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.
|
||||
// flight is the part upload running in the background, or nil when none
|
||||
// is. At most one, ever: when the buffer fills and this is not nil, the
|
||||
// write waits for it before handing off the next part. Keeping it to one
|
||||
// is what makes part numbers and ETags ordered for free.
|
||||
flight *inFlightPart
|
||||
|
||||
// spare is the second buffer while it is idle: handed back by a finished
|
||||
// part upload with its capacity intact, waiting to become the next part.
|
||||
// nil until the first hand-off, so a blob that never flushes never
|
||||
// allocates it.
|
||||
spare *bytes.Buffer
|
||||
|
||||
// flightErr is the failure of a background part upload, kept so it can be
|
||||
// reported by whatever touches the writer next: a Write, the wait for the
|
||||
// in-flight slot, or Commit. Sticky, so the cause stays visible instead of
|
||||
// degrading into a bare "writer closed" one call later.
|
||||
flightErr error
|
||||
|
||||
// aborted records that the hold-side multipart session has already been
|
||||
// aborted. A failing part aborts from its own goroutine, and Cancel,
|
||||
// Commit and the sweeper may all arrive afterwards; the session must only
|
||||
// be aborted once.
|
||||
aborted bool
|
||||
|
||||
// mu guards everything above against the abandoned-upload sweeper and
|
||||
// against the writer's own background part upload. Distribution hands a
|
||||
// given upload to one request at a time, so Write, Commit and Cancel never
|
||||
// contend with each other.
|
||||
//
|
||||
// Nothing that talks to the network on the happy path is done under this
|
||||
// lock: the hold calls and the S3 PUT of a part run in runPart with the
|
||||
// lock dropped, which is the entire point of the change (a held lock there
|
||||
// would serialise the upload again through Cancel and the sweeper, and
|
||||
// would pin the writer for minutes). The lock is taken only to swap
|
||||
// buffers and update state.
|
||||
mu sync.Mutex
|
||||
|
||||
// charged is the buffer budget this writer currently holds, in bytes. It is
|
||||
@@ -755,18 +820,45 @@ func (w *ProxyBlobWriter) projectedCap(n int) int64 {
|
||||
return min(max(need, 2*have), maxBufferSize)
|
||||
}
|
||||
|
||||
// otherBufferCap is the backing array held by the writer's second buffer: the
|
||||
// one currently being uploaded, or the one a finished upload handed back. Zero
|
||||
// until the first hand-off, because until then there is no second buffer.
|
||||
//
|
||||
// Reading Cap on a buffer that is in flight is safe: runPart only reads out of
|
||||
// it (Bytes), and the Reset that hands it back is done under w.mu.
|
||||
//
|
||||
// Callers hold w.mu.
|
||||
func (w *ProxyBlobWriter) otherBufferCap() int64 {
|
||||
if w.flight != nil {
|
||||
return int64(w.flight.buf.Cap())
|
||||
}
|
||||
if w.spare != nil {
|
||||
return int64(w.spare.Cap())
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// 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
|
||||
// What is charged is the backing arrays, not the bytes written: the buffer
|
||||
// being filled, plus the second buffer from the moment one exists at all. The
|
||||
// second is charged as it grows, through this same path, because it is created
|
||||
// empty at hand-off and only reaches full size as writes land in it.
|
||||
//
|
||||
// So a large upload now costs 2 * maxBufferSize, 32MB, not 16MB: that is the
|
||||
// price of holding one part in flight while the next one fills. A blob that
|
||||
// never fills a buffer never causes a hand-off, never allocates a second
|
||||
// buffer, and is charged only for what its one buffer grew to.
|
||||
//
|
||||
// Nothing is 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.
|
||||
// holds every byte of it. Everything goes back at once on Commit, Cancel or a
|
||||
// reap.
|
||||
//
|
||||
// Callers hold w.mu.
|
||||
func (w *ProxyBlobWriter) budgetDelta(n int) (int64, error) {
|
||||
want := w.projectedCap(n)
|
||||
want := w.otherBufferCap() + w.projectedCap(n)
|
||||
if want <= w.charged {
|
||||
return 0, nil
|
||||
}
|
||||
@@ -780,27 +872,52 @@ func (w *ProxyBlobWriter) budgetDelta(n int) (int64, error) {
|
||||
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.
|
||||
// releaseBudget returns everything this writer holds, both buffers included.
|
||||
// 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
|
||||
}
|
||||
|
||||
// closedErr is what a closed writer reports. A writer closed by a part upload
|
||||
// that failed reports that failure instead of the generic sentinel, so the
|
||||
// cause is not lost behind the closure it caused. Callers hold w.mu.
|
||||
func (w *ProxyBlobWriter) closedErr() error {
|
||||
if w.flightErr != nil {
|
||||
return w.flightErr
|
||||
}
|
||||
return errWriterClosed
|
||||
}
|
||||
|
||||
// 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.
|
||||
// A writer someone is actively inside is by definition not abandoned. Now
|
||||
// that the part upload runs off the lock this is a short window (a buffer
|
||||
// swap, some bookkeeping) rather than the length of a PUT, but skipping and
|
||||
// looking again on the next sweep is still the right answer.
|
||||
if !w.mu.TryLock() {
|
||||
return 0, false
|
||||
}
|
||||
defer w.mu.Unlock()
|
||||
|
||||
// A part on its way to S3 is a live upload, and it is liveness the clock
|
||||
// cannot see: lastActivity is only bumped when a part lands, so a writer
|
||||
// whose last Write was longer ago than the timeout may still be mid-PUT.
|
||||
// Reaping it would abort the multipart session out from under a goroutine
|
||||
// that is still feeding it, and the abandoned-upload sweep exists to
|
||||
// reclaim clients that hung up, not uploads that are working.
|
||||
//
|
||||
// This cannot hide a genuinely abandoned writer forever: runPart's context
|
||||
// is bounded by uploadPartTimeout, so the part always lands or fails, and
|
||||
// the next sweep sees a writer with no flight and a stale lastActivity.
|
||||
if w.flight != nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
idle := now.Sub(w.lastActivity)
|
||||
if w.closed || idle < ttl {
|
||||
return 0, false
|
||||
@@ -817,19 +934,62 @@ func (w *ProxyBlobWriter) reapIfIdle(now time.Time, ttl time.Duration) (time.Dur
|
||||
}
|
||||
|
||||
// Write writes data to the upload.
|
||||
// Buffers data and flushes a part once the buffer reaches maxBufferSize.
|
||||
//
|
||||
// Bytes are buffered until the buffer reaches maxBufferSize, at which point it
|
||||
// is handed to a background part upload and this call carries straight on into
|
||||
// a second buffer. At most one part is ever in flight: when the second buffer
|
||||
// fills too, the write blocks until the running part lands, takes its buffer
|
||||
// back, and hands the second one off.
|
||||
//
|
||||
// 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, handing off, 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.
|
||||
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.
|
||||
written := 0
|
||||
for written < len(p) {
|
||||
// Each pass takes at most what fits in the current buffer, so a hand-off
|
||||
// always happens at a buffer boundary. The buffer is empty again on the
|
||||
// next pass, so a pass can never accept zero bytes and spin.
|
||||
n, err := w.writeChunk(p[written:])
|
||||
written += n
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
}
|
||||
return written, nil
|
||||
}
|
||||
|
||||
// writeChunk is one pass of Write: it takes as much of p as fits in the
|
||||
// current buffer, charges the budget for it, appends it, and hands the buffer
|
||||
// off to a background part upload if that filled it.
|
||||
//
|
||||
// The budget acquire happens with the lock dropped. Blocking there is the
|
||||
// point (it is backpressure on the Docker client), but blocking with w.mu held
|
||||
// would make Cancel, the sweeper and the writer's own part goroutine queue
|
||||
// behind a write that is waiting on memory nobody has yet returned.
|
||||
func (w *ProxyBlobWriter) writeChunk(p []byte) (int, error) {
|
||||
w.mu.Lock()
|
||||
if w.closed {
|
||||
// Includes the one-part-late case: a background part that failed closed
|
||||
// the writer, and closedErr reports why rather than just that it is shut.
|
||||
err := w.closedErr()
|
||||
w.mu.Unlock()
|
||||
return 0, fmt.Errorf("writer closed")
|
||||
return 0, err
|
||||
}
|
||||
w.lastActivity = time.Now()
|
||||
delta, err := w.budgetDelta(len(p))
|
||||
|
||||
chunk := p
|
||||
if room := maxBufferSize - w.buffer.Len(); len(chunk) > room {
|
||||
chunk = chunk[:room]
|
||||
}
|
||||
|
||||
delta, err := w.budgetDelta(len(chunk))
|
||||
waitCtx := w.waitContext()
|
||||
w.mu.Unlock()
|
||||
|
||||
@@ -844,54 +1004,35 @@ func (w *ProxyBlobWriter) Write(p []byte) (int, error) {
|
||||
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.
|
||||
// Cancelled, reaped, or closed by a failing part 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 did.
|
||||
uploadBudget.release(delta)
|
||||
return 0, fmt.Errorf("writer closed")
|
||||
return 0, w.closedErr()
|
||||
}
|
||||
// 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
|
||||
// from one request at a time, so nothing can have charged the buffers 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))
|
||||
|
||||
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)
|
||||
// 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(chunk[:n])
|
||||
|
||||
// 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(chunk[:n])
|
||||
|
||||
// Flush once the buffer is full (S3 part size)
|
||||
if w.buffer.Len() >= maxBufferSize {
|
||||
if err := w.flushPart(); err != nil {
|
||||
return written, err
|
||||
}
|
||||
if w.buffer.Len() >= maxBufferSize {
|
||||
if err := w.handOffBuffer(); err != nil {
|
||||
return n, err
|
||||
}
|
||||
}
|
||||
|
||||
return written, nil
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// growBuffer sizes the buffer's backing array ahead of an n byte write so that
|
||||
@@ -927,50 +1068,200 @@ func (w *ProxyBlobWriter) growBuffer(n int) {
|
||||
w.buffer.Grow(maxBufferSize - w.buffer.Len())
|
||||
}
|
||||
|
||||
// flushPart uploads the current buffer as a part
|
||||
func (w *ProxyBlobWriter) flushPart() error {
|
||||
// handOffBuffer gives the full buffer to a background part upload and leaves
|
||||
// the writer a second, empty buffer to keep filling.
|
||||
//
|
||||
// Only one part is ever in flight, so if a previous one is still running this
|
||||
// waits for it first and reuses the buffer it hands back. That wait is where a
|
||||
// client outrunning S3 gets its backpressure, and it is also where a part that
|
||||
// failed a buffer ago is finally reported.
|
||||
//
|
||||
// Callers hold w.mu. The lock is dropped while waiting and held again on
|
||||
// return.
|
||||
func (w *ProxyBlobWriter) handOffBuffer() error {
|
||||
if w.buffer.Len() == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
if err := w.drainFlight(); err != nil {
|
||||
return err
|
||||
}
|
||||
if w.closed {
|
||||
// Cancelled or reaped while this call waited for the previous part.
|
||||
return w.closedErr()
|
||||
}
|
||||
|
||||
// The buffer the finished part handed back, capacity intact
|
||||
// (bytes.Buffer.Reset keeps the backing array), so a large upload allocates
|
||||
// its two buffers once and then just swaps them. nil only on the very first
|
||||
// hand-off: a blob that never fills a buffer never gets here, and so never
|
||||
// allocates a second one.
|
||||
next := w.spare
|
||||
w.spare = nil
|
||||
if next == nil {
|
||||
next = &bytes.Buffer{}
|
||||
}
|
||||
|
||||
f := &inFlightPart{
|
||||
number: w.partNumber,
|
||||
buf: w.buffer,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
w.partNumber++
|
||||
w.buffer = next
|
||||
w.flight = f
|
||||
|
||||
go w.runPart(f)
|
||||
return nil
|
||||
}
|
||||
|
||||
// drainFlight waits for the in-flight part, if there is one, and reports
|
||||
// whether it failed. The failure is sticky, so every later caller sees the
|
||||
// cause rather than just a closed writer.
|
||||
//
|
||||
// Callers hold w.mu. The lock is dropped for the wait, which it has to be:
|
||||
// runPart needs it to record its result and hand the buffer back.
|
||||
func (w *ProxyBlobWriter) drainFlight() error {
|
||||
if f := w.flight; f != nil {
|
||||
w.mu.Unlock()
|
||||
<-f.done
|
||||
w.mu.Lock()
|
||||
}
|
||||
return w.flightErr
|
||||
}
|
||||
|
||||
// abandonFlight waits for the in-flight part the way Cancel needs to: bounded,
|
||||
// so a wedged PUT cannot pin a client's DELETE, and reporting whether the
|
||||
// goroutine actually finished.
|
||||
//
|
||||
// Cancel must not abort the multipart session while a part PUT is still
|
||||
// running. An abort that overtakes the PUT can arrive before the hold has even
|
||||
// issued the upload ID, and then there is nothing left that knows the session
|
||||
// exists: it leaks in S3 until the bucket's own multipart expiry catches it.
|
||||
//
|
||||
// Callers hold w.mu. The lock is dropped for the wait.
|
||||
func (w *ProxyBlobWriter) abandonFlight() bool {
|
||||
f := w.flight
|
||||
if f == nil {
|
||||
return true
|
||||
}
|
||||
|
||||
timer := time.NewTimer(uploadFlightAbandonTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
w.mu.Unlock()
|
||||
finished := false
|
||||
select {
|
||||
case <-f.done:
|
||||
finished = true
|
||||
case <-timer.C:
|
||||
}
|
||||
w.mu.Lock()
|
||||
return finished
|
||||
}
|
||||
|
||||
// runPart uploads one part in the background and hands the buffer back.
|
||||
//
|
||||
// The context is detached and bounded rather than the request's: the PATCH
|
||||
// that filled this buffer has usually been answered by the time the PUT
|
||||
// finishes, and cancelling an upload because the request that produced its
|
||||
// bytes ended is exactly wrong.
|
||||
func (w *ProxyBlobWriter) runPart(f *inFlightPart) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), uploadPartTimeout)
|
||||
defer cancel()
|
||||
|
||||
etag, err := w.uploadPart(ctx, f.number, f.buf.Bytes())
|
||||
|
||||
w.mu.Lock()
|
||||
if err != nil {
|
||||
// Nobody was watching this goroutine, so the failure is reported one
|
||||
// part late: by the next Write, by the next hand-off's wait, or by
|
||||
// Commit. Close the writer here so no further bytes are accepted.
|
||||
w.flightErr = fmt.Errorf("part %d upload failed: %w", f.number, err)
|
||||
w.closed = true
|
||||
} else {
|
||||
// Ordered by construction: part number+1 is not handed off until this
|
||||
// append has happened, because hand-off waits on f.done first.
|
||||
w.parts = append(w.parts, CompletedPart{PartNumber: f.number, ETag: etag})
|
||||
slog.Debug("Part uploaded successfully", "component", "proxy_blob_store/runPart", "part_number", f.number, "etag", etag)
|
||||
}
|
||||
|
||||
// The buffer comes back to the writer with its capacity intact, ready to be
|
||||
// the next part's. No budget is released: the memory is still held.
|
||||
f.buf.Reset()
|
||||
w.spare = f.buf
|
||||
w.flight = nil
|
||||
|
||||
// A part landing is the writer doing work on the client's behalf, so it
|
||||
// counts as activity. Without this a writer whose only remaining job was a
|
||||
// slow PUT would keep ageing towards the sweeper's idle timeout.
|
||||
w.lastActivity = time.Now()
|
||||
|
||||
if err != nil {
|
||||
// Abort from here rather than leaving it to whoever notices the error:
|
||||
// the upload ID is known now, this goroutine is the last thing touching
|
||||
// the session, and the writer may never be touched again. abortIfStarted
|
||||
// is idempotent, so a later Cancel or Commit does not double-abort.
|
||||
abortCtx, abortCancel := context.WithTimeout(context.Background(), uploadAbortTimeout)
|
||||
w.abortIfStarted(abortCtx)
|
||||
abortCancel()
|
||||
}
|
||||
w.mu.Unlock()
|
||||
|
||||
close(f.done)
|
||||
}
|
||||
|
||||
// uploadPart starts the multipart session if this is the first part, asks the
|
||||
// hold for a presigned part URL, and PUTs the bytes to S3.
|
||||
//
|
||||
// No lock is held for any of it. That is what lets the next buffer fill while
|
||||
// this one is going up, and it is why body is passed in rather than read off
|
||||
// the writer: the caller owns those bytes for the duration.
|
||||
func (w *ProxyBlobWriter) uploadPart(ctx context.Context, partNumber int, body []byte) (string, error) {
|
||||
tempDigest := fmt.Sprintf("uploads/temp-%s", w.id)
|
||||
|
||||
// Start the multipart upload on the first flush rather than in Create. A
|
||||
w.mu.Lock()
|
||||
uploadID := w.uploadID
|
||||
w.mu.Unlock()
|
||||
|
||||
// Start the multipart upload on the first part rather than in Create. A
|
||||
// blob that never fills the buffer is committed with a single direct PUT
|
||||
// and needs no multipart session, no temp object and no server side copy.
|
||||
if w.uploadID == "" {
|
||||
uploadID, err := w.store.startMultipartUpload(ctx, tempDigest)
|
||||
// Only the first part can find this empty: parts are strictly serialised.
|
||||
if uploadID == "" {
|
||||
id, err := w.store.startMultipartUpload(ctx, tempDigest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to start multipart upload: %w", err)
|
||||
return "", fmt.Errorf("failed to start multipart upload: %w", err)
|
||||
}
|
||||
uploadID = id
|
||||
|
||||
w.mu.Lock()
|
||||
w.uploadID = uploadID
|
||||
w.mu.Unlock()
|
||||
}
|
||||
|
||||
// Get structured upload info for this part
|
||||
uploadInfo, err := w.store.getPartUploadInfo(ctx, tempDigest, w.uploadID, w.partNumber)
|
||||
uploadInfo, err := w.store.getPartUploadInfo(ctx, tempDigest, uploadID, partNumber)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get part upload info: %w", err)
|
||||
return "", fmt.Errorf("failed to get part upload info: %w", err)
|
||||
}
|
||||
|
||||
// Upload part to S3 presigned URL
|
||||
req, err := http.NewRequestWithContext(ctx, "PUT", uploadInfo.URL, bytes.NewReader(w.buffer.Bytes()))
|
||||
req, err := http.NewRequestWithContext(ctx, "PUT", uploadInfo.URL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/octet-stream")
|
||||
|
||||
resp, err := w.store.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("part upload failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
|
||||
return "", fmt.Errorf("part upload failed: status %d, body: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
// Store ETag for completion
|
||||
@@ -984,21 +1275,40 @@ func (w *ProxyBlobWriter) flushPart() error {
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err == nil && result.ETag != "" {
|
||||
etag = result.ETag
|
||||
} else {
|
||||
return fmt.Errorf("no ETag in response")
|
||||
return "", fmt.Errorf("no ETag in response")
|
||||
}
|
||||
}
|
||||
|
||||
w.parts = append(w.parts, CompletedPart{
|
||||
PartNumber: w.partNumber,
|
||||
ETag: etag,
|
||||
})
|
||||
return etag, nil
|
||||
}
|
||||
|
||||
slog.Debug("Part uploaded successfully", "component", "proxy_blob_store/flushPart", "part_number", w.partNumber, "etag", etag)
|
||||
// flushFinalPart uploads whatever is left in the buffer as the last part.
|
||||
//
|
||||
// Synchronous, unlike every other part: Commit has nothing left to overlap it
|
||||
// with, and the ETag has to be in w.parts before completeUpload names them.
|
||||
// Callers hold w.mu; the lock is dropped for the upload itself so a final PUT
|
||||
// does not pin the writer. That is safe because Commit has already marked the
|
||||
// writer closed and taken it out of globalUploads, so nothing else will accept
|
||||
// bytes for it or reap it while this runs.
|
||||
func (w *ProxyBlobWriter) flushFinalPart(ctx context.Context) error {
|
||||
if w.buffer.Len() == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reset buffer and increment part number
|
||||
w.buffer.Reset()
|
||||
number := w.partNumber
|
||||
w.partNumber++
|
||||
body := w.buffer.Bytes()
|
||||
|
||||
w.mu.Unlock()
|
||||
etag, err := w.uploadPart(ctx, number, body)
|
||||
w.mu.Lock()
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
w.parts = append(w.parts, CompletedPart{PartNumber: number, ETag: etag})
|
||||
w.buffer.Reset()
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -1009,10 +1319,13 @@ func (w *ProxyBlobWriter) ReadFrom(r io.Reader) (int64, error) {
|
||||
// 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
|
||||
var closedErr error
|
||||
if w.closed {
|
||||
closedErr = w.closedErr()
|
||||
}
|
||||
w.mu.Unlock()
|
||||
if closed {
|
||||
return 0, fmt.Errorf("writer closed")
|
||||
if closedErr != nil {
|
||||
return 0, closedErr
|
||||
}
|
||||
|
||||
// Read in chunks and flush when needed
|
||||
@@ -1049,15 +1362,15 @@ func (w *ProxyBlobWriter) Size() int64 {
|
||||
// Commit finalizes the upload.
|
||||
//
|
||||
// The digest the client sent is verified against the bytes actually received
|
||||
// before anything else happens, then the blob is finalized: a direct PUT to
|
||||
// the final key if it is all still buffered, otherwise a final part plus the
|
||||
// hold's completeUpload.
|
||||
// before anything else happens, then the part still going up (if any) is waited
|
||||
// for, then the blob is finalized: a direct PUT to 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")
|
||||
return distribution.Descriptor{}, w.closedErr()
|
||||
}
|
||||
w.closed = true
|
||||
w.lastActivity = time.Now()
|
||||
@@ -1076,12 +1389,23 @@ func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descript
|
||||
// nothing in storage. The digest is the blob's address in a shared,
|
||||
// content-addressed bucket, so a client that names its bytes wrongly would
|
||||
// otherwise overwrite or shadow someone else's layer.
|
||||
//
|
||||
// This runs before the wait for the in-flight part on purpose: the hash
|
||||
// covers every byte Write accepted, including the ones still going up, so
|
||||
// waiting would buy nothing and would only delay refusing a bad blob.
|
||||
// Aborting, on the other hand, has to wait (see drainFlight and
|
||||
// abandonFlight): an abort that overtakes a running PUT can leak the
|
||||
// session.
|
||||
if desc.Digest.Algorithm() != digest.Canonical {
|
||||
// The abort has to wait for the part, even though the verdict did not.
|
||||
_ = w.drainFlight()
|
||||
w.abortIfStarted(ctx)
|
||||
slog.Warn("Rejected blob with unsupported digest algorithm", "component", "proxy_blob_store/Commit", "algorithm", desc.Digest.Algorithm(), "id", w.id)
|
||||
return distribution.Descriptor{}, distribution.ErrBlobDigestUnsupported
|
||||
}
|
||||
if computed := w.digester.Digest(); computed != desc.Digest {
|
||||
// The abort has to wait for the part, even though the verdict did not.
|
||||
_ = w.drainFlight()
|
||||
w.abortIfStarted(ctx)
|
||||
slog.Warn("Rejected blob whose content does not match its digest", "component", "proxy_blob_store/Commit", "claimed", desc.Digest, "computed", computed, "size", w.size)
|
||||
return distribution.Descriptor{}, distribution.ErrBlobInvalidDigest{
|
||||
@@ -1090,6 +1414,16 @@ func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descript
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for the part still going up. Its ETag has to be in w.parts before
|
||||
// completeUpload names them, and a failure it hit has to be reported here
|
||||
// rather than swallowed. This is also what decides whether there is a
|
||||
// multipart session at all: the first part starts one, and it may not have
|
||||
// got that far yet.
|
||||
if err := w.drainFlight(); err != nil {
|
||||
w.abortIfStarted(ctx)
|
||||
return distribution.Descriptor{}, err
|
||||
}
|
||||
|
||||
// Nothing was ever flushed, so the whole blob is in memory and can go
|
||||
// straight to its final content-addressed key. This is the common case:
|
||||
// every config blob and the large majority of layers land here.
|
||||
@@ -1110,7 +1444,7 @@ func (w *ProxyBlobWriter) Commit(ctx context.Context, desc distribution.Descript
|
||||
// Flush any remaining buffered data
|
||||
if w.buffer.Len() > 0 {
|
||||
slog.Debug("Flushing final buffer", "component", "proxy_blob_store/Commit", "bytes", w.buffer.Len())
|
||||
if err := w.flushPart(); err != nil {
|
||||
if err := w.flushFinalPart(ctx); err != nil {
|
||||
// Try to abort multipart on error
|
||||
w.abortIfStarted(ctx)
|
||||
return distribution.Descriptor{}, fmt.Errorf("failed to flush final part: %w", err)
|
||||
@@ -1171,12 +1505,18 @@ func (w *ProxyBlobWriter) putDirect(ctx context.Context, dgst digest.Digest) err
|
||||
return nil
|
||||
}
|
||||
|
||||
// abortIfStarted aborts the multipart upload, if one was ever started. A writer
|
||||
// whose blob stayed inside the buffer has no session to abort.
|
||||
// abortIfStarted aborts the multipart upload, if one was ever started and has
|
||||
// not been aborted already. A writer whose blob stayed inside the buffer has no
|
||||
// session to abort.
|
||||
//
|
||||
// Exactly once, because there are now several ways to arrive here: a failing
|
||||
// part aborts from its own goroutine, and Cancel, Commit and the sweeper may
|
||||
// all follow it. Callers hold w.mu.
|
||||
func (w *ProxyBlobWriter) abortIfStarted(ctx context.Context) {
|
||||
if w.uploadID == "" {
|
||||
if w.uploadID == "" || w.aborted {
|
||||
return
|
||||
}
|
||||
w.aborted = true
|
||||
if err := w.store.abortMultipartUpload(ctx, w.uploadID); err != nil {
|
||||
slog.Warn("Failed to abort multipart upload", "component", "proxy_blob_store", "error", err)
|
||||
// Continue anyway - we want to mark upload as cancelled
|
||||
@@ -1198,6 +1538,16 @@ func (w *ProxyBlobWriter) Cancel(ctx context.Context) error {
|
||||
delete(globalUploads, w.id)
|
||||
globalUploadsMu.Unlock()
|
||||
|
||||
// Closed above, so no further bytes are taken; now let the part already on
|
||||
// its way to S3 finish before aborting the session it belongs to. Bounded:
|
||||
// a wedged PUT must not pin the client's DELETE, and if the wait does time
|
||||
// out the abort still goes out, since a session that may exist is worth one
|
||||
// best-effort abort.
|
||||
if !w.abandonFlight() {
|
||||
slog.Warn("Cancelling an upload whose part is still in flight",
|
||||
"component", "proxy_blob_store/Cancel", "id", w.id, "waited", uploadFlightAbandonTimeout)
|
||||
}
|
||||
|
||||
w.abortIfStarted(ctx)
|
||||
|
||||
slog.Debug("Upload cancelled", "component", "proxy_blob_store/Cancel", "id", w.id)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -674,6 +675,11 @@ type mockHoldServer struct {
|
||||
AbortError error
|
||||
PresignError error
|
||||
|
||||
// AbortHook, when set, runs as an abort request is handled. It exists so a
|
||||
// test can order the abort against other events (an in-flight part
|
||||
// finishing, say) rather than only counting aborts after the fact.
|
||||
AbortHook func()
|
||||
|
||||
// Response customization
|
||||
UploadID string
|
||||
}
|
||||
@@ -712,6 +718,12 @@ type mockS3Server struct {
|
||||
mu sync.Mutex
|
||||
Parts map[int][]byte
|
||||
|
||||
// PartsStarted records the part number of every multipart part PUT that
|
||||
// arrived, logged before PartHook runs. Parts records the bytes, and only
|
||||
// once the PUT is allowed to finish, so the two together tell a test which
|
||||
// parts are still in flight.
|
||||
PartsStarted []int
|
||||
|
||||
// DirectPuts records whole-blob PUTs to the final key (the non-multipart
|
||||
// path), separately from multipart parts.
|
||||
DirectPuts [][]byte
|
||||
@@ -720,10 +732,50 @@ type mockS3Server struct {
|
||||
// Error injection
|
||||
UploadError error
|
||||
|
||||
// PartHook, when set, runs on every multipart part PUT before its bytes are
|
||||
// recorded, with the server's own lock released. It can block, which is how
|
||||
// a test holds one part in flight while checking that the next buffer keeps
|
||||
// filling, and it can return an error to fail that particular part.
|
||||
PartHook func(partNumber int) error
|
||||
|
||||
// Response customization
|
||||
ETagInHeader bool // true = ETag in header, false = in JSON body
|
||||
}
|
||||
|
||||
// setPartHook installs the per-part hook under the server's lock, so setting it
|
||||
// from the test goroutine never races the handler reading it.
|
||||
func (m *mockS3Server) setPartHook(fn func(partNumber int) error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.PartHook = fn
|
||||
}
|
||||
|
||||
// startedParts is the part numbers whose PUT has arrived, finished or not.
|
||||
func (m *mockS3Server) startedParts() []int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return append([]int(nil), m.PartsStarted...)
|
||||
}
|
||||
|
||||
// finishedParts is the part numbers whose PUT completed, sorted.
|
||||
func (m *mockS3Server) finishedParts() []int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
nums := make([]int, 0, len(m.Parts))
|
||||
for n := range m.Parts {
|
||||
nums = append(nums, n)
|
||||
}
|
||||
sort.Ints(nums)
|
||||
return nums
|
||||
}
|
||||
|
||||
// part is the bytes recorded for a finished part.
|
||||
func (m *mockS3Server) part(n int) []byte {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.Parts[n]
|
||||
}
|
||||
|
||||
// newMockHoldServer creates a mock hold service
|
||||
func newMockHoldServer(t *testing.T, s3URL string) *mockHoldServer {
|
||||
m := &mockHoldServer{
|
||||
@@ -831,6 +883,9 @@ func newMockHoldServer(t *testing.T, s3URL string) *mockHoldServer {
|
||||
m.AbortCalls = append(m.AbortCalls, mockAbortCall{
|
||||
UploadID: body["uploadId"].(string),
|
||||
})
|
||||
if m.AbortHook != nil {
|
||||
m.AbortHook()
|
||||
}
|
||||
|
||||
w.WriteHeader(http.StatusOK)
|
||||
json.NewEncoder(w).Encode(map[string]any{})
|
||||
@@ -851,13 +906,20 @@ func newMockS3Server(t *testing.T, etagInHeader bool) *mockS3Server {
|
||||
ETagInHeader: etagInHeader,
|
||||
}
|
||||
|
||||
// The lock is taken in short stretches rather than held for the whole
|
||||
// handler: a part PUT can now be deliberately blocked, and blocking with
|
||||
// the server's lock held would freeze every other upload with it, which is
|
||||
// exactly the overlap these tests are trying to observe.
|
||||
m.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
uploadErr := m.UploadError
|
||||
hook := m.PartHook
|
||||
etagInHeader := m.ETagInHeader
|
||||
m.mu.Unlock()
|
||||
|
||||
if m.UploadError != nil {
|
||||
if uploadErr != nil {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
w.Write([]byte(m.UploadError.Error()))
|
||||
w.Write([]byte(uploadErr.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -866,20 +928,37 @@ func newMockS3Server(t *testing.T, etagInHeader bool) *mockS3Server {
|
||||
|
||||
// A whole-blob PUT lands on /blob and carries no part number.
|
||||
if r.URL.Path == "/blob" {
|
||||
m.mu.Lock()
|
||||
m.DirectPuts = append(m.DirectPuts, body)
|
||||
m.DirectContentTypes = append(m.DirectContentTypes, r.Header.Get("Content-Type"))
|
||||
m.mu.Unlock()
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
|
||||
// Parse part number from URL
|
||||
partNum, _ := strconv.Atoi(r.URL.Query().Get("partNumber"))
|
||||
|
||||
m.mu.Lock()
|
||||
m.PartsStarted = append(m.PartsStarted, partNum)
|
||||
m.mu.Unlock()
|
||||
|
||||
if hook != nil {
|
||||
if err := hook(partNum); err != nil {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
w.Write([]byte(err.Error()))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
m.Parts[partNum] = body
|
||||
m.mu.Unlock()
|
||||
|
||||
// Generate ETag
|
||||
etag := fmt.Sprintf(`"etag-part-%d"`, partNum)
|
||||
|
||||
if m.ETagInHeader {
|
||||
if etagInHeader {
|
||||
w.Header().Set("ETag", etag)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
} else {
|
||||
@@ -908,6 +987,25 @@ func createTestProxyBlobStore(t *testing.T, holdURL string) *ProxyBlobStore {
|
||||
return store
|
||||
}
|
||||
|
||||
// waitForParts blocks until the writer has no part upload in flight.
|
||||
//
|
||||
// Filling the buffer no longer uploads the part before Write returns: the full
|
||||
// buffer is handed to a background goroutine and Write carries on into a second
|
||||
// buffer. Any test that inspects the hold's or S3's call log, or the writer's
|
||||
// recorded parts, straight after a Write has to wait for that goroutine first.
|
||||
// Commit and Cancel do the waiting themselves, so they need no help.
|
||||
func waitForParts(w *ProxyBlobWriter) {
|
||||
for {
|
||||
w.mu.Lock()
|
||||
f := w.flight
|
||||
w.mu.Unlock()
|
||||
if f == nil {
|
||||
return
|
||||
}
|
||||
<-f.done
|
||||
}
|
||||
}
|
||||
|
||||
// generateTestData creates n bytes of predictable test data
|
||||
func generateTestData(n int) []byte {
|
||||
data := make([]byte, n)
|
||||
@@ -966,8 +1064,9 @@ func TestCreate_Success(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestInitiate_HoldErrorSurfacesAtFirstFlush pins where a failing initiateUpload
|
||||
// is now reported. Create no longer calls the hold, so the failure shows up on
|
||||
// the first write that fills the buffer.
|
||||
// is now reported. Create no longer calls the hold, and the first part goes up
|
||||
// in the background, so the failure shows up one part late: on the write that
|
||||
// follows it, not on the write that filled the buffer.
|
||||
func TestInitiate_HoldErrorSurfacesAtFirstFlush(t *testing.T) {
|
||||
s3Server := newMockS3Server(t, true)
|
||||
defer s3Server.Close()
|
||||
@@ -984,9 +1083,14 @@ func TestInitiate_HoldErrorSurfacesAtFirstFlush(t *testing.T) {
|
||||
}
|
||||
defer writer.Cancel(context.Background())
|
||||
|
||||
_, err = writer.Write(generateTestData(maxBufferSize))
|
||||
if _, err = writer.Write(generateTestData(maxBufferSize)); err != nil {
|
||||
t.Fatalf("The write that fills the buffer hands the part off and returns: %v", err)
|
||||
}
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
_, err = writer.Write(generateTestData(1024))
|
||||
if err == nil {
|
||||
t.Fatal("Expected error from the flush that starts the multipart upload")
|
||||
t.Fatal("Expected the failed initiate to surface on the next write")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "hold service unavailable") {
|
||||
@@ -1070,6 +1174,7 @@ func TestWrite_TriggerFlush(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("Write() failed: %v", err)
|
||||
}
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
// Verify flush occurred (1 part uploaded)
|
||||
s3Server.mu.Lock()
|
||||
@@ -1122,6 +1227,8 @@ func TestWrite_MultipleFlushes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
// Verify 2 flushes occurred, half a buffer remains
|
||||
s3Server.mu.Lock()
|
||||
partCount := len(s3Server.Parts)
|
||||
@@ -1201,6 +1308,7 @@ func TestFlushPart_Success(t *testing.T) {
|
||||
|
||||
// Get the internal writer to check parts
|
||||
pbw := writer.(*ProxyBlobWriter)
|
||||
waitForParts(pbw)
|
||||
if len(pbw.parts) != 1 {
|
||||
t.Errorf("Expected 1 part recorded, got %d", len(pbw.parts))
|
||||
}
|
||||
@@ -1239,6 +1347,7 @@ func TestFlushPart_ETagInJSON(t *testing.T) {
|
||||
|
||||
// Verify part was recorded with ETag from JSON
|
||||
pbw := writer.(*ProxyBlobWriter)
|
||||
waitForParts(pbw)
|
||||
if len(pbw.parts) != 1 {
|
||||
t.Errorf("Expected 1 part recorded, got %d", len(pbw.parts))
|
||||
}
|
||||
@@ -1265,10 +1374,15 @@ func TestFlushPart_HoldError(t *testing.T) {
|
||||
}
|
||||
defer writer.Cancel(context.Background())
|
||||
|
||||
// Write enough to trigger flush
|
||||
// Write enough to trigger flush. The part goes up in the background, so the
|
||||
// failure lands on the write after it, not on this one.
|
||||
data := generateTestData(maxBufferSize)
|
||||
_, err = writer.Write(data)
|
||||
if _, err = writer.Write(data); err != nil {
|
||||
t.Fatalf("The write that hands the part off should succeed: %v", err)
|
||||
}
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
_, err = writer.Write(generateTestData(1024))
|
||||
if err == nil {
|
||||
t.Fatal("Expected error when hold service fails")
|
||||
}
|
||||
@@ -1295,10 +1409,15 @@ func TestFlushPart_S3Error(t *testing.T) {
|
||||
}
|
||||
defer writer.Cancel(context.Background())
|
||||
|
||||
// Write enough to trigger flush
|
||||
// Write enough to trigger flush. The part goes up in the background, so the
|
||||
// failure lands on the write after it, not on this one.
|
||||
data := generateTestData(maxBufferSize)
|
||||
_, err = writer.Write(data)
|
||||
if _, err = writer.Write(data); err != nil {
|
||||
t.Fatalf("The write that hands the part off should succeed: %v", err)
|
||||
}
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
_, err = writer.Write(generateTestData(1024))
|
||||
if err == nil {
|
||||
t.Fatal("Expected error when S3 fails")
|
||||
}
|
||||
@@ -1329,10 +1448,15 @@ func TestFlushPart_NoETag(t *testing.T) {
|
||||
}
|
||||
defer writer.Cancel(context.Background())
|
||||
|
||||
// Write enough to trigger flush
|
||||
// Write enough to trigger flush. The part goes up in the background, so the
|
||||
// failure lands on the write after it, not on this one.
|
||||
data := generateTestData(maxBufferSize)
|
||||
_, err = writer.Write(data)
|
||||
if _, err = writer.Write(data); err != nil {
|
||||
t.Fatalf("The write that hands the part off should succeed: %v", err)
|
||||
}
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
_, err = writer.Write(generateTestData(1024))
|
||||
if err == nil {
|
||||
t.Fatal("Expected error when no ETag is returned")
|
||||
}
|
||||
@@ -1411,6 +1535,7 @@ func TestReadFrom_LargeFile(t *testing.T) {
|
||||
}
|
||||
|
||||
// Verify 2 flushes occurred
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
s3Server.mu.Lock()
|
||||
partCount := len(s3Server.Parts)
|
||||
s3Server.mu.Unlock()
|
||||
@@ -1543,6 +1668,7 @@ func TestCommit_WithRemainingBuffer(t *testing.T) {
|
||||
}
|
||||
|
||||
// At this point, 1 full part should be uploaded, half a buffer remains
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
s3Server.mu.Lock()
|
||||
partsBeforeCommit := len(s3Server.Parts)
|
||||
s3Server.mu.Unlock()
|
||||
@@ -1609,6 +1735,10 @@ func TestCommit_FlushError(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// The first part is already on its way; let it land so the injected error
|
||||
// hits the final part rather than that one.
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
|
||||
// Inject error for final flush
|
||||
holdServer.mu.Lock()
|
||||
holdServer.PartURLError = fmt.Errorf("flush error")
|
||||
@@ -1919,6 +2049,7 @@ func TestFullUploadFlow_Multipart(t *testing.T) {
|
||||
}
|
||||
|
||||
// Verify 2 parts uploaded during write (two full buffers)
|
||||
waitForParts(writer.(*ProxyBlobWriter))
|
||||
s3Server.mu.Lock()
|
||||
partsBeforeCommit := len(s3Server.Parts)
|
||||
s3Server.mu.Unlock()
|
||||
@@ -2886,14 +3017,25 @@ func TestWrite_NeverOvershootsTheThreshold(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
waitForParts(pbw)
|
||||
|
||||
if c := pbw.buffer.Cap(); c > maxBufferSize {
|
||||
t.Errorf("Buffer capacity overshot the threshold: cap %d > %d", c, maxBufferSize)
|
||||
}
|
||||
if c := pbw.otherBufferCap(); c > maxBufferSize {
|
||||
t.Errorf("Second 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)
|
||||
// The charge is exactly the two backing arrays, and both are capped
|
||||
// at the threshold, so a writer can never hold more than
|
||||
// maxWriterFootprint however the bytes were chunked.
|
||||
if want := pbw.otherBufferCap() + int64(pbw.buffer.Cap()); pbw.charged != want {
|
||||
t.Errorf("Expected the charge to be the two backing arrays, %d, got %d", want, pbw.charged)
|
||||
}
|
||||
if pbw.charged > maxWriterFootprint {
|
||||
t.Errorf("Charge %d exceeds a writer's peak of %d", pbw.charged, maxWriterFootprint)
|
||||
}
|
||||
|
||||
s3Server.mu.Lock()
|
||||
|
||||
@@ -10,13 +10,24 @@ import (
|
||||
)
|
||||
|
||||
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.
|
||||
maxWriterBuffers = 2
|
||||
|
||||
// maxWriterFootprint is a writer's peak buffer memory, 32MB. It is what a
|
||||
// large upload costs now that a part is uploaded in the background while
|
||||
// the next buffer fills; before that overlap a writer peaked at one buffer.
|
||||
maxWriterFootprint = maxWriterBuffers * maxBufferSize
|
||||
|
||||
// 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.
|
||||
// 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. 512MB is 16 large uploads at their peak (or 32 that never
|
||||
// grew a second buffer), 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
|
||||
@@ -34,6 +45,21 @@ const (
|
||||
// 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
|
||||
@@ -58,9 +84,9 @@ func newBufferBudget(limit int64) *bufferBudget {
|
||||
}
|
||||
|
||||
// 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.
|
||||
// the wait cap expires. A writer never charges more than maxWriterFootprint in
|
||||
// total, and the budget is never smaller than that, so a single writer can
|
||||
// always eventually be satisfied.
|
||||
func (b *bufferBudget) acquire(ctx context.Context, n int64) error {
|
||||
if n <= 0 {
|
||||
return nil
|
||||
@@ -105,14 +131,16 @@ var (
|
||||
// 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.
|
||||
// A budget below maxWriterFootprint is raised to it. Anything less could never
|
||||
// be acquired by even a single large writer, so it would not throttle uploads,
|
||||
// it would stall every one of them until the wait cap expired. The floor is two
|
||||
// buffers rather than one because a writer at its peak holds two: the one being
|
||||
// filled and the one being uploaded.
|
||||
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 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
|
||||
|
||||
@@ -196,16 +196,17 @@ func TestUploadBudget_BlockedWriteHonoursRequestCancellation(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// 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) {
|
||||
// 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 != maxBufferSize {
|
||||
t.Errorf("Expected a 4MB budget to be raised to %d, got %d", maxBufferSize, uploadBudget.limit)
|
||||
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)
|
||||
@@ -213,16 +214,17 @@ func TestConfigureUploads_ClampsBudgetToOneBuffer(t *testing.T) {
|
||||
|
||||
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)
|
||||
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: 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.
|
||||
// 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()
|
||||
@@ -250,8 +252,9 @@ func TestUploadBudget_HeldAcrossPartsAndReturnedOnCommit(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
if held := bytesHeld(); held != maxBufferSize {
|
||||
t.Errorf("Expected exactly one buffer charged across all parts, got %d bytes held", held)
|
||||
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{
|
||||
@@ -334,10 +337,13 @@ func TestSweep_ReapsIdleWritersAndLeavesLiveOnes(t *testing.T) {
|
||||
}
|
||||
pbw := writer.(*ProxyBlobWriter)
|
||||
|
||||
// Enough to flush a part, which is what opens the multipart session.
|
||||
// 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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user