diff --git a/backend/drive/upload.go b/backend/drive/upload.go index 10b3ab506..054024a7a 100644 --- a/backend/drive/upload.go +++ b/backend/drive/upload.go @@ -11,7 +11,6 @@ package drive import ( - "bytes" "context" "encoding/json" "fmt" @@ -22,6 +21,7 @@ import ( "github.com/rclone/rclone/fs" "github.com/rclone/rclone/fs/fserrors" + "github.com/rclone/rclone/lib/multipart" "github.com/rclone/rclone/lib/readers" "google.golang.org/api/drive/v3" "google.golang.org/api/googleapi" @@ -112,8 +112,13 @@ func (f *Fs) Upload(ctx context.Context, in io.Reader, size int64, contentType, } // Make an http.Request for the range passed in +// +// The body is wrapped in readers.NoCloser so the transport can't upgrade it +// to an io.Closer and close it after each attempt — the chunk buffer is +// pool-backed, so an early close would free its pages and a retry would then +// read a dead buffer. The upload loop owns the buffer's lifetime. func (rx *resumableUpload) makeRequest(ctx context.Context, start int64, body io.ReadSeeker, reqSize int64) *http.Request { - req, _ := http.NewRequestWithContext(ctx, "POST", rx.URI, body) + req, _ := http.NewRequestWithContext(ctx, "POST", rx.URI, readers.NoCloser(body)) req.ContentLength = reqSize totalSize := "*" if rx.ContentLength >= 0 { @@ -168,37 +173,40 @@ func (rx *resumableUpload) Upload(ctx context.Context) (*drive.File, error) { start := int64(0) var StatusCode int var err error - buf := make([]byte, int(rx.f.opt.ChunkSize)) + chunkSize := int64(rx.f.opt.ChunkSize) for finished := false; !finished; { - var reqSize int64 - var chunk io.ReadSeeker + reqSize := chunkSize if rx.ContentLength >= 0 { - // If size known use repeatable reader for smoother bwlimit if start >= rx.ContentLength { break } - reqSize = min(rx.ContentLength-start, int64(rx.f.opt.ChunkSize)) - chunk = readers.NewRepeatableLimitReaderBuffer(rx.Media, buf, reqSize) - } else { - // If size unknown read into buffer - var n int - n, err = readers.ReadFill(rx.Media, buf) + reqSize = min(rx.ContentLength-start, chunkSize) + } + + // Buffer the chunk in memory from the global pool so reads are + // repeatable for retries + rw := multipart.NewRW() + var n int64 + n, err = io.CopyN(rw, rx.Media, reqSize) + if rx.ContentLength < 0 { if err == io.EOF { // Send the last chunk with the correct ContentLength // otherwise Google doesn't know we've finished - rx.ContentLength = start + int64(n) + rx.ContentLength = start + n finished = true - } else if err != nil { - return nil, err + err = nil } - reqSize = int64(n) - chunk = bytes.NewReader(buf[:reqSize]) + reqSize = n + } + if err != nil { + _ = rw.Close() + return nil, err } // Transfer the chunk err = rx.f.pacer.Call(func() (bool, error) { fs.Debugf(rx.remote, "Sending chunk %d length %d", start, reqSize) - StatusCode, err = rx.transferChunk(ctx, start, chunk, reqSize) + StatusCode, err = rx.transferChunk(ctx, start, rw, reqSize) again, err := rx.f.shouldRetry(ctx, err) if StatusCode == statusResumeIncomplete || StatusCode == http.StatusCreated || StatusCode == http.StatusOK { again = false @@ -206,9 +214,13 @@ func (rx *resumableUpload) Upload(ctx context.Context) (*drive.File, error) { } return again, err }) + closeErr := rw.Close() if err != nil { return nil, err } + if closeErr != nil { + return nil, closeErr + } start += reqSize } diff --git a/backend/drive/upload_internal_test.go b/backend/drive/upload_internal_test.go new file mode 100644 index 000000000..9a90fe95d --- /dev/null +++ b/backend/drive/upload_internal_test.go @@ -0,0 +1,57 @@ +package drive + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/rclone/rclone/fs" + "github.com/rclone/rclone/lib/pacer" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestResumableUploadRetry checks that a chunk which fails with a 5xx is +// retried successfully. The chunk buffer is pool-backed and implements +// io.Closer, so if it reaches the http transport unwrapped the transport +// closes it after the failed attempt, returning its pages to the pool, and +// the retry then reads a freed buffer. +func TestResumableUploadRetry(t *testing.T) { + content := bytes.Repeat([]byte("resumable"), 512) + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + if attempts.Add(1) == 1 { + w.WriteHeader(http.StatusInternalServerError) + return + } + assert.Equal(t, content, body, "retried chunk should re-send the full chunk") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"fake-id","name":"remote"}`)) + })) + defer server.Close() + + f := &Fs{ + pacer: fs.NewPacer(context.Background(), pacer.NewGoogleDrive()), + client: server.Client(), + } + f.opt.ChunkSize = fs.SizeSuffix(len(content)) + rx := &resumableUpload{ + f: f, + remote: "remote", + URI: server.URL, + Media: bytes.NewReader(content), + MediaType: "application/octet-stream", + ContentLength: int64(len(content)), + } + info, err := rx.Upload(context.Background()) + require.NoError(t, err) + require.NotNil(t, info) + assert.Equal(t, "fake-id", info.Id) + assert.Equal(t, int32(2), attempts.Load(), "expected exactly one failed attempt and one retry") +}