From 0feebe6eb594e28c29022c65ffcd2ae0e8bebac7 Mon Sep 17 00:00:00 2001 From: Michael Anthony Knyszek Date: Tue, 15 Feb 2022 00:22:20 +0000 Subject: [PATCH] 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 Reviewed-by: Michael Pratt Run-TryBot: Michael Knyszek --- src/runtime/export_test.go | 1 + src/runtime/string.go | 117 +++++++++++++++++++++++++++++++---- src/runtime/string_test.go | 123 +++++++++++++++++++++++++++++++++++++ 3 files changed, 229 insertions(+), 12 deletions(-) diff --git a/src/runtime/export_test.go b/src/runtime/export_test.go index 708da264b72..c364e5bea98 100644 --- a/src/runtime/export_test.go +++ b/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 diff --git a/src/runtime/string.go b/src/runtime/string.go index 8b20c93fd78..845dcb50c49 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 { diff --git a/src/runtime/string_test.go b/src/runtime/string_test.go index 4eda12c35dd..1ea7f5e481c 100644 --- a/src/runtime/string_test.go +++ b/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) + } + } +}