You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by lo...@apache.org on 2022/11/05 16:55:50 UTC

[beam] 04/12: Use humanize library to parse byte sizes.

This is an automated email from the ASF dual-hosted git repository.

lostluck pushed a commit to branch beam23893pipelinehint
in repository https://gitbox.apache.org/repos/asf/beam.git

commit fa3513203c8a36cdcc8f8c427eb5840e323b3289
Author: lostluck <13...@users.noreply.github.com>
AuthorDate: Fri Nov 4 14:40:08 2022 -0700

    Use humanize library to parse byte sizes.
---
 sdks/go.mod                                    |  1 +
 sdks/go.sum                                    |  1 +
 sdks/go/pkg/beam/options/resource/hint.go      | 22 ++++++++++++++----
 sdks/go/pkg/beam/options/resource/hint_test.go | 31 ++++++++++++++++++++++----
 4 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/sdks/go.mod b/sdks/go.mod
index 27670370ecf..fcfb71e83f7 100644
--- a/sdks/go.mod
+++ b/sdks/go.mod
@@ -29,6 +29,7 @@ require (
 	cloud.google.com/go/pubsub v1.26.0
 	cloud.google.com/go/storage v1.27.0
 	github.com/docker/go-connections v0.4.0
+	github.com/dustin/go-humanize v1.0.0
 	github.com/go-sql-driver/mysql v1.6.0
 	github.com/golang/protobuf v1.5.2 // TODO(danoliveira): Fully replace this with google.golang.org/protobuf
 	github.com/google/go-cmp v0.5.9
diff --git a/sdks/go.sum b/sdks/go.sum
index 55d19289d8e..e01740082be 100644
--- a/sdks/go.sum
+++ b/sdks/go.sum
@@ -392,6 +392,7 @@ github.com/docker/libtrust v0.0.0-20150114040149-fa567046d9b1/go.mod h1:cyGadeNE
 github.com/docker/spdystream v0.0.0-20160310174837-449fdfce4d96/go.mod h1:Qh8CwZgvJUkLughtfhJv5dyTYa91l1fOUCrgjqmcifM=
 github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
+github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/elazarl/goproxy v0.0.0-20180725130230-947c36da3153/go.mod h1:/Zj4wYkgs4iZTTu3o/KG3Itv/qCCa8VVMlb3i9OVuzc=
 github.com/emicklei/go-restful v0.0.0-20170410110728-ff4f55a20633/go.mod h1:otzb+WCGbkyDHkqmQmT5YD2WR4BBwUdeQoFo8l/7tVs=
diff --git a/sdks/go/pkg/beam/options/resource/hint.go b/sdks/go/pkg/beam/options/resource/hint.go
index 2772737921a..de78950b522 100644
--- a/sdks/go/pkg/beam/options/resource/hint.go
+++ b/sdks/go/pkg/beam/options/resource/hint.go
@@ -23,6 +23,8 @@ import (
 	"bytes"
 	"fmt"
 	"strconv"
+
+	"github.com/dustin/go-humanize"
 )
 
 // Hints contains a list of hints for a given scope.
@@ -99,11 +101,23 @@ type Hint interface {
 //
 // See https://beam.apache.org/documentation/runtime/resource-hints/ for more information about
 // resource hints.
-func MinRamBytes(v int64) Hint {
-	if v < 0 {
-		panic(fmt.Sprintf("negative min ram hint: %v", v))
+func MinRamBytes(v uint64) Hint {
+	return minRamHint{value: int64(v)}
+}
+
+// ParseMinRam converts various byte units, including MB, GB, MiB, and GiB into a hint.
+// An invalid byte size format will cause ParseMinRam to panic.
+//
+// Hints are advisory only and runners may not respect them.
+//
+// See https://beam.apache.org/documentation/runtime/resource-hints/ for more information about
+// resource hints.
+func ParseMinRam(v string) Hint {
+	b, err := humanize.ParseBytes(v)
+	if err != nil {
+		panic(fmt.Sprintf("resource.ParseMinRam: unable to parse %q: %v", v, err))
 	}
-	return minRamHint{value: v}
+	return MinRamBytes(b)
 }
 
 type minRamHint struct {
diff --git a/sdks/go/pkg/beam/options/resource/hint_test.go b/sdks/go/pkg/beam/options/resource/hint_test.go
index 1d1ccce7a0a..45c742ecae0 100644
--- a/sdks/go/pkg/beam/options/resource/hint_test.go
+++ b/sdks/go/pkg/beam/options/resource/hint_test.go
@@ -77,15 +77,38 @@ func TestMinRamBytesHint_Payload(t *testing.T) {
 	}
 }
 
-func TestMinRamBytes_Panic(t *testing.T) {
+func TestParseMinRamHint(t *testing.T) {
+	tests := []struct {
+		value   string
+		payload string
+	}{
+		{"0", "0"},
+		{"2", "2"},
+		{"11", "11"},
+		{"2003", "2003"},
+		{"1.23MB", "1230000"},
+		{"1.23MiB", "1289748"},
+		{"4GB", "4000000000"},
+		{"2GiB", "2147483648"},
+		{"1.4KiB", "1433"},
+	}
+
+	for _, test := range tests {
+		h := ParseMinRam(test.value)
+		if got, want := h.Payload(), []byte(test.payload); !bytes.Equal(got, want) {
+			t.Errorf("%v.Payload() = %v, want %v", h, string(got), string(want))
+		}
+	}
+}
+
+func TestParseMinRamHint_panic(t *testing.T) {
 	defer func() {
 		if r := recover(); r == nil {
-			t.Error("expected panic, but didn't get one.")
+			t.Error("want ParseMinRam to panic")
 		}
 	}()
-	MinRamBytes(-1)
+	ParseMinRam("a bad byte string")
 }
-
 // We copy the URN from the proto for use as a constant rather than perform a direct look up
 // each time, or increase initialization time. However we do need to validate that they are
 // correct, and match the standard hint urns, so that's done here.