sdk/go/arvados/byte_size.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvados
   6
   7 import (
   8         "encoding/json"
   9         "fmt"
  10         "math"
  11         "strconv"
  12         "strings"
  13 )
  14
  15 type ByteSize int64
  16
  17 // ByteSizeOrPercent indicates either a number of bytes or a
  18 // percentage from 1 to 100.
  19 type ByteSizeOrPercent ByteSize
  20
  21 var prefixValue = map[string]int64{
  22         "":   1,
  23         "K":  1000,
  24         "Ki": 1 << 10,
  25         "M":  1000000,
  26         "Mi": 1 << 20,
  27         "G":  1000000000,
  28         "Gi": 1 << 30,
  29         "T":  1000000000000,
  30         "Ti": 1 << 40,
  31         "P":  1000000000000000,
  32         "Pi": 1 << 50,
  33         "E":  1000000000000000000,
  34         "Ei": 1 << 60,
  35 }
  36
  37 func (n *ByteSize) UnmarshalJSON(data []byte) error {
  38         if len(data) == 0 || data[0] != '"' {
  39                 var i int64
  40                 err := json.Unmarshal(data, &i)
  41                 if err != nil {
  42                         return err
  43                 }
  44                 *n = ByteSize(i)
  45                 return nil
  46         }
  47         var s string
  48         err := json.Unmarshal(data, &s)
  49         if err != nil {
  50                 return err
  51         }
  52         split := strings.LastIndexAny(s, "0123456789.+-eE") + 1
  53         if split == 0 {
  54                 return fmt.Errorf("invalid byte size %q", s)
  55         }
  56         if s[split-1] == 'E' {
  57                 // We accepted an E as if it started the exponent part
  58                 // of a json number, but if the next char isn't +, -,
  59                 // or digit, then the E must have meant Exa. Instead
  60                 // of "4.5E"+"iB" we want "4.5"+"EiB".
  61                 split--
  62         }
  63         var val json.Number
  64         dec := json.NewDecoder(strings.NewReader(s[:split]))
  65         dec.UseNumber()
  66         err = dec.Decode(&val)
  67         if err != nil {
  68                 return err
  69         }
  70         if split == len(s) {
  71                 return nil
  72         }
  73         prefix := strings.Trim(s[split:], " ")
  74         if strings.HasSuffix(prefix, "B") {
  75                 prefix = prefix[:len(prefix)-1]
  76         }
  77         pval, ok := prefixValue[prefix]
  78         if !ok {
  79                 return fmt.Errorf("invalid unit %q", strings.Trim(s[split:], " "))
  80         }
  81         if intval, err := val.Int64(); err == nil {
  82                 if pval > 1 && (intval*pval)/pval != intval {
  83                         return fmt.Errorf("size %q overflows int64", s)
  84                 }
  85                 *n = ByteSize(intval * pval)
  86                 return nil
  87         } else if floatval, err := val.Float64(); err == nil {
  88                 if floatval*float64(pval) > math.MaxInt64 {
  89                         return fmt.Errorf("size %q overflows int64", s)
  90                 }
  91                 *n = ByteSize(int64(floatval * float64(pval)))
  92                 return nil
  93         } else {
  94                 return fmt.Errorf("bug: json.Number for %q is not int64 or float64: %s", s, err)
  95         }
  96 }
  97
  98 func (n ByteSizeOrPercent) MarshalJSON() ([]byte, error) {
  99         if n < 0 && n >= -100 {
 100                 return []byte(fmt.Sprintf("\"%d%%\"", -n)), nil
 101         } else {
 102                 return json.Marshal(int64(n))
 103         }
 104 }
 105
 106 func (n *ByteSizeOrPercent) UnmarshalJSON(data []byte) error {
 107         if len(data) == 0 || data[0] != '"' {
 108                 return (*ByteSize)(n).UnmarshalJSON(data)
 109         }
 110         var s string
 111         err := json.Unmarshal(data, &s)
 112         if err != nil {
 113                 return err
 114         }
 115         if s := strings.TrimSpace(s); len(s) > 0 && s[len(s)-1] == '%' {
 116                 pct, err := strconv.ParseInt(strings.TrimSpace(s[:len(s)-1]), 10, 64)
 117                 if err != nil {
 118                         return err
 119                 }
 120                 if pct < 0 || pct > 100 {
 121                         return fmt.Errorf("invalid value %q (percentage must be between 0 and 100)", s)
 122                 }
 123                 *n = ByteSizeOrPercent(-pct)
 124                 return nil
 125         }
 126         return (*ByteSize)(n).UnmarshalJSON(data)
 127 }
 128
 129 // ByteSize returns the absolute byte size specified by n, or 0 if n
 130 // specifies a percent.
 131 func (n ByteSizeOrPercent) ByteSize() ByteSize {
 132         if n >= -100 && n < 0 {
 133                 return 0
 134         } else {
 135                 return ByteSize(n)
 136         }
 137 }
 138
 139 // ByteSize returns the percentage specified by n, or 0 if n specifies
 140 // an absolute byte size.
 141 func (n ByteSizeOrPercent) Percent() int64 {
 142         if n >= -100 && n < 0 {
 143                 return int64(-n)
 144         } else {
 145                 return 0
 146         }
 147 }