service/s3/s3manager: Use sync.Pool for reuse of part buffers for streaming payloads

Author: skmcgrail

CHANGELOG_PENDING.md

@@ -29,6 +29,9 @@ SDK Enhancements
   * Related to [aws/aws-sdk-go#2310](https://github.com/aws/aws-sdk-go/pull/2310)
   * Fixes [#251](https://github.com/aws/aws-sdk-go-v2/issues/251)
 * `aws/request` : Retryer is now a named field on Request. ([#393](https://github.com/aws/aws-sdk-go-v2/pull/393))
-  
+* `service/s3/s3manager`: Adds `sync.Pool` to allow reuse of part buffers for streaming payloads ([#404](https://github.com/aws/aws-sdk-go-v2/pull/404))
+  * Fixes [#402](https://github.com/aws/aws-sdk-go-v2/issues/402) 
+  * Uses the new behavior introduced in V1 [#2863](https://github.com/aws/aws-sdk-go/pull/2863) which allows the reuse of the sync.Pool across multiple Upload request that match part sizes.
+   
 SDK Bugs
 ---

service/s3/s3manager/download.go

@@ -433,7 +433,7 @@ func (d *downloader) downloadChunk(chunk dlchunk) error {
 		// Check if the returned error is an errReadingBody.
 		// If err is errReadingBody this indicates that an error
 		// occurred while copying the http response body.
-		// If this occurs we unwrap the error to set the underling error
+		// If this occurs we unwrap the error to set the underlying error
 		// and attempt any remaining retries.
 		if bodyErr, ok := err.(*errReadingBody); ok {
 			err = bodyErr.Unwrap()

service/s3/s3manager/upload.go

@@ -162,6 +162,9 @@ type Uploader struct {
 
 	// Defines the buffer strategy used when uploading a part
 	BufferProvider ReadSeekerWriteToProvider
+
+	// partPool allows for the re-usage of streaming payload part buffers between upload calls
+	partPool *partPool
 }
 
 // NewUploader creates a new Uploader instance to upload objects to S3. Pass In
@@ -179,8 +182,12 @@ type Uploader struct {
 //          u.PartSize = 64 * 1024 * 1024 // 64MB per part
 //     })
 func NewUploader(cfg aws.Config, options ...func(*Uploader)) *Uploader {
+	return newUploader(s3.New(cfg), options...)
+}
+
+func newUploader(client s3iface.ClientAPI, options ...func(*Uploader)) *Uploader {
 	u := &Uploader{
-		S3:                s3.New(cfg),
+		S3:                client,
 		PartSize:          DefaultUploadPartSize,
 		Concurrency:       DefaultUploadConcurrency,
 		LeavePartsOnError: false,
@@ -192,6 +199,8 @@ func NewUploader(cfg aws.Config, options ...func(*Uploader)) *Uploader {
 		option(u)
 	}
 
+	u.partPool = newPartPool(u.PartSize)
+
 	return u
 }
 
@@ -214,20 +223,7 @@ func NewUploader(cfg aws.Config, options ...func(*Uploader)) *Uploader {
 //          u.PartSize = 64 * 1024 * 1024 // 64MB per part
 //     })
 func NewUploaderWithClient(svc s3iface.ClientAPI, options ...func(*Uploader)) *Uploader {
-	u := &Uploader{
-		S3:                svc,
-		PartSize:          DefaultUploadPartSize,
-		Concurrency:       DefaultUploadConcurrency,
-		LeavePartsOnError: false,
-		MaxUploadParts:    MaxUploadParts,
-		BufferProvider:    defaultUploadBufferProvider(),
-	}
-
-	for _, option := range options {
-		option(u)
-	}
-
-	return u
+	return newUploader(svc, options...)
 }
 
 // Upload uploads an object to S3, intelligently buffering large files into
@@ -287,6 +283,7 @@ func (u Uploader) UploadWithContext(ctx context.Context, input *UploadInput, opt
 	for _, opt := range opts {
 		opt(&i.cfg)
 	}
+
 	i.cfg.RequestOptions = append(i.cfg.RequestOptions, request.WithAppendUserAgent("S3Manager"))
 
 	return i.upload()
@@ -391,6 +388,15 @@ func (u *uploader) init() error {
 	if u.cfg.PartSize == 0 {
 		u.cfg.PartSize = DefaultUploadPartSize
 	}
+	if u.cfg.MaxUploadParts == 0 {
+		u.cfg.MaxUploadParts = MaxUploadParts
+	}
+
+	// If PartSize was changed or partPool was never setup then we need to allocated a new pool
+	// so that we return []byte slices of the correct size
+	if u.cfg.partPool == nil || u.cfg.partPool.partSize != u.cfg.PartSize {
+		u.cfg.partPool = newPartPool(u.cfg.PartSize)
+	}
 
 	// Try to get the total size for some optimizations
 	return u.initSize()
@@ -460,11 +466,13 @@ func (u *uploader) nextReader() (io.ReadSeeker, int, func(), error) {
 		return reader, int(n), cleanup, err
 
 	default:
-		part := make([]byte, u.cfg.PartSize)
+		part := u.cfg.partPool.Get().([]byte)
 		n, err := readFillBuf(r, part)
 		u.readerPos += int64(n)
 
-		cleanup := func() {}
+		cleanup := func() {
+			u.cfg.partPool.Put(part)
+		}
 
 		return bytes.NewReader(part[0:n]), n, cleanup, err
 	}
@@ -751,3 +759,18 @@ func (u *multiuploader) complete() *s3.CompleteMultipartUploadOutput {
 
 	return resp.CompleteMultipartUploadOutput
 }
+
+type partPool struct {
+	partSize int64
+	sync.Pool
+}
+
+func newPartPool(partSize int64) *partPool {
+	p := &partPool{partSize: partSize}
+
+	p.New = func() interface{} {
+		return make([]byte, p.partSize)
+	}
+
+	return p
+}