From c8dcb507b7f4bdefa010d8c31f7e14af7dd2b48c Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Thu, 2 Jul 2026 18:21:19 -0700 Subject: [PATCH] Fix seekable transfermanager uploads --- .../s3/transfermanager/api_op_UploadObject.go | 52 ++++++++++++++++++- .../api_op_UploadObject_test.go | 42 +++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/feature/s3/transfermanager/api_op_UploadObject.go b/feature/s3/transfermanager/api_op_UploadObject.go index 23c0eb2db78..6f8cbefffeb 100644 --- a/feature/s3/transfermanager/api_op_UploadObject.go +++ b/feature/s3/transfermanager/api_op_UploadObject.go @@ -828,10 +828,17 @@ type uploader struct { partPool bytesBufferPool objectSize int64 multipleRead bool + readerStart int64 + readerPos int64 progressEmitter *singleObjectProgressEmitter } +type readerAtSeeker interface { + io.ReaderAt + io.Seeker +} + func (u *uploader) upload(ctx context.Context) (*UploadObjectOutput, error) { if err := u.init(); err != nil { return nil, fmt.Errorf("unable to initialize upload: %w", err) @@ -854,8 +861,10 @@ func (u *uploader) upload(ctx context.Context) (*UploadObjectOutput, error) { return nil, err } - u.partPool = newDefaultSlicePool(u.options.PartSizeBytes, u.options.Concurrency+1) // only create the caching pool for multipart upload - defer u.partPool.Close() + if _, ok := u.in.Body.(readerAtSeeker); !ok { + u.partPool = newDefaultSlicePool(u.options.PartSizeBytes, u.options.Concurrency+1) // only create the caching pool for multipart upload + defer u.partPool.Close() + } mu := multiUploader{ uploader: u, } @@ -878,6 +887,13 @@ func (u *uploader) initSize() error { u.objectSize = -1 switch r := u.in.Body.(type) { case io.Seeker: + start, err := r.Seek(0, io.SeekCurrent) + if err != nil { + return err + } + u.readerStart = start + u.readerPos = start + n, err := types.SeekerLen(r) if err != nil { return err @@ -931,6 +947,10 @@ func (u *uploader) singleUpload(ctx context.Context, r io.Reader, sz int, cleanU // nextReader reads the next chunk of data from input Body func (u *uploader) nextReader(ctx context.Context) (io.Reader, int, func(), error) { + if r, ok := u.in.Body.(readerAtSeeker); ok { + return u.nextSectionReader(r) + } + if !u.multipleRead { u.multipleRead = true // read first part up to a maximum of PartSize to avoid allocating 8MB buffer out of the gate @@ -969,6 +989,34 @@ func (u *uploader) nextReader(ctx context.Context) (io.Reader, int, func(), erro return bytes.NewReader(part[0:n]), n, cleanup, err } +func (u *uploader) nextSectionReader(r readerAtSeeker) (io.Reader, int, func(), error) { + if !u.multipleRead { + u.multipleRead = true + if u.objectSize < u.options.MultipartUploadThreshold { + n := u.objectSize + reader := io.NewSectionReader(r, u.readerPos, n) + u.readerPos += n + return reader, int(n), func() {}, io.EOF + } + } + + bytesLeft := u.objectSize - (u.readerPos - u.readerStart) + if bytesLeft <= 0 { + return bytes.NewReader(nil), 0, func() {}, io.EOF + } + + n := u.options.PartSizeBytes + var err error + if bytesLeft <= n { + n = bytesLeft + err = io.EOF + } + + reader := io.NewSectionReader(r, u.readerPos, n) + u.readerPos += n + return reader, int(n), func() {}, err +} + func (u *uploader) freshContext(ctx context.Context) (context.Context, context.CancelFunc) { if u.options.FailTimeout <= 0 { return ctx, func() {} diff --git a/feature/s3/transfermanager/api_op_UploadObject_test.go b/feature/s3/transfermanager/api_op_UploadObject_test.go index 3733072fbb7..77d55a28777 100644 --- a/feature/s3/transfermanager/api_op_UploadObject_test.go +++ b/feature/s3/transfermanager/api_op_UploadObject_test.go @@ -284,6 +284,48 @@ func TestUploadOrderMultiDifferentPartSize(t *testing.T) { } } +func TestUploadOrderMultiSeekableBodyUsesSectionReaders(t *testing.T) { + c, ops, args := s3testing.NewUploadLoggingClient(nil) + mgr := New(c, func(options *Options) { + options.PartSizeBytes = 4 + options.MultipartUploadThreshold = 6 + options.Concurrency = 1 + }) + + body := bytes.NewReader([]byte("0123456789abc")) + if _, err := body.Seek(2, io.SeekStart); err != nil { + t.Fatalf("failed to seek body: %v", err) + } + + _, err := mgr.UploadObject(context.Background(), &UploadObjectInput{ + Bucket: aws.String("Bucket"), + Key: aws.String("Key"), + Body: body, + }) + if err != nil { + t.Errorf("expect no error, got %v", err) + } + + if diff := cmpDiff([]string{"CreateMultipartUpload", "UploadPart", "UploadPart", "UploadPart", "CompleteMultipartUpload"}, *ops); len(diff) > 0 { + t.Error(diff) + } + + for i, expect := range []string{"2345", "6789", "abc"} { + part := (*args)[i+1].(*s3.UploadPartInput) + if _, ok := part.Body.(*io.SectionReader); !ok { + t.Errorf("expect part %d body to be %T, got %T", i+1, &io.SectionReader{}, part.Body) + } + + actual, err := io.ReadAll(part.Body) + if err != nil { + t.Fatalf("failed to read part %d body: %v", i+1, err) + } + if string(actual) != expect { + t.Errorf("expect part %d body %q, got %q", i+1, expect, string(actual)) + } + } +} + func TestUploadOrderMultiWithPartSizeEqualToThreshold(t *testing.T) { c, ops, args := s3testing.NewUploadLoggingClient(nil) mgr := New(c, func(options *Options) {