Repository: aws/aws-sdk-go (go)

Issue:
memory usage of s3 upload manager for Go 1.13.5 and 1.13.6
### Version of AWS SDK for Go?
v1.28.0 

### Version of Go (`go version`)?
1.13.5, 1.13.6

### What issue did you see?
Excessive memory consumption when concurrently uploading to s3 using the s3 upload manager. The program runs out of memory and crashes when uploading from an AWS EC2 linux instance with 1 GB RAM. The program does not crash or leak memory when running with Go 1.12.14.

Profiling Go 1.12.14 on instance: The top nodes don't contain any memory usage related to aws-sdk-go

Profiling Go 1.13.6 on instance: Program crashes (as it runs out of memory) and unable to see profiling output, however running the program locally (16 GB RAM) I can see all the memory is consumed by the s3 upload processes unlike the above.

(pprof) top
Showing nodes accounting for 266243.44kB, 99.40% of 267842.77kB total
Dropped 501 nodes (cum <= 1339.21kB)
Showing top 10 nodes out of 11
      flat  flat%   sum%        cum   cum%
266240.12kB 99.40% 99.40% 266240.12kB 99.40%  github.com/aws/aws-sdk-go/service/s3/s3manager.newPartPool.func1
       1kB 0.00037% 99.40% 235524.98kB 87.93%  github.com/aws/aws-sdk-go/service/s3/s3manager.(*multiuploader).upload
    0.59kB 0.00022% 99.40% 266246.66kB 99.40%  github.com/aws/aws-sdk-go/service/s3/s3manager.Uploader.UploadWithContext
    0.38kB 0.00014% 99.40% 266245.95kB 99.40%  github.com/aws/aws-sdk-go/service/s3/s3manager.(*uploader).upload
    0.09kB 3.5e-05% 99.40% 266241.81kB 99.40%  github.com/aws/aws-sdk-go/service/s3/s3manager.(*uploader).nextReader
         0     0% 99.40% 266241.62kB 99.40%  github.com/aws/aws-sdk-go/service/s3/s3manager.(*partPool).Get
         0     0% 99.40% 266246.66kB 99.40%  github.com/aws/aws-sdk-go/service/s3/s3manager.Uploader.Upload
 
### Steps to reproduce
Run the code below using Go 1.13.6 from an AWS EC2 instance (I did this with 10 uploads each around 200 MB in size)

```go
package main

import (
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

	_ "net/http/pprof"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/pkg/profile"
)

func (s s3uploader) uploadToS3(path string, res io.Reader) error {
	_, err := s.u.Upload(&s3manager.UploadInput{
		Bucket: aws.String(os.Getenv("S3_STORAGE_BUCKET")),
		Key:    aws.String(path),
		Body:   res,
	})
	if err != nil {
		return err
	}
	return nil
}

func (s s3uploader) syncOne(source string) error {
	u, err := url.Parse(source)
	if err != nil {
		return err
	}
	var res *http.Response
	res, err = http.Get(source)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	log.Printf("uploading: %s", source)
	if err := s.uploadToS3(u.Path, res.Body); err != nil {
		return err
	}
	return nil
}

func (s s3uploader) sync(urls []string, concurrency int) error {
	work := make(chan string)
	var wg sync.WaitGroup
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for url := range work {
				if err := s.syncOne(url); err != nil {
					time.Sleep(time.Millisecond*time.Duration(rand.Intn(5000)) + time.Second)
					log.Printf("Retrying url: %s", url)
				}
			}
		}()
	}

	for _, url := range urls {
		work <- url
	}
	close(work)
	wg.Wait()
	return nil
}

type s3uploader struct {
	u *s3manager.Uploader
}

func main() {
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()
	defer profile.Start(profile.MemProfile).Stop()
	conf := aws.Config{Region: aws.String(os.Getenv("AWS_REGION"))}
	s := session.New(&conf)
	uploader := s3manager.NewUploader(s, func(u *s3manager.Uploader) {
		u.PartSize = 10 * 1024 * 1024 // 10MB per part
	})
	var ul s3uploader
	ul.u = uploader

	const concurrency = 5
	if err := ul.sync(os.Args[1:], concurrency); err != nil {
		log.Printf("  ...syncing failed: %s", err)
		return
	}

}
```



   


Pull request:
Improve memory allocation behavior by replacing sync.Pool with custom pool implementation
* Fixes #3075

# Pool Benchmarks

```
name                  time/op
Pools/sync.Pool/0-72  23.3ms ± 6%
Pools/sync.Pool/1-72  44.1ms ± 3%
Pools/sync.Pool/2-72  51.6ms ±20%
Pools/sync.Pool/3-72   128ms ± 1%
Pools/custom/0-72     23.0ms ± 5%
Pools/custom/1-72     44.5ms ± 1%
Pools/custom/2-72     47.1ms ± 2%
Pools/custom/3-72      131ms ± 7%

name                  alloc/op
Pools/sync.Pool/0-72  10.6MB ± 0%
Pools/sync.Pool/1-72  15.9MB ± 0%
Pools/sync.Pool/2-72  26.6MB ± 0%
Pools/sync.Pool/3-72   164MB ± 0%
Pools/custom/0-72     10.6MB ± 0%
Pools/custom/1-72     10.7MB ± 0%
Pools/custom/2-72     16.0MB ± 0%
Pools/custom/3-72     59.0MB ± 0%

name                  allocs/op
Pools/sync.Pool/0-72   1.03k ± 0%
Pools/sync.Pool/1-72   1.29k ± 0%
Pools/sync.Pool/2-72   1.82k ± 0%
Pools/sync.Pool/3-72   13.4k ± 0%
Pools/custom/0-72      1.03k ± 0%
Pools/custom/1-72      1.28k ± 0%
Pools/custom/2-72      1.80k ± 0%
Pools/custom/3-72      13.4k ± 0%
```

# Multi-Part Upload
* 150 GB File
* 16 MB PartSize
* 72 Concurrency

Expected Buffer Allocations: (72+1)*16 = 1168 MB

## Before (sync.Pool)
As shown sync.Pool drastically performs over allocations, 2624 MB. This is over 124% more in allocations!

![Old Pool](https://user-images.githubusercontent.com/549813/76017558-a6eece00-5ed3-11ea-9af8-1c07d1f89d25.png)

## After
We see that we now allocate the exact amount of expected buffers, 1168 MB.

![New Pool](https://user-images.githubusercontent.com/549813/76018127-9723b980-5ed4-11ea-8e92-c547ba6e87d9.png)

# Performance Comparison
```
$ ~/go/bin/benchstat /mnt/ram/results.old /mnt/ram/results.new
name                                                         old time/op    new time/op    delta
Upload/540MB_File/72_Concurrency/5MB_PartSize/Unbuffered-72     1.50s ±27%     1.57s ±31%   +4.52%  (p=0.006 n=84+89)

name                                                         old alloc/op   new alloc/op   delta
Upload/540MB_File/72_Concurrency/5MB_PartSize/Unbuffered-72     595MB ± 0%     407MB ± 0%  -31.57%  (p=0.000 n=96+99)

name                                                         old allocs/op  new allocs/op  delta
Upload/540MB_File/72_Concurrency/5MB_PartSize/Unbuffered-72      151k ± 0%      151k ± 0%   +0.06%  (p=0.009 n=86+88)
```



Patch:
diff --git a/CHANGELOG_PENDING.md b/CHANGELOG_PENDING.md
index de4535194..f74e69e45 100644
--- a/CHANGELOG_PENDING.md
+++ b/CHANGELOG_PENDING.md
@@ -2,6 +2,9 @@
 
 ### SDK Enhancements
 * `aws/credentials`: Clarify `token` usage in `NewStaticCredentials` documentation.
-    * Related to [#3162](https://github.com/aws/aws-sdk-go/issues/3162).
+  * Related to [#3162](https://github.com/aws/aws-sdk-go/issues/3162).
+* `service/s3/s3manager`: Improve memory allocation behavior by replacing sync.Pool with custom pool implementation ([#3183](https://github.com/aws/aws-sdk-go/pull/3183))
+  * Improves memory allocations that occur when the provided `io.Reader` to upload does not satisfy both the `io.ReaderAt` and `io.ReadSeeker` interfaces.
+  * Fixes [#3075](https://github.com/aws/aws-sdk-go/issues/3075)
 
 ### SDK Bugs
diff --git a/service/s3/s3manager/pool.go b/service/s3/s3manager/pool.go
new file mode 100644
index 000000000..05113286d
--- /dev/null
+++ b/service/s3/s3manager/pool.go
@@ -0,0 +1,244 @@
+package s3manager
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
+
+type byteSlicePool interface {
+	Get(aws.Context) (*[]byte, error)
+	Put(*[]byte)
+	ModifyCapacity(int)
+	SliceSize() int64
+	Close()
+}
+
+type maxSlicePool struct {
+	// allocator is defined as a function pointer to allow
+	// for test cases to instrument custom tracers when allocations
+	// occur.
+	allocator sliceAllocator
+
+	slices         chan *[]byte
+	allocations    chan struct{}
+	capacityChange chan struct{}
+
+	max       int
+	sliceSize int64
+
+	mtx sync.RWMutex
+}
+
+func newMaxSlicePool(sliceSize int64) *maxSlicePool {
+	p := &maxSlicePool{sliceSize: sliceSize}
+	p.allocator = p.newSlice
+
+	return p
+}
+
+var errZeroCapacity = fmt.Errorf("get called on zero capacity pool")
+
+func (p *maxSlicePool) Get(ctx aws.Context) (*[]byte, error) {
+	// check if context is canceled before attempting to get a slice
+	// this ensures priority is given to the cancel case first
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
+	p.mtx.RLock()
+
+	for {
+		select {
+		case bs, ok := <-p.slices:
+			p.mtx.RUnlock()
+			if !ok {
+				// attempt to get on a zero capacity pool
+				return nil, errZeroCapacity
+			}
+			return bs, nil
+		case _, ok := <-p.allocations:
+			p.mtx.RUnlock()
+			if !ok {
+				// attempt to get on a zero capacity pool
+				return nil, errZeroCapacity
+			}
+			return p.allocator(), nil
+		case <-ctx.Done():
+			p.mtx.RUnlock()
+			return nil, ctx.Err()
+		default:
+			// In the event that there are no slices or allocations available
+			// This prevents some deadlock situations that can occur around sync.RWMutex
+			// When a lock request occurs on ModifyCapacity, no new readers are allowed to acquire a read lock.
+			// By releasing the read lock here and waiting for a notification, we prevent a deadlock situation where
+			// Get could hold the read lock indefinitely waiting for capacity, ModifyCapacity is waiting for a write lock,
+			// and a Put is blocked trying to get a read-lock which is blocked by ModifyCapacity.
+
+			// Short-circuit if the pool capacity is zero.
+			if p.max == 0 {
+				p.mtx.RUnlock()
+				return nil, errZeroCapacity
+			}
+
+			// Since we will be releasing the read-lock we need to take the reference to the channel.
+			// Since channels are references we will still get notified if slices are added, or if
+			// the channel is closed due to a capacity modification. This specifically avoids a data race condition
+			// where ModifyCapacity both closes a channel and initializes a new one while we don't have a read-lock.
+			c := p.capacityChange
+
+			p.mtx.RUnlock()
+
+			select {
+			case _ = <-c:
+				p.mtx.RLock()
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+		}
+	}
+}
+
+func (p *maxSlicePool) Put(bs *[]byte) {
+	p.mtx.RLock()
+	defer p.mtx.RUnlock()
+
+	if p.max == 0 {
+		return
+	}
+
+	select {
+	case p.slices <- bs:
+		p.notifyCapacity()
+	default:
+		// If the new channel when attempting to add the slice then we drop the slice.
+		// The logic here is to prevent a deadlock situation if channel is already at max capacity.
+		// Allows us to reap allocations that are returned and are no longer needed.
+	}
+}
+
+func (p *maxSlicePool) ModifyCapacity(delta int) {
+	if delta == 0 {
+		return
+	}
+
+	p.mtx.Lock()
+	defer p.mtx.Unlock()
+
+	p.max += delta
+
+	if p.max == 0 {
+		p.empty()
+		return
+	}
+
+	if p.capacityChange != nil {
+		close(p.capacityChange)
+	}
+	p.capacityChange = make(chan struct{}, p.max)
+
+	origAllocations := p.allocations
+	p.allocations = make(chan struct{}, p.max)
+
+	newAllocs := len(origAllocations) + delta
+	for i := 0; i < newAllocs; i++ {
+		p.allocations <- struct{}{}
+	}
+
+	if origAllocations != nil {
+		close(origAllocations)
+	}
+
+	origSlices := p.slices
+	p.slices = make(chan *[]byte, p.max)
+	if origSlices == nil {
+		return
+	}
+
+	close(origSlices)
+	for bs := range origSlices {
+		select {
+		case p.slices <- bs:
+		default:
+			// If the new channel blocks while adding slices from the old channel
+			// then we drop the slice. The logic here is to prevent a deadlock situation
+			// if the new channel has a smaller capacity then the old.
+		}
+	}
+}
+
+func (p *maxSlicePool) notifyCapacity() {
+	select {
+	case p.capacityChange <- struct{}{}:
+	default:
+		// This *shouldn't* happen as the channel is both buffered to the max pool capacity size and is resized
+		// on capacity modifications. This is just a safety to ensure that a blocking situation can't occur.
+	}
+}
+
+func (p *maxSlicePool) SliceSize() int64 {
+	return p.sliceSize
+}
+
+func (p *maxSlicePool) Close() {
+	p.mtx.Lock()
+	defer p.mtx.Unlock()
+	p.empty()
+}
+
+func (p *maxSlicePool) empty() {
+	p.max = 0
+
+	if p.capacityChange != nil {
+		close(p.capacityChange)
+		p.capacityChange = nil
+	}
+
+	if p.allocations != nil {
+		close(p.allocations)
+		for range p.allocations {
+			// drain channel
+		}
+		p.allocations = nil
+	}
+
+	if p.slices != nil {
+		close(p.slices)
+		for range p.slices {
+			// drain channel
+		}
+		p.slices = nil
+	}
+}
+
+func (p *maxSlicePool) newSlice() *[]byte {
+	bs := make([]byte, p.sliceSize)
+	return &bs
+}
+
+type returnCapacityPoolCloser struct {
+	byteSlicePool
+	returnCapacity int
+}
+
+func (n *returnCapacityPoolCloser) ModifyCapacity(delta int) {
+	if delta > 0 {
+		n.returnCapacity = -1 * delta
+	}
+	n.byteSlicePool.ModifyCapacity(delta)
+}
+
+func (n *returnCapacityPoolCloser) Close() {
+	if n.returnCapacity < 0 {
+		n.byteSlicePool.ModifyCapacity(n.returnCapacity)
+	}
+}
+
+type sliceAllocator func() *[]byte
+
+var newByteSlicePool = func(sliceSize int64) byteSlicePool {
+	return newMaxSlicePool(sliceSize)
+}
diff --git a/service/s3/s3manager/upload.go b/service/s3/s3manager/upload.go
index 5b6199fb4..8ed2975d8 100644
--- a/service/s3/s3manager/upload.go
+++ b/service/s3/s3manager/upload.go
@@ -366,6 +366,7 @@ func (u *uploader) upload() (*UploadOutput, error) {
 	if err := u.init(); err != nil {
 		return nil, awserr.New("ReadRequestBody", "unable to initialize upload", err)
 	}
+	defer u.cfg.partPool.Close()
 
 	if u.cfg.PartSize < MinUploadPartSize {
 		msg := fmt.Sprintf("part size must be at least %d bytes", MinUploadPartSize)
@@ -404,8 +405,13 @@ func (u *uploader) init() error {
 
 	// 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.Size() != u.cfg.PartSize {
+	poolCap := u.cfg.Concurrency + 1
+	if u.cfg.partPool == nil || u.cfg.partPool.SliceSize() != u.cfg.PartSize {
 		u.cfg.partPool = newByteSlicePool(u.cfg.PartSize)
+		u.cfg.partPool.ModifyCapacity(poolCap)
+	} else {
+		u.cfg.partPool = &returnCapacityPoolCloser{byteSlicePool: u.cfg.partPool}
+		u.cfg.partPool.ModifyCapacity(poolCap)
 	}
 
 	return nil
@@ -441,10 +447,6 @@ func (u *uploader) initSize() error {
 // does not need to be wrapped in a mutex because nextReader is only called
 // from the main thread.
 func (u *uploader) nextReader() (io.ReadSeeker, int, func(), error) {
-	type readerAtSeeker interface {
-		io.ReaderAt
-		io.ReadSeeker
-	}
 	switch r := u.in.Body.(type) {
 	case readerAtSeeker:
 		var err error
@@ -476,7 +478,11 @@ func (u *uploader) nextReader() (io.ReadSeeker, int, func(), error) {
 		return reader, int(n), cleanup, err
 
 	default:
-		part := u.cfg.partPool.Get()
+		part, err := u.cfg.partPool.Get(u.ctx)
+		if err != nil {
+			return nil, 0, func() {}, err
+		}
+
 		n, err := readFillBuf(r, *part)
 		u.readerPos += int64(n)
 
@@ -764,40 +770,7 @@ func (u *multiuploader) complete() *s3.CompleteMultipartUploadOutput {
 	return resp
 }
 
-type byteSlicePool interface {
-	Get() *[]byte
-	Put(*[]byte)
-	Size() int64
-}
-
-type partPool struct {
-	partSize int64
-	sync.Pool
-}
-
-func (p *partPool) Get() *[]byte {
-	return p.Pool.Get().(*[]byte)
-}
-
-func (p *partPool) Put(b *[]byte) {
-	p.Pool.Put(b)
-}
-
-func (p *partPool) Size() int64 {
-	return p.partSize
-}
-
-func newPartPool(partSize int64) *partPool {
-	p := &partPool{partSize: partSize}
-
-	p.New = func() interface{} {
-		bs := make([]byte, p.partSize)
-		return &bs
-	}
-
-	return p
-}
-
-var newByteSlicePool = func(partSize int64) byteSlicePool {
-	return newPartPool(partSize)
+type readerAtSeeker interface {
+	io.ReaderAt
+	io.ReadSeeker
 }


Test patch:
diff --git a/service/s3/s3manager/pool_test.go b/service/s3/s3manager/pool_test.go
new file mode 100644
index 000000000..55de684f6
--- /dev/null
+++ b/service/s3/s3manager/pool_test.go
@@ -0,0 +1,174 @@
+// +build go1.7
+
+package s3manager
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
+
+func TestMaxSlicePool(t *testing.T) {
+	pool := newMaxSlicePool(0)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			// increase pool capacity by 2
+			pool.ModifyCapacity(2)
+
+			// remove 2 items
+			bsOne, err := pool.Get(context.Background())
+			if err != nil {
+				t.Errorf("failed to get slice from pool: %v", err)
+			}
+			bsTwo, err := pool.Get(context.Background())
+			if err != nil {
+				t.Errorf("failed to get slice from pool: %v", err)
+			}
+
+			done := make(chan struct{})
+			go func() {
+				defer close(done)
+
+				// attempt to remove a 3rd in parallel
+				bs, err := pool.Get(context.Background())
+				if err != nil {
+					t.Errorf("failed to get slice from pool: %v", err)
+				}
+				pool.Put(bs)
+
+				// attempt to remove a 4th that has been canceled
+				ctx, cancel := context.WithCancel(context.Background())
+				cancel()
+				bs, err = pool.Get(ctx)
+				if err == nil {
+					pool.Put(bs)
+					t.Errorf("expected no slice to be returned")
+					return
+				}
+			}()
+
+			pool.Put(bsOne)
+
+			<-done
+
+			pool.ModifyCapacity(-1)
+
+			pool.Put(bsTwo)
+
+			pool.ModifyCapacity(-1)
+
+			// any excess returns should drop
+			rando := make([]byte, 0)
+			pool.Put(&rando)
+		}()
+	}
+	wg.Wait()
+
+	if e, a := 0, len(pool.slices); e != a {
+		t.Errorf("expected %v, got %v", e, a)
+	}
+	if e, a := 0, len(pool.allocations); e != a {
+		t.Errorf("expected %v, got %v", e, a)
+	}
+	if e, a := 0, pool.max; e != a {
+		t.Errorf("expected %v, got %v", e, a)
+	}
+
+	_, err := pool.Get(context.Background())
+	if err == nil {
+		t.Errorf("expected error on zero capacity pool")
+	}
+
+	pool.Close()
+}
+
+type recordedPartPool struct {
+	recordedAllocs      uint64
+	recordedGets        uint64
+	recordedOutstanding int64
+	*maxSlicePool
+}
+
+func newRecordedPartPool(sliceSize int64) *recordedPartPool {
+	sp := newMaxSlicePool(sliceSize)
+
+	rp := &recordedPartPool{}
+
+	allocator := sp.allocator
+	sp.allocator = func() *[]byte {
+		atomic.AddUint64(&rp.recordedAllocs, 1)
+		return allocator()
+	}
+
+	rp.maxSlicePool = sp
+
+	return rp
+}
+
+func (r *recordedPartPool) Get(ctx aws.Context) (*[]byte, error) {
+	atomic.AddUint64(&r.recordedGets, 1)
+	atomic.AddInt64(&r.recordedOutstanding, 1)
+	return r.maxSlicePool.Get(ctx)
+}
+
+func (r *recordedPartPool) Put(b *[]byte) {
+	atomic.AddInt64(&r.recordedOutstanding, -1)
+	r.maxSlicePool.Put(b)
+}
+
+func swapByteSlicePool(f func(sliceSize int64) byteSlicePool) func() {
+	orig := newByteSlicePool
+
+	newByteSlicePool = f
+
+	return func() {
+		newByteSlicePool = orig
+	}
+}
+
+type syncSlicePool struct {
+	sync.Pool
+	sliceSize int64
+}
+
+func newSyncSlicePool(sliceSize int64) *syncSlicePool {
+	p := &syncSlicePool{sliceSize: sliceSize}
+	p.New = func() interface{} {
+		bs := make([]byte, p.sliceSize)
+		return &bs
+	}
+	return p
+}
+
+func (s *syncSlicePool) Get(ctx aws.Context) (*[]byte, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		return s.Pool.Get().(*[]byte), nil
+	}
+}
+
+func (s *syncSlicePool) Put(bs *[]byte) {
+	s.Pool.Put(bs)
+}
+
+func (s *syncSlicePool) ModifyCapacity(_ int) {
+	return
+}
+
+func (s *syncSlicePool) SliceSize() int64 {
+	return s.sliceSize
+}
+
+func (s *syncSlicePool) Close() {
+	return
+}
diff --git a/service/s3/s3manager/upload_internal_test.go b/service/s3/s3manager/upload_internal_test.go
index 687059c60..502df232c 100644
--- a/service/s3/s3manager/upload_internal_test.go
+++ b/service/s3/s3manager/upload_internal_test.go
@@ -9,6 +9,7 @@ import (
 	"io/ioutil"
 	random "math/rand"
 	"net/http"
+	"strconv"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -21,49 +22,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/s3/internal/s3testing"
 )
 
-type recordedPartPool struct {
-	allocations uint64
-	gets        uint64
-	outstanding int64
-	*partPool
-}
-
-func newRecordedPartPool(partSize int64) *recordedPartPool {
-	rp := &recordedPartPool{}
-	pp := newPartPool(partSize)
-
-	n := pp.New
-	pp.New = func() interface{} {
-		atomic.AddUint64(&rp.allocations, 1)
-		return n()
-	}
-
-	rp.partPool = pp
-
-	return rp
-}
-
-func (r *recordedPartPool) Get() *[]byte {
-	atomic.AddUint64(&r.gets, 1)
-	atomic.AddInt64(&r.outstanding, 1)
-	return r.partPool.Get()
-}
-
-func (r *recordedPartPool) Put(b *[]byte) {
-	atomic.AddInt64(&r.outstanding, -1)
-	r.partPool.Put(b)
-}
-
-func swapByteSlicePool(f func(partSize int64) byteSlicePool) func() {
-	orig := newByteSlicePool
-
-	newByteSlicePool = f
-
-	return func() {
-		newByteSlicePool = orig
-	}
-}
-
 type testReader struct {
 	br *bytes.Reader
 	m  sync.Mutex
@@ -77,16 +35,28 @@ func (r *testReader) Read(p []byte) (n int, err error) {
 
 func TestUploadByteSlicePool(t *testing.T) {
 	cases := map[string]struct {
-		PartSize int64
-		FileSize int64
+		PartSize      int64
+		FileSize      int64
+		Concurrency   int
+		ExAllocations uint64
 	}{
-		"single part": {
-			PartSize: sdkio.MebiByte * 5,
-			FileSize: sdkio.MebiByte * 5,
+		"single part, single concurrency": {
+			PartSize:      sdkio.MebiByte * 5,
+			FileSize:      sdkio.MebiByte * 5,
+			ExAllocations: 2,
+			Concurrency:   1,
 		},
-		"multi-part": {
-			PartSize: sdkio.MebiByte * 5,
-			FileSize: sdkio.MebiByte * 10,
+		"multi-part, single concurrency": {
+			PartSize:      sdkio.MebiByte * 5,
+			FileSize:      sdkio.MebiByte * 10,
+			ExAllocations: 2,
+			Concurrency:   1,
+		},
+		"multi-part, multiple concurrency": {
+			PartSize:      sdkio.MebiByte * 5,
+			FileSize:      sdkio.MebiByte * 20,
+			ExAllocations: 3,
+			Concurrency:   2,
 		},
 	}
 
@@ -94,8 +64,8 @@ func TestUploadByteSlicePool(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			var p *recordedPartPool
 
-			unswap := swapByteSlicePool(func(partSize int64) byteSlicePool {
-				p = newRecordedPartPool(partSize)
+			unswap := swapByteSlicePool(func(sliceSize int64) byteSlicePool {
+				p = newRecordedPartPool(sliceSize)
 				return p
 			})
 			defer unswap()
@@ -131,7 +101,7 @@ func TestUploadByteSlicePool(t *testing.T) {
 
 			uploader := NewUploaderWithClient(svc, func(u *Uploader) {
 				u.PartSize = tt.PartSize
-				u.Concurrency = 50
+				u.Concurrency = tt.Concurrency
 			})
 
 			expected := s3testing.GetTestBytes(int(tt.FileSize))
@@ -144,13 +114,16 @@ func TestUploadByteSlicePool(t *testing.T) {
 				t.Errorf("expected no error, but got %v", err)
 			}
 
-			if v := atomic.LoadInt64(&p.outstanding); v != 0 {
+			if v := atomic.LoadInt64(&p.recordedOutstanding); v != 0 {
 				t.Fatalf("expected zero outsnatding pool parts, got %d", v)
 			}
 
-			gets, allocs := atomic.LoadUint64(&p.gets), atomic.LoadUint64(&p.allocations)
+			gets, allocs := atomic.LoadUint64(&p.recordedGets), atomic.LoadUint64(&p.recordedAllocs)
 
 			t.Logf("total gets %v, total allocations %v", gets, allocs)
+			if e, a := tt.ExAllocations, allocs; a > e {
+				t.Errorf("expected %v allocations, got %v", e, a)
+			}
 		})
 	}
 }
@@ -185,8 +158,8 @@ func TestUploadByteSlicePool_Failures(t *testing.T) {
 				t.Run(operation, func(t *testing.T) {
 					var p *recordedPartPool
 
-					unswap := swapByteSlicePool(func(partSize int64) byteSlicePool {
-						p = newRecordedPartPool(partSize)
+					unswap := swapByteSlicePool(func(sliceSize int64) byteSlicePool {
+						p = newRecordedPartPool(sliceSize)
 						return p
 					})
 					defer unswap()
@@ -245,7 +218,7 @@ func TestUploadByteSlicePool_Failures(t *testing.T) {
 						t.Fatalf("expected error but got none")
 					}
 
-					if v := atomic.LoadInt64(&p.outstanding); v != 0 {
+					if v := atomic.LoadInt64(&p.recordedOutstanding); v != 0 {
 						t.Fatalf("expected zero outsnatding pool parts, got %d", v)
 					}
 				})
@@ -253,3 +226,194 @@ func TestUploadByteSlicePool_Failures(t *testing.T) {
 		})
 	}
 }
+
+func TestUploadByteSlicePoolConcurrentMultiPartSize(t *testing.T) {
+	var (
+		pools []*recordedPartPool
+		mtx   sync.Mutex
+	)
+
+	unswap := swapByteSlicePool(func(sliceSize int64) byteSlicePool {
+		mtx.Lock()
+		defer mtx.Unlock()
+		b := newRecordedPartPool(sliceSize)
+		pools = append(pools, b)
+		return b
+	})
+	defer unswap()
+
+	sess := unit.Session.Copy()
+	svc := s3.New(sess)
+	svc.Handlers.Unmarshal.Clear()
+	svc.Handlers.UnmarshalMeta.Clear()
+	svc.Handlers.UnmarshalError.Clear()
+	svc.Handlers.Send.Clear()
+	svc.Handlers.Send.PushFront(func(r *request.Request) {
+		if r.Body != nil {
+			io.Copy(ioutil.Discard, r.Body)
+		}
+
+		r.HTTPResponse = &http.Response{
+			StatusCode: 200,
+			Body:       ioutil.NopCloser(bytes.NewReader([]byte{})),
+		}
+
+		switch data := r.Data.(type) {
+		case *s3.CreateMultipartUploadOutput:
+			data.UploadId = aws.String("UPLOAD-ID")
+		case *s3.UploadPartOutput:
+			data.ETag = aws.String(fmt.Sprintf("ETAG%d", random.Int()))
+		case *s3.CompleteMultipartUploadOutput:
+			data.Location = aws.String("https://location")
+			data.VersionId = aws.String("VERSION-ID")
+		case *s3.PutObjectOutput:
+			data.VersionId = aws.String("VERSION-ID")
+		}
+	})
+
+	uploader := NewUploaderWithClient(svc, func(u *Uploader) {
+		u.PartSize = 5 * sdkio.MebiByte
+		u.Concurrency = 2
+	})
+
+	var wg sync.WaitGroup
+	for i := 0; i < 2; i++ {
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			expected := s3testing.GetTestBytes(int(15 * sdkio.MebiByte))
+			_, err := uploader.Upload(&UploadInput{
+				Bucket: aws.String("bucket"),
+				Key:    aws.String("key"),
+				Body:   &testReader{br: bytes.NewReader(expected)},
+			})
+			if err != nil {
+				t.Errorf("expected no error, but got %v", err)
+			}
+		}()
+		go func() {
+			defer wg.Done()
+			expected := s3testing.GetTestBytes(int(15 * sdkio.MebiByte))
+			_, err := uploader.Upload(&UploadInput{
+				Bucket: aws.String("bucket"),
+				Key:    aws.String("key"),
+				Body:   &testReader{br: bytes.NewReader(expected)},
+			}, func(u *Uploader) {
+				u.PartSize = 6 * sdkio.MebiByte
+			})
+			if err != nil {
+				t.Errorf("expected no error, but got %v", err)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	if e, a := 3, len(pools); e != a {
+		t.Errorf("expected %v, got %v", e, a)
+	}
+
+	for _, p := range pools {
+		if v := atomic.LoadInt64(&p.recordedOutstanding); v != 0 {
+			t.Fatalf("expected zero outsnatding pool parts, got %d", v)
+		}
+
+		t.Logf("total gets %v, total allocations %v",
+			atomic.LoadUint64(&p.recordedGets),
+			atomic.LoadUint64(&p.recordedAllocs))
+	}
+}
+
+func BenchmarkPools(b *testing.B) {
+	cases := []struct {
+		PartSize      int64
+		FileSize      int64
+		Concurrency   int
+		ExAllocations uint64
+	}{
+		0: {
+			PartSize:    sdkio.MebiByte * 5,
+			FileSize:    sdkio.MebiByte * 5,
+			Concurrency: 1,
+		},
+		1: {
+			PartSize:    sdkio.MebiByte * 5,
+			FileSize:    sdkio.MebiByte * 10,
+			Concurrency: 1,
+		},
+		2: {
+			PartSize:    sdkio.MebiByte * 5,
+			FileSize:    sdkio.MebiByte * 20,
+			Concurrency: 2,
+		},
+		3: {
+			PartSize:    sdkio.MebiByte * 5,
+			FileSize:    sdkio.MebiByte * 250,
+			Concurrency: 10,
+		},
+	}
+
+	sess := unit.Session.Copy()
+	svc := s3.New(sess)
+	svc.Handlers.Unmarshal.Clear()
+	svc.Handlers.UnmarshalMeta.Clear()
+	svc.Handlers.UnmarshalError.Clear()
+	svc.Handlers.Send.Clear()
+	svc.Handlers.Send.PushFront(func(r *request.Request) {
+		if r.Body != nil {
+			io.Copy(ioutil.Discard, r.Body)
+		}
+
+		r.HTTPResponse = &http.Response{
+			StatusCode: 200,
+			Body:       ioutil.NopCloser(bytes.NewReader([]byte{})),
+		}
+
+		switch data := r.Data.(type) {
+		case *s3.CreateMultipartUploadOutput:
+			data.UploadId = aws.String("UPLOAD-ID")
+		case *s3.UploadPartOutput:
+			data.ETag = aws.String(fmt.Sprintf("ETAG%d", random.Int()))
+		case *s3.CompleteMultipartUploadOutput:
+			data.Location = aws.String("https://location")
+			data.VersionId = aws.String("VERSION-ID")
+		case *s3.PutObjectOutput:
+			data.VersionId = aws.String("VERSION-ID")
+		}
+	})
+
+	pools := map[string]func(sliceSize int64) byteSlicePool{
+		"sync.Pool": func(sliceSize int64) byteSlicePool {
+			return newSyncSlicePool(sliceSize)
+		},
+		"custom": func(sliceSize int64) byteSlicePool {
+			return newMaxSlicePool(sliceSize)
+		},
+	}
+
+	for name, poolFunc := range pools {
+		b.Run(name, func(b *testing.B) {
+			unswap := swapByteSlicePool(poolFunc)
+			defer unswap()
+			for i, c := range cases {
+				b.Run(strconv.Itoa(i), func(b *testing.B) {
+					uploader := NewUploaderWithClient(svc, func(u *Uploader) {
+						u.PartSize = c.PartSize
+						u.Concurrency = c.Concurrency
+					})
+
+					expected := s3testing.GetTestBytes(int(c.FileSize))
+					b.ResetTimer()
+					_, err := uploader.Upload(&UploadInput{
+						Bucket: aws.String("bucket"),
+						Key:    aws.String("key"),
+						Body:   &testReader{br: bytes.NewReader(expected)},
+					})
+					if err != nil {
+						b.Fatalf("expected no error, but got %v", err)
+					}
+				})
+			}
+		})
+	}
+}