From a203fa414dcc368afd1d7ebd91db45ea79858a37 Mon Sep 17 00:00:00 2001
From: Jarek Kowalski <jaak@jkowalski.net>
Date: Sun, 22 Aug 2021 12:33:36 -0700
Subject: [PATCH] gather: MakeContiguous support for arbitrary chunk sizes
 (#1247)

---
 internal/gather/gather_write_buffer.go      | 10 ++++++++--
 internal/gather/gather_write_buffer_test.go | 13 +++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/internal/gather/gather_write_buffer.go b/internal/gather/gather_write_buffer.go
index a2dfb284e..289ee829e 100644
--- a/internal/gather/gather_write_buffer.go
+++ b/internal/gather/gather_write_buffer.go
@@ -41,8 +41,14 @@ func (b *WriteBuffer) MakeContiguous(length int) []byte {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
-	b.alloc = contiguousAllocator
-	v := b.allocChunk()[0:length]
+	var v []byte
+
+	if length > contiguousAllocator.chunkSize {
+		v = make([]byte, length)
+	} else {
+		b.alloc = contiguousAllocator
+		v = b.allocChunk()[0:length]
+	}
 
 	b.inner.Slices = [][]byte{v}
 
diff --git a/internal/gather/gather_write_buffer_test.go b/internal/gather/gather_write_buffer_test.go
index 25d736112..d8df26563 100644
--- a/internal/gather/gather_write_buffer_test.go
+++ b/internal/gather/gather_write_buffer_test.go
@@ -4,6 +4,8 @@
 	"bytes"
 	"fmt"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestGatherWriteBuffer(t *testing.T) {
@@ -68,3 +70,14 @@ func TestGatherDefaultWriteBuffer(t *testing.T) {
 		t.Errorf("invalid number of slices %v, want %v", got, want)
 	}
 }
+
+func TestGatherWriteBufferContig(t *testing.T) {
+	var w WriteBuffer
+	defer w.Close()
+
+	// allocate more than contig allocator can provide
+	theCap := contiguousAllocator.chunkSize + 10
+	b := w.MakeContiguous(theCap)
+	require.Equal(t, theCap, len(b))
+	require.Equal(t, theCap, cap(b))
+}