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>

1 files changed, 105 insertions, 12 deletions

diff --git a/src/runtime/string.go b/src/runtime/string.go
index 8b20c93fd7..845dcb50c4 100644
--- a/src/runtime/string.go
+++ b/src/runtime/string.go
@@ -351,14 +351,14 @@ func hasPrefix(s, prefix string) bool {
 }
 
 const (
-	maxUint = ^uint(0)
-	maxInt  = int(maxUint >> 1)
+	maxUint64 = ^uint64(0)
+	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.
-func atoi(s string) (int, bool) {
+// representable by a value of type int64.
+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 {