filer: retry small non-seekable uploads to UC Volumes

.nextchanges/cli/7607.md

@@ -0,0 +1 @@
+* Retry `fs cp` uploads to UC Volumes from another Volume or from DBFS (files up to 8 MiB) when the Files API returns a transient error such as HTTP 503. ([#7607](https://github.com/databricks/cli/pull/7607))

libs/filer/files_client.go

@@ -1,9 +1,11 @@
 package filer
 
 import (
+	"bytes"
 	"cmp"
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"io/fs"
 	"net/http"
@@ -282,7 +284,8 @@ func (w *FilesClient) Write(ctx context.Context, name string, reader io.Reader,
 	// upload engine, which sends small files in a single PUT and splits large ones
 	// into parts. Non-seekable streams and the flag-off case use the single-shot
 	// UploadFile endpoint below: the former to avoid buffering the whole stream in
-	// memory, the latter to keep the default behavior a plain PUT. The engine
+	// memory (uploadFile buffers only small ones, to retry them), the latter to
+	// keep the default behavior a plain PUT. The engine
 	// recovers an io.ReaderAt for concurrent positioned reads when the source
 	// provides one (a local file).
 	if MultipartUploadEnabled(ctx) && isSeekable(reader) {
@@ -312,24 +315,50 @@ func (w *FilesClient) Write(ctx context.Context, name string, reader io.Reader,
 	return err
 }
 
+// filesReplayBufferSize is the largest non-seekable upload that is buffered in
+// memory so it can be retried. fs cp between remote paths (Volumes or DBFS to
+// Volumes) streams the source download into the upload, so without the buffer
+// a single 503 fails the copy. Larger streams get one attempt, so memory per
+// upload stays bounded. A variable so tests can lower it.
+var filesReplayBufferSize = 8 << 20
+
+// replayable returns a reader that can be rewound to start before each upload
+// attempt. A non-seekable reader (including a pipe, which is an *os.File whose
+// Seek fails) is read into memory if it is at most filesReplayBufferSize bytes;
+// otherwise stream holds the bytes read so far followed by the
+// rest of reader, to be sent once.
+func replayable(reader io.Reader) (rs io.ReadSeeker, start int64, stream io.Reader, err error) {
+	if s, ok := reader.(io.ReadSeeker); ok {
+		if start, err := s.Seek(0, io.SeekCurrent); err == nil {
+			return s, start, nil, nil
+		}
+	}
+	head, err := io.ReadAll(io.LimitReader(reader, int64(filesReplayBufferSize)+1))
+	if err != nil {
+		return nil, 0, nil, fmt.Errorf("reading upload body: %w", err)
+	}
+	if len(head) > filesReplayBufferSize {
+		return nil, 0, io.MultiReader(bytes.NewReader(head), reader), nil
+	}
+	return bytes.NewReader(head), 0, nil, nil
+}
+
 // uploadFile sends reader in a single PUT. files/v2 rejects a retrier on this
-// call because it cannot replay a streamed body, so a seekable reader is rewound
-// to its starting offset and retried here. Other readers get one attempt; that
-// includes pipes, which are *os.File and so implement io.Seeker but fail Seek.
+// call because it cannot replay a streamed body, so the filer retries it here
+// with a reader from replayable, rewound to its starting offset each attempt.
 func (w *FilesClient) uploadFile(ctx context.Context, absPath string, reader io.Reader, overwrite bool) error {
-	seeker, ok := reader.(io.Seeker)
-	if !ok {
-		return w.putFile(ctx, absPath, io.NopCloser(reader), overwrite)
-	}
-	start, err := seeker.Seek(0, io.SeekCurrent)
+	seeker, start, stream, err := replayable(reader)
 	if err != nil {
-		return w.putFile(ctx, absPath, io.NopCloser(reader), overwrite)
+		return err
+	}
+	if stream != nil {
+		return w.putFile(ctx, absPath, io.NopCloser(stream), overwrite)
 	}
 	return ops.Execute(ctx, func(ctx context.Context) error {
 		if _, err := seeker.Seek(start, io.SeekStart); err != nil {
 			return err
 		}
-		body := &attemptBody{r: reader}
+		body := &attemptBody{r: seeker}
 		err := w.putFile(ctx, absPath, body, overwrite)
 		body.Close()
 		return err

libs/filer/files_client_test.go

@@ -312,14 +312,63 @@ func TestFilesClientRetriesThrottledUploadFromStartOffset(t *testing.T) {
 	assert.Equal(t, []string{"abc", "abc", "abc"}, uploads)
 }
 
-func TestFilesClientDoesNotRetryNonSeekableUpload(t *testing.T) {
+// uploadsWithOneThrottle records each PUT body; the first PUT is throttled.
+func uploadsWithOneThrottle(server *testserver.Server) *[]string {
+	var uploads []string
+	put := throttled(1, ok)
+	server.Handle("PUT", "/api/2.0/fs/files/{path...}", func(req testserver.Request) any {
+		uploads = append(uploads, string(req.Body))
+		return put(req)
+	})
+	return &uploads
+}
+
+func TestFilesClientRetriesSmallNonSeekableUpload(t *testing.T) {
+	ctx := t.Context()
+	server := testserver.New(t)
+	uploads := uploadsWithOneThrottle(server)
+	f := newFilesClientWithFastRetries(t, server)
+	setReplayBufferSize(t, 3)
+
+	// Exactly filesReplayBufferSize bytes still fit in the buffer.
+	require.NoError(t, f.Write(ctx, "hello", onlyReader{strings.NewReader("abc")}, CreateParentDirectories))
+	assert.Equal(t, []string{"abc", "abc"}, *uploads)
+}
+
+func TestFilesClientRetriesSmallPipeUpload(t *testing.T) {
+	ctx := t.Context()
+	server := testserver.New(t)
+	uploads := uploadsWithOneThrottle(server)
+	f := newFilesClientWithFastRetries(t, server)
+
+	// A pipe is an *os.File whose Seek fails.
+	pr, pw, err := os.Pipe()
+	require.NoError(t, err)
+	_, err = pw.WriteString("abc")
+	require.NoError(t, err)
+	require.NoError(t, pw.Close())
+	defer pr.Close()
+
+	require.NoError(t, f.Write(ctx, "hello", pr, CreateParentDirectories))
+	assert.Equal(t, []string{"abc", "abc"}, *uploads)
+}
+
+func TestFilesClientDoesNotRetryLargeNonSeekableUpload(t *testing.T) {
 	ctx := t.Context()
 	server := testserver.New(t)
-	server.Handle("PUT", "/api/2.0/fs/files/{path...}", throttled(1, ok))
+	uploads := uploadsWithOneThrottle(server)
 	f := newFilesClientWithFastRetries(t, server)
+	setReplayBufferSize(t, 2)
 
 	err := f.Write(ctx, "hello", onlyReader{strings.NewReader("abc")}, CreateParentDirectories)
 	assert.ErrorContains(t, err, "AWS S3 is throttling requests")
+	assert.Equal(t, []string{"abc"}, *uploads)
+}
+
+func setReplayBufferSize(t *testing.T, n int) {
+	old := filesReplayBufferSize
+	filesReplayBufferSize = n
+	t.Cleanup(func() { filesReplayBufferSize = old })
 }
 
 func TestFilesClientStopsRetryingAfterLimit(t *testing.T) {