service/s3/s3manager/pool: prefer using allocated slices over allocat…

…ing new ones if possible

Problem:
S3 uploader could allocate more memory than it is actually needed.
Added unit test should showcase the problem

service/s3/s3manager/pool.go

@@ -60,6 +60,14 @@ func (p *maxSlicePool) Get(ctx aws.Context) (*[]byte, error) {
 				return nil, errZeroCapacity
 			}
 			return bs, nil
+		case <-ctx.Done():
+			p.mtx.RUnlock()
+			return nil, ctx.Err()
+		default:
+			// pass
+		}
+
+		select {
 		case _, ok := <-p.allocations:
 			p.mtx.RUnlock()
 			if !ok {

service/s3/s3manager/pool_test.go

@@ -90,6 +90,33 @@ func TestMaxSlicePool(t *testing.T) {
 	pool.Close()
 }
 
+func TestPoolShouldPreferAllocatedSlicesOverNewAllocations(t *testing.T) {
+	pool := newMaxSlicePool(0)
+	defer pool.Close()
+
+	// Prepare pool: make it so that pool contains 1 allocated slice and 1 allocation permit
+	pool.ModifyCapacity(2)
+	initialSlice, err := pool.Get(context.Background())
+	if err != nil {
+		t.Errorf("failed to get slice from pool: %v", err)
+	}
+	pool.Put(initialSlice)
+
+	for i := 0; i < 100; i++ {
+		newSlice, err := pool.Get(context.Background())
+		if err != nil {
+			t.Errorf("failed to get slice from pool: %v", err)
+			return
+		}
+
+		if newSlice != initialSlice {
+			t.Errorf("pool allocated a new slice despite it having pre-allocated one")
+			return
+		}
+		pool.Put(newSlice)
+	}
+}
+
 type recordedPartPool struct {
 	recordedAllocs      uint64
 	recordedGets        uint64