From 38efd985078d9710c45e18798174ba6898b5cf64 Mon Sep 17 00:00:00 2001
From: Aleksandar Aleksandrov <aleksandar@messagebird.com>
Date: Thu, 1 Mar 2018 18:07:31 +0100
Subject: [PATCH] Add the chunk_size optional parameter to gcs storage

---
 physical/gcs/gcs.go                           | 25 ++++++++++++++++---
 .../storage/google-cloud.html.md              |  7 ++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/physical/gcs/gcs.go b/physical/gcs/gcs.go
index c81e0489699f..85656d899ec1 100644
--- a/physical/gcs/gcs.go
+++ b/physical/gcs/gcs.go
@@ -21,9 +21,6 @@ import (
 	"google.golang.org/api/option"
 )
 
-// Verify GCSBackend satisfies the correct interfaces
-var _ physical.Backend = (*GCSBackend)(nil)
-
 // GCSBackend is a physical backend that stores data
 // within an Google Cloud Storage bucket.
 type GCSBackend struct {
@@ -33,6 +30,15 @@ type GCSBackend struct {
 	logger     log.Logger
 }
 
+var (
+	// Verify GCSBackend satisfies the correct interfaces
+	_ physical.Backend = (*GCSBackend)(nil)
+
+	// Number of bytes the writer will attempt to write in a single request.
+	// Defaults to 8Mb, as defined in the gcs library
+	chunkSize = 8 * 1024 * 1024
+)
+
 // NewGCSBackend constructs a Google Cloud Storage backend using a pre-existing
 // bucket. Credentials can be provided to the backend, sourced
 // from environment variables or a service account file
@@ -70,6 +76,18 @@ func NewGCSBackend(conf map[string]string, logger log.Logger) (physical.Backend,
 		}
 	}
 
+	chunkSizeStr, ok := conf["chunk_size"]
+	if ok {
+		chunkSize, err = strconv.Atoi(chunkSizeStr)
+		if err != nil {
+			return nil, errwrap.Wrapf("failed parsing chunk_size parameter: {{err}}", err)
+		}
+		chunkSize *= 1024
+		if logger.IsDebug() {
+			logger.Debug("physical/gcs: chunk_size set", "chunk_size", chunkSize)
+		}
+	}
+
 	g := GCSBackend{
 		bucketName: bucketName,
 		client:     client,
@@ -111,6 +129,7 @@ func (g *GCSBackend) Put(ctx context.Context, entry *physical.Entry) error {
 
 	bucket := g.client.Bucket(g.bucketName)
 	writer := bucket.Object(entry.Key).NewWriter(context.Background())
+	writer.ChunkSize = chunkSize
 
 	g.permitPool.Acquire()
 	defer g.permitPool.Release()
diff --git a/website/source/docs/configuration/storage/google-cloud.html.md b/website/source/docs/configuration/storage/google-cloud.html.md
index 2e6a98b66a74..7e8940c8e8e6 100644
--- a/website/source/docs/configuration/storage/google-cloud.html.md
+++ b/website/source/docs/configuration/storage/google-cloud.html.md
@@ -42,6 +42,13 @@ storage "gcs" {
 - `max_parallel` `(string: "128")` – Specifies the maximum number of concurrent
   requests.
 
+- `chunk_size` `(string: "8192")` – Specifies the maximum kilobytes of each object
+  the gcs writer will attempt to send to the server in a single request.
+  If set to 0, it will attempt to send the whole object at once, but it will
+  not retry any failures either (not recommended). If you are not storing large
+  objects in Vault, it is recommended to set this to a low value (minimum is 256)
+  since it will drastically reduce the amount of memory Vault uses.
+
 ## `gcs` Examples
 
 ### Default Example