From 307504713f5f6a8846cdc762898446ac87310119 Mon Sep 17 00:00:00 2001
From: Arthur Baars <arthur@semmle.com>
Date: Sun, 14 Feb 2016 18:15:15 +0000
Subject: [PATCH] Storagedriver: GCS: add chunksize parameter

Signed-off-by: Arthur Baars <arthur@semmle.com>
---
 docs/configuration.md                   |  1 +
 docs/storage-drivers/gcs.md             | 11 +++++++
 registry/storage/driver/gcs/gcs.go      | 38 +++++++++++++++++++++++--
 registry/storage/driver/gcs/gcs_test.go |  1 +
 4 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/docs/configuration.md b/docs/configuration.md
index ef01d2f6e..f5dc05779 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -86,6 +86,7 @@ information about each option that appears later in this page.
         bucket: bucketname
         keyfile: /path/to/keyfile
         rootdirectory: /gcs/object/name/prefix
+        chunksize: 5242880
       s3:
         accesskey: awsaccesskey
         secretkey: awssecretkey
diff --git a/docs/storage-drivers/gcs.md b/docs/storage-drivers/gcs.md
index 33cc94b37..0aa9b30dd 100644
--- a/docs/storage-drivers/gcs.md
+++ b/docs/storage-drivers/gcs.md
@@ -52,6 +52,17 @@ An implementation of the `storagedriver.StorageDriver` interface which uses Goog
     <td>
       This is a prefix that will be applied to all Google Cloud Storage keys to allow you to segment data in your bucket if necessary.
   </tr>
+  </tr>
+   <tr>
+    <td>
+      <code>chunksize</code>
+    </td>
+    <td>
+      no (default 5242880)
+    </td>
+    <td>
+      This is the chunk size used for uploading large blobs, must be a multiple of 256*1024.
+  </tr>
 
 </table>
 
diff --git a/registry/storage/driver/gcs/gcs.go b/registry/storage/driver/gcs/gcs.go
index 14600dee5..abe0b9f68 100644
--- a/registry/storage/driver/gcs/gcs.go
+++ b/registry/storage/driver/gcs/gcs.go
@@ -22,6 +22,7 @@ import (
 	"math/rand"
 	"net/http"
 	"net/url"
+	"reflect"
 	"regexp"
 	"sort"
 	"strconv"
@@ -50,7 +51,7 @@ const (
 
 	uploadSessionContentType = "application/x-docker-upload-session"
 	minChunkSize             = 256 * 1024
-	maxChunkSize             = 20 * minChunkSize
+	defaultChunkSize         = 20 * minChunkSize
 
 	maxTries = 5
 )
@@ -65,6 +66,7 @@ type driverParameters struct {
 	privateKey    []byte
 	client        *http.Client
 	rootDirectory string
+	chunkSize     int
 }
 
 func init() {
@@ -87,6 +89,7 @@ type driver struct {
 	email         string
 	privateKey    []byte
 	rootDirectory string
+	chunkSize     int
 }
 
 // FromParameters constructs a new Driver with a given parameters map
@@ -103,6 +106,31 @@ func FromParameters(parameters map[string]interface{}) (storagedriver.StorageDri
 		rootDirectory = ""
 	}
 
+	chunkSize := defaultChunkSize
+	chunkSizeParam, ok := parameters["chunksize"]
+	if ok {
+		switch v := chunkSizeParam.(type) {
+		case string:
+			vv, err := strconv.Atoi(v)
+			if err != nil {
+				return nil, fmt.Errorf("chunksize parameter must be an integer, %v invalid", chunkSizeParam)
+			}
+			chunkSize = vv
+		case int, uint, int32, uint32, uint64, int64:
+			chunkSize = int(reflect.ValueOf(v).Convert(reflect.TypeOf(chunkSize)).Int())
+		default:
+			return nil, fmt.Errorf("invalid valud for chunksize: %#v", chunkSizeParam)
+		}
+
+		if chunkSize < minChunkSize {
+			return nil, fmt.Errorf("The chunksize %#v parameter should be a number that is larger than or equal to %d", chunkSize, minChunkSize)
+		}
+
+		if chunkSize%minChunkSize != 0 {
+			return nil, fmt.Errorf("chunksize should be a multiple of %d", minChunkSize)
+		}
+	}
+
 	var ts oauth2.TokenSource
 	jwtConf := new(jwt.Config)
 	if keyfile, ok := parameters["keyfile"]; ok {
@@ -121,7 +149,6 @@ func FromParameters(parameters map[string]interface{}) (storagedriver.StorageDri
 		if err != nil {
 			return nil, err
 		}
-
 	}
 
 	params := driverParameters{
@@ -130,6 +157,7 @@ func FromParameters(parameters map[string]interface{}) (storagedriver.StorageDri
 		email:         jwtConf.Email,
 		privateKey:    jwtConf.PrivateKey,
 		client:        oauth2.NewClient(context.Background(), ts),
+		chunkSize:     chunkSize,
 	}
 
 	return New(params)
@@ -141,12 +169,16 @@ func New(params driverParameters) (storagedriver.StorageDriver, error) {
 	if rootDirectory != "" {
 		rootDirectory += "/"
 	}
+	if params.chunkSize <= 0 || params.chunkSize%minChunkSize != 0 {
+		return nil, fmt.Errorf("Invalid chunksize: %d is not a positive multiple of %d", params.chunkSize, minChunkSize)
+	}
 	d := &driver{
 		bucket:        params.bucket,
 		rootDirectory: rootDirectory,
 		email:         params.email,
 		privateKey:    params.privateKey,
 		client:        params.client,
+		chunkSize:     params.chunkSize,
 	}
 
 	return &base.Base{
@@ -263,7 +295,7 @@ func (d *driver) Writer(context ctx.Context, path string, append bool) (storaged
 		client: d.client,
 		bucket: d.bucket,
 		name:   d.pathToKey(path),
-		buffer: make([]byte, maxChunkSize),
+		buffer: make([]byte, d.chunkSize),
 	}
 
 	if append {
diff --git a/registry/storage/driver/gcs/gcs_test.go b/registry/storage/driver/gcs/gcs_test.go
index 4852bf2c4..f2808d5fc 100644
--- a/registry/storage/driver/gcs/gcs_test.go
+++ b/registry/storage/driver/gcs/gcs_test.go
@@ -75,6 +75,7 @@ func init() {
 			email:         email,
 			privateKey:    privateKey,
 			client:        oauth2.NewClient(ctx.Background(), ts),
+			chunkSize:     defaultChunkSize,
 		}
 
 		return New(parameters)