From 42cd33cbce18d754608f8cac76133209f7ef29f5 Mon Sep 17 00:00:00 2001
From: Maciej Zimnoch <maciej@scylladb.com>
Date: Thu, 20 Aug 2020 17:17:15 +0200
Subject: [PATCH] gensupport: Allow user to provide his own buffer used for
 uploading

---
 googleapi/googleapi.go        | 16 ++++++++++++++++
 internal/gensupport/buffer.go |  5 +++++
 internal/gensupport/media.go  | 11 ++++++++---
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/googleapi/googleapi.go b/googleapi/googleapi.go
index d1784f1a349..69b7badc3ab 100644
--- a/googleapi/googleapi.go
+++ b/googleapi/googleapi.go
@@ -211,6 +211,20 @@ type MediaOption interface {
 	setOptions(o *MediaOptions)
 }
 
+type bufferOption []byte
+
+func (bp bufferOption) setOptions(o *MediaOptions) {
+	o.Buffer = bp
+}
+
+// WithBuffer returns MediaOption which sets buffer used for media uploads.
+// It needs to be at least MinUploadChunkSize long.
+// If used together with ChunkSize, buffer needs to be at least that long.
+// If not set, each upload will allocate its own memory.
+func WithBuffer(buffer []byte) MediaOption {
+	return bufferOption(buffer)
+}
+
 type contentTypeOption string
 
 func (ct contentTypeOption) setOptions(o *MediaOptions) {
@@ -251,6 +265,8 @@ type MediaOptions struct {
 	ForceEmptyContentType bool
 
 	ChunkSize int
+
+	Buffer []byte
 }
 
 // ProcessMediaOptions stores options from opts in a MediaOptions.
diff --git a/internal/gensupport/buffer.go b/internal/gensupport/buffer.go
index 3d0817ede98..c7b1bfb4fe9 100644
--- a/internal/gensupport/buffer.go
+++ b/internal/gensupport/buffer.go
@@ -28,6 +28,11 @@ func NewMediaBuffer(media io.Reader, chunkSize int) *MediaBuffer {
 	return &MediaBuffer{media: media, chunk: make([]byte, 0, chunkSize)}
 }
 
+// NewMediaBuffer initializes a MediaBuffer.
+func NewMediaBufferWithBuffer(media io.Reader, chunkSize int, buffer []byte) *MediaBuffer {
+	return &MediaBuffer{media: media, chunk: buffer}
+}
+
 // Chunk returns the current buffered chunk, the offset in the underlying media
 // from which the chunk is drawn, and the size of the chunk.
 // Successive calls to Chunk return the same chunk between calls to Next.
diff --git a/internal/gensupport/media.go b/internal/gensupport/media.go
index 0288cc30427..581e07fb5d7 100644
--- a/internal/gensupport/media.go
+++ b/internal/gensupport/media.go
@@ -200,11 +200,16 @@ func typeHeader(contentType string) textproto.MIMEHeader {
 //
 // After PrepareUpload has been called, media should no longer be used: the
 // media content should be accessed via one of the return values.
-func PrepareUpload(media io.Reader, chunkSize int) (r io.Reader, mb *MediaBuffer, singleChunk bool) {
+func PrepareUpload(media io.Reader, chunkSize int, buffer []byte) (r io.Reader, mb *MediaBuffer, singleChunk bool) {
 	if chunkSize == 0 { // do not chunk
 		return media, nil, true
 	}
-	mb = NewMediaBuffer(media, chunkSize)
+	if buffer != nil {
+		mb = NewMediaBufferWithBuffer(media, chunkSize, buffer)
+	} else {
+		mb = NewMediaBuffer(media, chunkSize)
+	}
+
 	_, _, _, err := mb.Chunk()
 	// If err is io.EOF, we can upload this in a single request. Otherwise, err is
 	// either nil or a non-EOF error. If it is the latter, then the next call to
@@ -234,7 +239,7 @@ func NewInfoFromMedia(r io.Reader, options []googleapi.MediaOption) *MediaInfo {
 	if !opts.ForceEmptyContentType {
 		r, mi.mType = DetermineContentType(r, opts.ContentType)
 	}
-	mi.media, mi.buffer, mi.singleChunk = PrepareUpload(r, opts.ChunkSize)
+	mi.media, mi.buffer, mi.singleChunk = PrepareUpload(r, opts.ChunkSize, opts.Buffer)
 	return mi
 }