support "g" (GiB) as a size unit

Signed-off-by: Shinya Hayashi <[email protected]>
peng225 · Nov 5, 2023 · a4d2963 · a4d2963
1 parent efd1181
commit a4d2963
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/internal/argparser/argparser_test.go b/internal/argparser/argparser_test.go
@@ -35,6 +35,11 @@ func TestParseMultipartThresh(t *testing.T) {
 		},
 		{
 			multipartThreshStr:      "8g",
+			expectedMultipartThresh: 8 * 1024 * 1024 * 1024,
+			expectedErr:             false,
+		},
+		{
+			multipartThreshStr:      "5t",
 			expectedMultipartThresh: 0,
 			expectedErr:             true,
 		},

diff --git a/internal/argparser/size.go b/internal/argparser/size.go
@@ -38,8 +38,9 @@ func parseSizeUnit(s string) (int, error) {
 	unit := map[string]int{
 		"k": 1024,
 		"m": 1024 * 1024,
+		"g": 1024 * 1024 * 1024,
 	}
-	r := regexp.MustCompile("^[1-9][0-9]*[km]*$")
+	r := regexp.MustCompile("^[1-9][0-9]*[kmg]*$")
 	if r.MatchString(s) {
 		if size, err := strconv.Atoi(s); err == nil {
 			return size, nil
@@ -51,7 +52,7 @@ func parseSizeUnit(s string) (int, error) {
 			return baseNum * unit[s[len(s)-1:]], nil
 		}
 	}
-	return 0, fmt.Errorf("Illegal size format: %v\n", s)
+	return 0, fmt.Errorf("illegal size format: %v\n", s)
 }
 
 func ParseMultipartThresh(s string) (int, error) {

diff --git a/internal/cmd/root.go b/internal/cmd/root.go
@@ -162,10 +162,10 @@ func handleCommonFlags() {
 func defineCommonFlags(cmd *cobra.Command) {
 	cmd.Flags().IntVar(&numObj, "num_obj", 10, "The maximum number of objects per process.")
 	cmd.Flags().IntVar(&numWorker, "num_worker", 1, "The number of workers per process.")
-	cmd.Flags().StringVar(&sizePattern, "size", "4k", "The size of object. Should be in the form like \"8k\" or \"4k-2m\". Only \"k\" and \"m\" is allowed as an unit.")
+	cmd.Flags().StringVar(&sizePattern, "size", "4k", `The size of object. Should be in the form like "8k" or "4k-2m". Only "k", "m" and "g" is allowed as an unit.`)
 	cmd.Flags().DurationVar(&execTime, "time", time.Second*3, "Time duration for run the workload.")
 	cmd.Flags().StringSliceVar(&bucketNames, "bucket", nil, "The name list of the buckets. e.g. \"bucket1,bucket2\"")
 	cmd.Flags().StringVar(&opeRatioStr, "ope_ratio", "1,1,1,0", "The ration of put, get, delete and list operations. e.g. \"2,3,1,1\"")
 	cmd.Flags().StringVar(&endpoint, "endpoint", "", "The endpoint URL and TCP port number. e.g. \"http://127.0.0.1:9000\"")
-	cmd.Flags().StringVar(&multipartThreshStr, "multipart_thresh", "100m", "The threshold of the object size to switch to the multipart upload. Only \"k\" and \"m\" is allowed as an unit.")
+	cmd.Flags().StringVar(&multipartThreshStr, "multipart_thresh", "100m", `The threshold of the object size to switch to the multipart upload. Only "k", "m" and "g" is allowed as an unit.`)
 }