From 35da91aa89eb18b8a724915ddcbc9d8c727abaec Mon Sep 17 00:00:00 2001
From: Olivier Poitrey <rs@rhapsodyk.net>
Date: Mon, 10 Oct 2016 14:25:17 -0700
Subject: [PATCH] Add ability to set the init width produced by BufferPool

---
 bufferpool.go      | 17 +++++++++++++----
 bufferpool_test.go | 40 +++++++++++++++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/bufferpool.go b/bufferpool.go
index 8c8ac64..c0737fd 100644
--- a/bufferpool.go
+++ b/bufferpool.go
@@ -8,12 +8,21 @@ import (
 // channel.
 type BufferPool struct {
 	c chan *bytes.Buffer
+	w int
 }
 
-// NewBufferPool creates a new BufferPool bounded to the given size.
-func NewBufferPool(size int) (bp *BufferPool) {
+// NewBufferPool creates a new BufferPool bounded to the given maxSize
+// with a buffer width initialized to 0.
+func NewBufferPool(maxSize int) (bp *BufferPool) {
+	return NewBufferPoolWidth(maxSize, 0)
+}
+
+// NewBufferPoolWidth creates a new BufferPool bounded to the given maxSize
+// with new buffer's backend byte slices sized based on width.
+func NewBufferPoolWidth(maxSize, width int) (bp *BufferPool) {
 	return &BufferPool{
-		c: make(chan *bytes.Buffer, size),
+		c: make(chan *bytes.Buffer, maxSize),
+		w: width,
 	}
 }
 
@@ -25,7 +34,7 @@ func (bp *BufferPool) Get() (b *bytes.Buffer) {
 	// reuse existing buffer
 	default:
 		// create new buffer
-		b = bytes.NewBuffer([]byte{})
+		b = bytes.NewBuffer(make([]byte, 0, bp.w))
 	}
 	return
 }
diff --git a/bufferpool_test.go b/bufferpool_test.go
index b344bdc..f1b276c 100644
--- a/bufferpool_test.go
+++ b/bufferpool_test.go
@@ -6,16 +6,46 @@ import (
 )
 
 func TestBufferPool(t *testing.T) {
-	var size int = 4
+	maxSize := 4
 
-	bufPool := NewBufferPool(size)
+	bufPool := NewBufferPool(maxSize)
 
 	// Test Get/Put
 	b := bufPool.Get()
+	if b.Cap() != 0 {
+		t.Errorf("bufferpool width invalid: got %v want %v", b.Cap(), 0)
+	}
+	bufPool.Put(b)
+
+	// Add some additional buffers beyond the pool size.
+	for i := 0; i < maxSize*2; i++ {
+		bufPool.Put(bytes.NewBuffer([]byte{}))
+	}
+
+	// Close the channel so we can iterate over it.
+	close(bufPool.c)
+
+	// Check the size of the pool.
+	if len(bufPool.c) != maxSize {
+		t.Fatalf("bufferpool size invalid: got %v want %v", len(bufPool.c), maxSize)
+	}
+}
+
+func TestBufferPoolWidth(t *testing.T) {
+	maxSize := 4
+	width := 10
+
+	bufPool := NewBufferPoolWidth(maxSize, width)
+
+	// Test Get/Put
+	b := bufPool.Get()
+	if b.Cap() != width {
+		t.Errorf("bufferpool width invalid: got %v want %v", b.Cap(), width)
+	}
 	bufPool.Put(b)
 
 	// Add some additional buffers beyond the pool size.
-	for i := 0; i < size*2; i++ {
+	for i := 0; i < maxSize*2; i++ {
 		bufPool.Put(bytes.NewBuffer([]byte{}))
 	}
 
@@ -23,8 +53,8 @@ func TestBufferPool(t *testing.T) {
 	close(bufPool.c)
 
 	// Check the size of the pool.
-	if len(bufPool.c) != size {
-		t.Fatalf("bufferpool size invalid: got %v want %v", len(bufPool.c), size)
+	if len(bufPool.c) != maxSize {
+		t.Fatalf("bufferpool size invalid: got %v want %v", len(bufPool.c), maxSize)
 	}
 
 }