serve s3: fix CompleteMultipartUpload failing with InvalidPart while parts are being sent

Completion now waits for acknowledged parts still being written instead of failing.

Fixes #10257

cmd/serve/s3/multipart.go

@@ -912,11 +912,13 @@ func (b *s3Backend) CompleteMultipartUpload(ctx context.Context, bucketName, obj
 
 	// close commits the upload, failing with ErrInvalidPart and the
 	// upload left open if the streamed parts don't form the complete
-	// object. Any other failure is of the commit, so the upload can't be
-	// retried and forgetUpload tears it down. (After a failed commit the
-	// abort it does is a no-op.)
-	if err := up.close(); err != nil {
-		if !errors.Is(err, gofakes3.ErrInvalidPart) {
+	// object, with ErrSlowDown if acknowledged parts are still being
+	// written, or with the context's error if the request is cancelled
+	// while waiting for them. Any other failure is of the commit, so the
+	// upload can't be retried and forgetUpload tears it down. (After a
+	// failed commit the abort it does is a no-op.)
+	if err := up.close(ctx); err != nil {
+		if !errors.Is(err, gofakes3.ErrInvalidPart) && !errors.Is(err, gofakes3.ErrSlowDown) && err != ctx.Err() {
 			b.forgetUpload(uploadID, up)
 		}
 		return "", "", err
@@ -1138,11 +1140,30 @@ func (up *multipartUpload) validate(input *gofakes3.CompleteMultipartUploadReque
 // the upload is left open. The check and the commit share one critical
 // section so no late part can slip into the stream between them.
 //
+// Parts which were acknowledged but are still being written to the sink
+// are waited for. If that takes longer than multipartWaitTimeout,
+// gofakes3.ErrSlowDown is returned and the upload is left open. If ctx is
+// done first, ctx.Err() is returned and the upload is left open.
+//
 // Returns ErrNoSuchUpload if the upload was aborted by a concurrent
 // AbortMultipartUpload, in which case nothing has been committed - the
 // upload must not be reported as complete.
-func (up *multipartUpload) close() error {
+func (up *multipartUpload) close(ctx context.Context) error {
+	timer := time.NewTimer(multipartWaitTimeout)
+	defer timer.Stop()
 	up.mu.Lock()
+	for up.pumping && !up.closed {
+		changed := up.changed
+		up.mu.Unlock()
+		select {
+		case <-changed:
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-timer.C:
+			return gofakes3.ErrSlowDown
+		}
+		up.mu.Lock()
+	}
 	if up.closed {
 		aborted := up.aborted
 		up.mu.Unlock()

cmd/serve/s3/multipart_test.go

@@ -565,13 +565,13 @@ func TestMultipartCloseAfterAbort(t *testing.T) {
 	up := newMultipartUpload("bucket", "key", "bucket/key", "bucket/key", nil, 0)
 	up.fh = &stubSink{}
 	require.NoError(t, up.abort())
-	require.ErrorIs(t, up.close(), gofakes3.ErrNoSuchUpload)
+	require.ErrorIs(t, up.close(context.Background()), gofakes3.ErrNoSuchUpload)
 
 	// close stays idempotent after a successful close.
 	up = newMultipartUpload("bucket", "key", "bucket/key", "bucket/key", nil, 0)
 	up.fh = &stubSink{}
-	require.NoError(t, up.close())
-	require.NoError(t, up.close())
+	require.NoError(t, up.close(context.Background()))
+	require.NoError(t, up.close(context.Background()))
 }
 
 // TestMultipartCloseIncomplete checks that close refuses to commit while a
@@ -593,7 +593,7 @@ func TestMultipartCloseIncomplete(t *testing.T) {
 	require.Equal(t, turnBuffer, turn)
 	require.NoError(t, up.bufferPart(context.Background(), 2, int64(len(contents)), md5Sum[:], rw))
 
-	require.ErrorIs(t, up.close(), gofakes3.ErrInvalidPart)
+	require.ErrorIs(t, up.close(context.Background()), gofakes3.ErrInvalidPart)
 	assert.False(t, sink.closed)
 
 	// The upload is still open so abort tears it down.
@@ -1644,6 +1644,162 @@ func (s *flakySink) Write(p []byte) (int, error) {
 	return n, errBoom
 }
 
+// gatedSink is a sink which blocks writes past its first gateAt bytes until
+// release is closed, signalling blocked when the first such write starts.
+type gatedSink struct {
+	bufferSink
+	gateAt  int
+	once    sync.Once
+	blocked chan struct{}
+	release chan struct{}
+}
+
+func newGatedSink(gateAt int) *gatedSink {
+	return &gatedSink{gateAt: gateAt, blocked: make(chan struct{}), release: make(chan struct{})}
+}
+
+func (s *gatedSink) Write(p []byte) (int, error) {
+	s.mu.Lock()
+	written := s.buf.Len()
+	s.mu.Unlock()
+	if written+len(p) > s.gateAt {
+		s.once.Do(func() { close(s.blocked) })
+		<-s.release
+	}
+	return s.bufferSink.Write(p)
+}
+
+// startPumpBehindRetry recreates a client retrying part 1 because its
+// request was slow to answer: part 2 arrives first and is buffered, then
+// the request for part 1 streams it and blocks pumping part 2 into sink,
+// and the retry of part 1 is acknowledged. It returns the result of the
+// original part 1 request.
+func startPumpBehindRetry(t *testing.T, b *s3Backend, bucket, object string, uploadID gofakes3.UploadID, sink *gatedSink, part1, part2 []byte) <-chan error {
+	ctx := context.Background()
+	_, err := b.UploadPart(ctx, bucket, object, uploadID, 2, int64(len(part2)), bytes.NewReader(part2))
+	require.NoError(t, err)
+	part1Done := make(chan error, 1)
+	go func() {
+		_, err := b.UploadPart(ctx, bucket, object, uploadID, 1, int64(len(part1)), bytes.NewReader(part1))
+		part1Done <- err
+	}()
+	<-sink.blocked
+	_, err = b.UploadPart(ctx, bucket, object, uploadID, 1, int64(len(part1)), bytes.NewReader(part1))
+	require.NoError(t, err)
+	return part1Done
+}
+
+// TestMultipartCloseWaitsForPump checks that completing an upload whose
+// parts have all been acknowledged waits for parts still being written to
+// the sink, rather than rejecting the completion.
+func TestMultipartCloseWaitsForPump(t *testing.T) {
+	b, _, bucket := newPutTestBackend(t, "", nil)
+	const object = "close-while-pumping.bin"
+	uploadID, _ := newBufferSinkUpload(t, b, bucket, object)
+	up, err := b.loadUpload(context.Background(), uploadID)
+	require.NoError(t, err)
+	part1 := []byte(random.String(1024))
+	part2 := []byte(random.String(1024))
+	sink := newGatedSink(len(part1))
+	up.fh = sink
+	part1Done := startPumpBehindRetry(t, b, bucket, object, uploadID, sink, part1, part2)
+
+	closed := make(chan error, 1)
+	go func() { closed <- up.close(context.Background()) }()
+	select {
+	case err := <-closed:
+		t.Fatalf("close returned %v while part 2 was still being written", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(sink.release)
+	require.NoError(t, <-closed)
+	require.NoError(t, <-part1Done)
+	assert.Equal(t, append(part1, part2...), sink.Bytes())
+}
+
+// TestMultipartCompleteWaitingForPumpGivesUp checks that a completion which
+// waits too long for acknowledged parts to be written tells the client to
+// slow down and keeps the upload, so the client can retry it.
+func TestMultipartCompleteWaitingForPumpGivesUp(t *testing.T) {
+	b, _, bucket := newPutTestBackend(t, "", nil)
+	oldTimeout := multipartWaitTimeout
+	multipartWaitTimeout = 100 * time.Millisecond
+	t.Cleanup(func() { multipartWaitTimeout = oldTimeout })
+	ctx := context.Background()
+	const object = "complete-gives-up.bin"
+	uploadID, _ := newBufferSinkUpload(t, b, bucket, object)
+	up, err := b.loadUpload(ctx, uploadID)
+	require.NoError(t, err)
+	part1 := []byte(random.String(1024))
+	part2 := []byte(random.String(1024))
+	sink := newGatedSink(len(part1))
+	up.fh = sink
+	part1Done := startPumpBehindRetry(t, b, bucket, object, uploadID, sink, part1, part2)
+	md5s := [][16]byte{md5.Sum(part1), md5.Sum(part2)}
+
+	_, _, err = b.CompleteMultipartUpload(ctx, bucket, object, uploadID, &gofakes3.CompleteMultipartUploadRequest{
+		Parts: []gofakes3.CompletedPart{
+			{PartNumber: 1, ETag: hex.EncodeToString(md5s[0][:])},
+			{PartNumber: 2, ETag: hex.EncodeToString(md5s[1][:])},
+		},
+	})
+	assert.True(t, gofakes3.HasErrorCode(err, gofakes3.ErrSlowDown), "want SlowDown, got %v", err)
+
+	close(sink.release)
+	require.NoError(t, <-part1Done)
+	_, err = b.loadUpload(ctx, uploadID)
+	require.NoError(t, err, "the upload was torn down instead of being kept for a retry")
+	require.NoError(t, up.close(context.Background()))
+	assert.Equal(t, append(part1, part2...), sink.Bytes())
+}
+
+// TestMultipartCompleteWaitingForPumpCancelled checks that a completion
+// whose request is cancelled while acknowledged parts are being written
+// returns at once, rather than after multipartWaitTimeout, and keeps the
+// upload so the client can retry it.
+func TestMultipartCompleteWaitingForPumpCancelled(t *testing.T) {
+	b, _, bucket := newPutTestBackend(t, "", nil)
+	ctx := context.Background()
+	const object = "complete-cancelled.bin"
+	uploadID, _ := newBufferSinkUpload(t, b, bucket, object)
+	up, err := b.loadUpload(ctx, uploadID)
+	require.NoError(t, err)
+	part1 := []byte(random.String(1024))
+	part2 := []byte(random.String(1024))
+	sink := newGatedSink(len(part1))
+	up.fh = sink
+	part1Done := startPumpBehindRetry(t, b, bucket, object, uploadID, sink, part1, part2)
+	md5s := [][16]byte{md5.Sum(part1), md5.Sum(part2)}
+
+	completeCtx, cancel := context.WithCancel(ctx)
+	completed := make(chan error, 1)
+	go func() {
+		_, _, err := b.CompleteMultipartUpload(completeCtx, bucket, object, uploadID, &gofakes3.CompleteMultipartUploadRequest{
+			Parts: []gofakes3.CompletedPart{
+				{PartNumber: 1, ETag: hex.EncodeToString(md5s[0][:])},
+				{PartNumber: 2, ETag: hex.EncodeToString(md5s[1][:])},
+			},
+		})
+		completed <- err
+	}()
+	time.Sleep(50 * time.Millisecond)
+	cancel()
+	select {
+	case err := <-completed:
+		assert.ErrorIs(t, err, context.Canceled)
+	case <-time.After(5 * time.Second):
+		t.Fatal("the cancelled completion kept waiting for the pump")
+	}
+
+	close(sink.release)
+	require.NoError(t, <-part1Done)
+	_, err = b.loadUpload(ctx, uploadID)
+	require.NoError(t, err, "the upload was torn down instead of being kept for a retry")
+	require.NoError(t, up.close(ctx))
+	assert.Equal(t, append(part1, part2...), sink.Bytes())
+}
+
 // TestMultipartPumpFailure checks that an upload is failed when a part
 // which was buffered, and so already acknowledged to the client, can't be
 // written to the sink, rather than a resend of it being appended to what