gensupport: Allow user to provide his own buffer used for uploading

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.

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.

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
 }