1// Copyright (C) 2019 Storj Labs, Inc. 2// See LICENSE for copying information. 3 4package memory 5 6import ( 7 "errors" 8 "fmt" 9 "strconv" 10 "strings" 11) 12 13// base 2 and base 10 sizes. 14const ( 15 B Size = 1 << (10 * iota) 16 KiB 17 MiB 18 GiB 19 TiB 20 PiB 21 EiB 22 23 KB Size = 1e3 24 MB Size = 1e6 25 GB Size = 1e9 26 TB Size = 1e12 27 PB Size = 1e15 28 EB Size = 1e18 29) 30 31// Size implements flag.Value for collecting memory size in bytes. 32type Size int64 33 34// Int returns bytes size as int. 35func (size Size) Int() int { return int(size) } 36 37// Int32 returns bytes size as int32. 38func (size Size) Int32() int32 { return int32(size) } 39 40// Int64 returns bytes size as int64. 41func (size Size) Int64() int64 { return int64(size) } 42 43// Float64 returns bytes size as float64. 44func (size Size) Float64() float64 { return float64(size) } 45 46// KiB returns size in kibibytes. 47func (size Size) KiB() float64 { return size.Float64() / KiB.Float64() } 48 49// MiB returns size in mebibytes. 50func (size Size) MiB() float64 { return size.Float64() / MiB.Float64() } 51 52// GiB returns size in gibibytes. 53func (size Size) GiB() float64 { return size.Float64() / GiB.Float64() } 54 55// TiB returns size in tebibytes. 56func (size Size) TiB() float64 { return size.Float64() / TiB.Float64() } 57 58// PiB returns size in pebibytes. 59func (size Size) PiB() float64 { return size.Float64() / PiB.Float64() } 60 61// EiB returns size in exbibytes. 62func (size Size) EiB() float64 { return size.Float64() / EiB.Float64() } 63 64// KB returns size in kilobytes. 65func (size Size) KB() float64 { return size.Float64() / KB.Float64() } 66 67// MB returns size in megabytes. 68func (size Size) MB() float64 { return size.Float64() / MB.Float64() } 69 70// GB returns size in gigabytes. 71func (size Size) GB() float64 { return size.Float64() / GB.Float64() } 72 73// TB returns size in terabytes. 74func (size Size) TB() float64 { return size.Float64() / TB.Float64() } 75 76// PB returns size in petabytes. 77func (size Size) PB() float64 { return size.Float64() / PB.Float64() } 78 79// EB returns size in exabytes. 80func (size Size) EB() float64 { return size.Float64() / EB.Float64() } 81 82// String converts size to a string using base-2 prefixes, unless the number 83// appears to be in base 10. 84func (size Size) String() string { 85 if countZeros(int64(size), 1000) > countZeros(int64(size), 1024) { 86 return size.Base10String() 87 } 88 return size.Base2String() 89} 90 91// countZeros considers a number num in base base. It counts zeros of that 92// number in that base from least significant to most, stopping when a non-zero 93// value is hit. 94func countZeros(num, base int64) (count int) { 95 for num != 0 && num%base == 0 { 96 num /= base 97 count++ 98 } 99 return count 100} 101 102// Base2String converts size to a string using base-2 prefixes. 103func (size Size) Base2String() string { 104 if size == 0 { 105 return "0 B" 106 } 107 108 switch { 109 case abs(size) >= EiB*2/3: 110 return fmt.Sprintf("%.1f EiB", size.EiB()) 111 case abs(size) >= PiB*2/3: 112 return fmt.Sprintf("%.1f PiB", size.PiB()) 113 case abs(size) >= TiB*2/3: 114 return fmt.Sprintf("%.1f TiB", size.TiB()) 115 case abs(size) >= GiB*2/3: 116 return fmt.Sprintf("%.1f GiB", size.GiB()) 117 case abs(size) >= MiB*2/3: 118 return fmt.Sprintf("%.1f MiB", size.MiB()) 119 case abs(size) >= KiB*2/3: 120 return fmt.Sprintf("%.1f KiB", size.KiB()) 121 } 122 123 return strconv.FormatInt(size.Int64(), 10) + " B" 124} 125 126// Base10String converts size to a string using base-10 prefixes. 127func (size Size) Base10String() string { 128 if size == 0 { 129 return "0 B" 130 } 131 132 switch { 133 case abs(size) >= EB*2/3: 134 return fmt.Sprintf("%.2f EB", size.EB()) 135 case abs(size) >= PB*2/3: 136 return fmt.Sprintf("%.2f PB", size.PB()) 137 case abs(size) >= TB*2/3: 138 return fmt.Sprintf("%.2f TB", size.TB()) 139 case abs(size) >= GB*2/3: 140 return fmt.Sprintf("%.2f GB", size.GB()) 141 case abs(size) >= MB*2/3: 142 return fmt.Sprintf("%.2f MB", size.MB()) 143 case abs(size) >= KB*2/3: 144 return fmt.Sprintf("%.2f KB", size.KB()) 145 } 146 147 return strconv.FormatInt(size.Int64(), 10) + " B" 148} 149 150func abs(size Size) Size { 151 if size > 0 { 152 return size 153 } 154 return -size 155} 156 157func isLetter(b byte) bool { 158 return ('a' <= b && b <= 'z') || ('A' <= b && b <= 'Z') 159} 160 161// Set updates value from string. 162func (size *Size) Set(s string) error { 163 if s == "" { 164 return errors.New("empty size") 165 } 166 167 p := len(s) 168 for isLetter(s[p-1]) { 169 p-- 170 171 if p < 0 { 172 return errors.New("p out of bounds") 173 } 174 } 175 176 value, suffix := s[:p], s[p:] 177 suffix = strings.ToUpper(suffix) 178 if suffix == "" || suffix[len(suffix)-1] != 'B' { 179 suffix += "B" 180 } 181 182 value = strings.TrimSpace(value) 183 v, err := strconv.ParseFloat(value, 64) 184 if err != nil { 185 return err 186 } 187 188 switch suffix { 189 case "EB": 190 *size = Size(v * EB.Float64()) 191 case "EIB": 192 *size = Size(v * EiB.Float64()) 193 case "PB": 194 *size = Size(v * PB.Float64()) 195 case "PIB": 196 *size = Size(v * PiB.Float64()) 197 case "TB": 198 *size = Size(v * TB.Float64()) 199 case "TIB": 200 *size = Size(v * TiB.Float64()) 201 case "GB": 202 *size = Size(v * GB.Float64()) 203 case "GIB": 204 *size = Size(v * GiB.Float64()) 205 case "MB": 206 *size = Size(v * MB.Float64()) 207 case "MIB": 208 *size = Size(v * MiB.Float64()) 209 case "KB": 210 *size = Size(v * KB.Float64()) 211 case "KIB": 212 *size = Size(v * KiB.Float64()) 213 case "B", "": 214 *size = Size(v) 215 default: 216 return fmt.Errorf("unknown suffix %q", suffix) 217 } 218 219 return nil 220} 221 222// Type implements pflag.Value. 223func (Size) Type() string { return "memory.Size" } 224 225// MarshalText returns size as a string. 226func (size Size) MarshalText() (string, error) { 227 return size.String(), nil 228} 229 230// UnmarshalText parses text as a string. 231func (size *Size) UnmarshalText(text []byte) error { 232 return size.Set(string(text)) 233} 234 235// MarshalJSON returns size as a json string. 236func (size Size) MarshalJSON() ([]byte, error) { 237 return []byte(strconv.Quote(size.String())), nil 238} 239 240// UnmarshalJSON parses text from a json string. 241func (size *Size) UnmarshalJSON(text []byte) error { 242 unquoted, err := strconv.Unquote(string(text)) 243 if err != nil { 244 return err 245 } 246 return size.Set(unquoted) 247} 248