runtime: add byte count parser for GOMEMLIMIT

This change adds a parser for the GOMEMLIMIT environment variable's
input. This environment variable accepts a number followed by an
optional prefix expressing the unit. Acceptable units include
B, KiB, MiB, GiB, TiB, where *iB is a power-of-two byte unit.

For #48409.

Change-Id: I6a3b4c02b175bfcf9c4debee6118cf5dda93bb6f
Reviewed-on: https://go-review.googlesource.com/c/go/+/393400
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Michael Pratt <mpratt@google.com>
Run-TryBot: Michael Knyszek <mknyszek@google.com>

src/runtime/export_test.go

@@ -33,6 +33,7 @@ var Fastlog2 = fastlog2
 
 var Atoi = atoi
 var Atoi32 = atoi32
+var ParseByteCount = parseByteCount
 
 var Nanotime = nanotime
 var NetpollBreak = netpollBreak

src/runtime/string.go

@@ -351,14 +351,14 @@ func hasPrefix(s, prefix string) bool {
 }
 
 const (
-	maxUint = ^uint(0)
+	maxUint64 = ^uint64(0)
-	maxInt  = int(maxUint >> 1)
+	maxInt64  = int64(maxUint64 >> 1)
 )
 
-// atoi parses an int from a string s.
+// atoi64 parses an int64 from a string s.
 // The bool result reports whether s is a number
-// representable by a value of type int.
+// representable by a value of type int64.
-func atoi(s string) (int, bool) {
+func atoi64(s string) (int64, bool) {
 	if s == "" {
 		return 0, false
 	}
@@ -369,18 +369,18 @@ func atoi(s string) (int, bool) {
 		s = s[1:]
 	}
 
-	un := uint(0)
+	un := uint64(0)
 	for i := 0; i < len(s); i++ {
 		c := s[i]
 		if c < '0' || c > '9' {
 			return 0, false
 		}
-		if un > maxUint/10 {
+		if un > maxUint64/10 {
 			// overflow
 			return 0, false
 		}
 		un *= 10
-		un1 := un + uint(c) - '0'
+		un1 := un + uint64(c) - '0'
 		if un1 < un {
 			// overflow
 			return 0, false
@@ -388,14 +388,14 @@ func atoi(s string) (int, bool) {
 		un = un1
 	}
 
-	if !neg && un > uint(maxInt) {
+	if !neg && un > uint64(maxInt64) {
 		return 0, false
 	}
-	if neg && un > uint(maxInt)+1 {
+	if neg && un > uint64(maxInt64)+1 {
 		return 0, false
 	}
 
-	n := int(un)
+	n := int64(un)
 	if neg {
 		n = -n
 	}
@@ -403,15 +403,108 @@ func atoi(s string) (int, bool) {
 	return n, true
 }
 
+// atoi is like atoi64 but for integers
+// that fit into an int.
+func atoi(s string) (int, bool) {
+	if n, ok := atoi64(s); n == int64(int(n)) {
+		return int(n), ok
+	}
+	return 0, false
+}
+
 // atoi32 is like atoi but for integers
 // that fit into an int32.
 func atoi32(s string) (int32, bool) {
-	if n, ok := atoi(s); n == int(int32(n)) {
+	if n, ok := atoi64(s); n == int64(int32(n)) {
 		return int32(n), ok
 	}
 	return 0, false
 }
 
+// parseByteCount parses a string that represents a count of bytes.
+//
+// s must match the following regular expression:
+//
+//     ^[0-9]+(([KMGT]i)?B)?$
+//
+// In other words, an integer byte count with an optional unit
+// suffix. Acceptable suffixes include one of
+// - KiB, MiB, GiB, TiB which represent binary IEC/ISO 80000 units, or
+// - B, which just represents bytes.
+//
+// Returns an int64 because that's what its callers want and recieve,
+// but the result is always non-negative.
+func parseByteCount(s string) (int64, bool) {
+	// The empty string is not valid.
+	if s == "" {
+		return 0, false
+	}
+	// Handle the easy non-suffix case.
+	last := s[len(s)-1]
+	if last >= '0' && last <= '9' {
+		n, ok := atoi64(s)
+		if !ok || n < 0 {
+			return 0, false
+		}
+		return n, ok
+	}
+	// Failing a trailing digit, this must always end in 'B'.
+	// Also at this point there must be at least one digit before
+	// that B.
+	if last != 'B' || len(s) < 2 {
+		return 0, false
+	}
+	// The one before that must always be a digit or 'i'.
+	if c := s[len(s)-2]; c >= '0' && c <= '9' {
+		// Trivial 'B' suffix.
+		n, ok := atoi64(s[:len(s)-1])
+		if !ok || n < 0 {
+			return 0, false
+		}
+		return n, ok
+	} else if c != 'i' {
+		return 0, false
+	}
+	// Finally, we need at least 4 characters now, for the unit
+	// prefix and at least one digit.
+	if len(s) < 4 {
+		return 0, false
+	}
+	power := 0
+	switch s[len(s)-3] {
+	case 'K':
+		power = 1
+	case 'M':
+		power = 2
+	case 'G':
+		power = 3
+	case 'T':
+		power = 4
+	default:
+		// Invalid suffix.
+		return 0, false
+	}
+	m := uint64(1)
+	for i := 0; i < power; i++ {
+		m *= 1024
+	}
+	n, ok := atoi64(s[:len(s)-3])
+	if !ok || n < 0 {
+		return 0, false
+	}
+	un := uint64(n)
+	if un > maxUint64/m {
+		// Overflow.
+		return 0, false
+	}
+	un *= m
+	if un > uint64(maxInt64) {
+		// Overflow.
+		return 0, false
+	}
+	return int64(un), true
+}
+
 //go:nosplit
 func findnull(s *byte) int {
 	if s == nil {

src/runtime/string_test.go

@@ -454,3 +454,126 @@ func TestAtoi32(t *testing.T) {
 		}
 	}
 }
+
+func TestParseByteCount(t *testing.T) {
+	for _, test := range []struct {
+		in  string
+		out int64
+		ok  bool
+	}{
+		// Good numeric inputs.
+		{"1", 1, true},
+		{"12345", 12345, true},
+		{"012345", 12345, true},
+		{"98765432100", 98765432100, true},
+		{"9223372036854775807", 1<<63 - 1, true},
+
+		// Good trivial suffix inputs.
+		{"1B", 1, true},
+		{"12345B", 12345, true},
+		{"012345B", 12345, true},
+		{"98765432100B", 98765432100, true},
+		{"9223372036854775807B", 1<<63 - 1, true},
+
+		// Good binary suffix inputs.
+		{"1KiB", 1 << 10, true},
+		{"05KiB", 5 << 10, true},
+		{"1MiB", 1 << 20, true},
+		{"10MiB", 10 << 20, true},
+		{"1GiB", 1 << 30, true},
+		{"100GiB", 100 << 30, true},
+		{"1TiB", 1 << 40, true},
+		{"99TiB", 99 << 40, true},
+
+		// Good zero inputs.
+		//
+		// -0 is an edge case, but no harm in supporting it.
+		{"-0", 0, true},
+		{"0", 0, true},
+		{"0B", 0, true},
+		{"0KiB", 0, true},
+		{"0MiB", 0, true},
+		{"0GiB", 0, true},
+		{"0TiB", 0, true},
+
+		// Bad inputs.
+		{"", 0, false},
+		{"-1", 0, false},
+		{"a12345", 0, false},
+		{"a12345B", 0, false},
+		{"12345x", 0, false},
+		{"0x12345", 0, false},
+
+		// Bad numeric inputs.
+		{"9223372036854775808", 0, false},
+		{"9223372036854775809", 0, false},
+		{"18446744073709551615", 0, false},
+		{"20496382327982653440", 0, false},
+		{"18446744073709551616", 0, false},
+		{"18446744073709551617", 0, false},
+		{"9999999999999999999999", 0, false},
+
+		// Bad trivial suffix inputs.
+		{"9223372036854775808B", 0, false},
+		{"9223372036854775809B", 0, false},
+		{"18446744073709551615B", 0, false},
+		{"20496382327982653440B", 0, false},
+		{"18446744073709551616B", 0, false},
+		{"18446744073709551617B", 0, false},
+		{"9999999999999999999999B", 0, false},
+
+		// Bad binary suffix inputs.
+		{"1Ki", 0, false},
+		{"05Ki", 0, false},
+		{"10Mi", 0, false},
+		{"100Gi", 0, false},
+		{"99Ti", 0, false},
+		{"22iB", 0, false},
+		{"B", 0, false},
+		{"iB", 0, false},
+		{"KiB", 0, false},
+		{"MiB", 0, false},
+		{"GiB", 0, false},
+		{"TiB", 0, false},
+		{"-120KiB", 0, false},
+		{"-891MiB", 0, false},
+		{"-704GiB", 0, false},
+		{"-42TiB", 0, false},
+		{"99999999999999999999KiB", 0, false},
+		{"99999999999999999MiB", 0, false},
+		{"99999999999999GiB", 0, false},
+		{"99999999999TiB", 0, false},
+		{"555EiB", 0, false},
+
+		// Mistaken SI suffix inputs.
+		{"0KB", 0, false},
+		{"0MB", 0, false},
+		{"0GB", 0, false},
+		{"0TB", 0, false},
+		{"1KB", 0, false},
+		{"05KB", 0, false},
+		{"1MB", 0, false},
+		{"10MB", 0, false},
+		{"1GB", 0, false},
+		{"100GB", 0, false},
+		{"1TB", 0, false},
+		{"99TB", 0, false},
+		{"1K", 0, false},
+		{"05K", 0, false},
+		{"10M", 0, false},
+		{"100G", 0, false},
+		{"99T", 0, false},
+		{"99999999999999999999KB", 0, false},
+		{"99999999999999999MB", 0, false},
+		{"99999999999999GB", 0, false},
+		{"99999999999TB", 0, false},
+		{"99999999999TiB", 0, false},
+		{"555EB", 0, false},
+	} {
+		out, ok := runtime.ParseByteCount(test.in)
+		if test.out != out || test.ok != ok {
+			t.Errorf("parseByteCount(%q) = (%v, %v) want (%v, %v)",
+				test.in, out, ok, test.out, test.ok)
+		}
+	}
+}