From 117757a5cb8ae931f5eab4de24763d8c0421c522 Mon Sep 17 00:00:00 2001
From: "libo.huang" <huanglibo2010@gmail.com>
Date: Fri, 31 Dec 2021 14:13:16 +0800
Subject: [PATCH] feat: add option to disable combining the pending part

Signed-off-by: Libo Huang <huanglibo2010@gmail.com>
---
 registry/storage/driver/s3-aws/s3.go      | 23 ++++++++++++++++++++++-
 registry/storage/driver/s3-aws/s3_test.go |  6 ++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/registry/storage/driver/s3-aws/s3.go b/registry/storage/driver/s3-aws/s3.go
index b9388ff4..16414b1a 100644
--- a/registry/storage/driver/s3-aws/s3.go
+++ b/registry/storage/driver/s3-aws/s3.go
@@ -97,6 +97,7 @@ type DriverParameters struct {
 	MultipartCopyChunkSize      int64
 	MultipartCopyMaxConcurrency int64
 	MultipartCopyThresholdSize  int64
+	MultipartCombineSmallPart   bool
 	RootDirectory               string
 	StorageClass                string
 	UserAgent                   string
@@ -146,6 +147,7 @@ type driver struct {
 	MultipartCopyChunkSize      int64
 	MultipartCopyMaxConcurrency int64
 	MultipartCopyThresholdSize  int64
+	MultipartCombineSmallPart   bool
 	RootDirectory               string
 	StorageClass                string
 	ObjectACL                   string
@@ -356,6 +358,23 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 		return nil, fmt.Errorf("the useDualStack parameter should be a boolean")
 	}
 
+	mutlipartCombineSmallPart := true
+	combine := parameters["multipartcombinesmallpart"]
+	switch combine := combine.(type) {
+	case string:
+		b, err := strconv.ParseBool(combine)
+		if err != nil {
+			return nil, fmt.Errorf("the multipartcombinesmallpart parameter should be a boolean")
+		}
+		mutlipartCombineSmallPart = b
+	case bool:
+		mutlipartCombineSmallPart = combine
+	case nil:
+		// do nothing
+	default:
+		return nil, fmt.Errorf("the multipartcombinesmallpart parameter should be a boolean")
+	}
+
 	sessionToken := ""
 
 	params := DriverParameters{
@@ -373,6 +392,7 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 		multipartCopyChunkSize,
 		multipartCopyMaxConcurrency,
 		multipartCopyThresholdSize,
+		mutlipartCombineSmallPart,
 		fmt.Sprint(rootDirectory),
 		storageClass,
 		fmt.Sprint(userAgent),
@@ -497,6 +517,7 @@ func New(params DriverParameters) (*Driver, error) {
 		MultipartCopyChunkSize:      params.MultipartCopyChunkSize,
 		MultipartCopyMaxConcurrency: params.MultipartCopyMaxConcurrency,
 		MultipartCopyThresholdSize:  params.MultipartCopyThresholdSize,
+		MultipartCombineSmallPart:   params.MultipartCombineSmallPart,
 		RootDirectory:               params.RootDirectory,
 		StorageClass:                params.StorageClass,
 		ObjectACL:                   params.ObjectACL,
@@ -1386,7 +1407,7 @@ func (w *writer) flushPart() error {
 		// nothing to write
 		return nil
 	}
-	if len(w.pendingPart) < int(w.driver.ChunkSize) {
+	if w.driver.MultipartCombineSmallPart && len(w.pendingPart) < int(w.driver.ChunkSize) {
 		// closing with a small pending part
 		// combine ready and pending to avoid writing a small part
 		w.readyPart = append(w.readyPart, w.pendingPart...)
diff --git a/registry/storage/driver/s3-aws/s3_test.go b/registry/storage/driver/s3-aws/s3_test.go
index 4d4631a7..31662b17 100644
--- a/registry/storage/driver/s3-aws/s3_test.go
+++ b/registry/storage/driver/s3-aws/s3_test.go
@@ -43,6 +43,7 @@ func init() {
 	regionEndpoint := os.Getenv("REGION_ENDPOINT")
 	sessionToken := os.Getenv("AWS_SESSION_TOKEN")
 	useDualStack := os.Getenv("S3_USE_DUALSTACK")
+	combineSmallPart := os.Getenv("MULTIPART_COMBINE_SMALL_PART")
 	if err != nil {
 		panic(err)
 	}
@@ -84,6 +85,10 @@ func init() {
 		useDualStackBool := false
 		if useDualStack != "" {
 			useDualStackBool, err = strconv.ParseBool(useDualStack)
+		}
+		multipartCombineSmallPart := true
+		if combineSmallPart != "" {
+			multipartCombineSmallPart, err = strconv.ParseBool(combineSmallPart)
 			if err != nil {
 				return nil, err
 			}
@@ -104,6 +109,7 @@ func init() {
 			defaultMultipartCopyChunkSize,
 			defaultMultipartCopyMaxConcurrency,
 			defaultMultipartCopyThresholdSize,
+			multipartCombineSmallPart,
 			rootDirectory,
 			storageClass,
 			driverName + "-test",