Transfer Manager v1 uploader optimization (#3284)

Author: wty-Bryant

.changelog/458251ad35f0459cb2d45d165a1d9420.json

@@ -0,0 +1,8 @@
+{
+    "id": "458251ad-35f0-459c-b2d4-5d165a1d9420",
+    "type": "feature",
+    "description": "Optimize allocation for transfer manager v1 uploader so no extra memory is used in buffer pool for single upload.",
+    "modules": [
+        "feature/s3/manager"
+    ]
+}

feature/s3/manager/upload.go

@@ -501,6 +501,19 @@ func (u *uploader) nextReader() (io.ReadSeeker, int, func(), error) {
 		return reader, int(n), cleanup, err
 
 	default:
+		if u.readerPos == 0 {
+			r := io.LimitReader(u.in.Body, u.cfg.PartSize)
+			firstPart, err := io.ReadAll(r)
+			if err != nil {
+				return nil, 0, func() {}, err
+			}
+			n := len(firstPart)
+			u.readerPos += int64(n)
+			if int64(n) < u.cfg.PartSize {
+				return bytes.NewReader(firstPart), n, func() {}, io.EOF
+			}
+			return bytes.NewReader(firstPart), n, func() {}, nil
+		}
 		part, err := u.cfg.partPool.Get(u.ctx)
 		if err != nil {
 			return nil, 0, func() {}, err

feature/s3/manager/upload_test.go

@@ -113,6 +113,160 @@ func TestUploadOrderMulti(t *testing.T) {
 	}
 }
 
+func TestUploadOrderMultiTriggerredBySinglePartSize(t *testing.T) {
+	c, invocations, args := s3testing.NewUploadLoggingClient(nil)
+	u := manager.NewUploader(c)
+
+	resp, err := u.Upload(context.Background(), &s3.PutObjectInput{
+		Bucket:               aws.String("Bucket"),
+		Key:                  aws.String("Key - value"),
+		Body:                 bytes.NewBuffer(make([]byte, 5*1024*1024)),
+		ServerSideEncryption: "aws:kms",
+		SSEKMSKeyId:          aws.String("KmsId"),
+		ContentType:          aws.String("content/type"),
+	})
+
+	if err != nil {
+		t.Errorf("Expected no error but received %v", err)
+	}
+
+	if diff := cmpDiff([]string{"CreateMultipartUpload", "UploadPart",
+		"CompleteMultipartUpload"}, *invocations); len(diff) > 0 {
+		t.Error(err)
+	}
+
+	if e, a := `https://mock.amazonaws.com/key`, resp.Location; e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+
+	if "UPLOAD-ID" != resp.UploadID {
+		t.Errorf("expect %q, got %q", "UPLOAD-ID", resp.UploadID)
+	}
+
+	if "VERSION-ID" != *resp.VersionID {
+		t.Errorf("expect %q, got %q", "VERSION-ID", *resp.VersionID)
+	}
+
+	// Validate input values
+	v := aws.ToString((*args)[1].(*s3.UploadPartInput).UploadId)
+	if "UPLOAD-ID" != v {
+		t.Errorf("Expected %q, but received %q", "UPLOAD-ID", v)
+	}
+	v = aws.ToString((*args)[2].(*s3.CompleteMultipartUploadInput).UploadId)
+	if "UPLOAD-ID" != v {
+		t.Errorf("Expected %q, but received %q", "UPLOAD-ID", v)
+	}
+
+	parts := (*args)[2].(*s3.CompleteMultipartUploadInput).MultipartUpload.Parts
+
+	num := parts[0].PartNumber
+	etag := aws.ToString(parts[0].ETag)
+
+	if aws.ToInt32(num) != 1 {
+		t.Errorf("expect 1, got %d", num)
+	}
+
+	if matched, err := regexp.MatchString(`^ETAG\d+$`, etag); !matched || err != nil {
+		t.Errorf("Failed regexp expression `^ETAG\\d+$` , got %s", etag)
+	}
+
+	// Custom headers
+	cmu := (*args)[0].(*s3.CreateMultipartUploadInput)
+
+	if e, a := types.ServerSideEncryption("aws:kms"), cmu.ServerSideEncryption; e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+
+	if e, a := "KmsId", aws.ToString(cmu.SSEKMSKeyId); e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+
+	if e, a := "content/type", aws.ToString(cmu.ContentType); e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+}
+
+func TestUploadOrderMultiJustExceedSinglePart(t *testing.T) {
+	c, invocations, args := s3testing.NewUploadLoggingClient(nil)
+	u := manager.NewUploader(c)
+
+	resp, err := u.Upload(context.Background(), &s3.PutObjectInput{
+		Bucket:               aws.String("Bucket"),
+		Key:                  aws.String("Key - value"),
+		Body:                 bytes.NewBuffer(make([]byte, 5*1024*1024+1)),
+		ServerSideEncryption: "aws:kms",
+		SSEKMSKeyId:          aws.String("KmsId"),
+		ContentType:          aws.String("content/type"),
+	})
+
+	if err != nil {
+		t.Errorf("Expected no error but received %v", err)
+	}
+
+	if diff := cmpDiff([]string{"CreateMultipartUpload", "UploadPart", "UploadPart",
+		"CompleteMultipartUpload"}, *invocations); len(diff) > 0 {
+		t.Error(err)
+	}
+
+	if e, a := `https://mock.amazonaws.com/key`, resp.Location; e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+
+	if "UPLOAD-ID" != resp.UploadID {
+		t.Errorf("expect %q, got %q", "UPLOAD-ID", resp.UploadID)
+	}
+
+	if "VERSION-ID" != *resp.VersionID {
+		t.Errorf("expect %q, got %q", "VERSION-ID", *resp.VersionID)
+	}
+
+	// Validate input values
+
+	// UploadPart
+	for i := 1; i < 3; i++ {
+		v := aws.ToString((*args)[i].(*s3.UploadPartInput).UploadId)
+		if "UPLOAD-ID" != v {
+			t.Errorf("Expected %q, but received %q", "UPLOAD-ID", v)
+		}
+	}
+
+	// CompleteMultipartUpload
+	v := aws.ToString((*args)[3].(*s3.CompleteMultipartUploadInput).UploadId)
+	if "UPLOAD-ID" != v {
+		t.Errorf("Expected %q, but received %q", "UPLOAD-ID", v)
+	}
+
+	parts := (*args)[3].(*s3.CompleteMultipartUploadInput).MultipartUpload.Parts
+
+	for i := 0; i < 2; i++ {
+		num := parts[i].PartNumber
+		etag := aws.ToString(parts[i].ETag)
+
+		if int32(i+1) != aws.ToInt32(num) {
+			t.Errorf("expect %d, got %d", i+1, num)
+		}
+
+		if matched, err := regexp.MatchString(`^ETAG\d+$`, etag); !matched || err != nil {
+			t.Errorf("Failed regexp expression `^ETAG\\d+$`")
+		}
+	}
+
+	// Custom headers
+	cmu := (*args)[0].(*s3.CreateMultipartUploadInput)
+
+	if e, a := types.ServerSideEncryption("aws:kms"), cmu.ServerSideEncryption; e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+
+	if e, a := "KmsId", aws.ToString(cmu.SSEKMSKeyId); e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+
+	if e, a := "content/type", aws.ToString(cmu.ContentType); e != a {
+		t.Errorf("expect %q, got %q", e, a)
+	}
+}
+
 func TestUploadOrderMultiDifferentPartSize(t *testing.T) {
 	s, ops, args := s3testing.NewUploadLoggingClient(nil)
 	mgr := manager.NewUploader(s, func(u *manager.Uploader) {
@@ -411,13 +565,13 @@ func TestUploadOrderMultiFailureLeaveParts(t *testing.T) {
 }
 
 type failreader struct {
-	times     int
-	failCount int
+	failBytes int64
+	readBytes int64
 }
 
 func (f *failreader) Read(b []byte) (int, error) {
-	f.failCount++
-	if f.failCount >= f.times {
+	f.readBytes += int64(len(b))
+	if f.readBytes > f.failBytes {
 		return 0, fmt.Errorf("random failure")
 	}
 	return len(b), nil
@@ -429,7 +583,7 @@ func TestUploadOrderReadFail1(t *testing.T) {
 	_, err := mgr.Upload(context.Background(), &s3.PutObjectInput{
 		Bucket: aws.String("Bucket"),
 		Key:    aws.String("Key"),
-		Body:   &failreader{times: 1},
+		Body:   &failreader{failBytes: 1},
 	})
 	if err == nil {
 		t.Fatalf("expect error to not be nil")
@@ -452,7 +606,7 @@ func TestUploadOrderReadFail2(t *testing.T) {
 	_, err := mgr.Upload(context.Background(), &s3.PutObjectInput{
 		Bucket: aws.String("Bucket"),
 		Key:    aws.String("Key"),
-		Body:   &failreader{times: 2},
+		Body:   &failreader{failBytes: 5 * 1024 * 1024},
 	})
 	if err == nil {
 		t.Fatalf("expect error to not be nil")