1package atc 2 3import ( 4 "encoding/json" 5 "errors" 6 "math" 7 "regexp" 8 "strconv" 9 "strings" 10) 11 12const MemoryRegex = "^([0-9]+)([G|M|K|g|m|k]?[b|B])?$" 13 14func (c *ContainerLimits) UnmarshalJSON(limit []byte) error { 15 var data interface{} 16 17 err := json.Unmarshal(limit, &data) 18 if err != nil { 19 return err 20 } 21 22 climits, err := ParseContainerLimits(data) 23 if err != nil { 24 return err 25 } 26 27 c.CPU = climits.CPU 28 c.Memory = climits.Memory 29 30 return nil 31} 32 33func ParseContainerLimits(data interface{}) (ContainerLimits, error) { 34 mapData, ok := data.(map[string]interface{}) 35 if !ok { 36 mapData = make(map[string]interface{}) 37 } 38 39 var c ContainerLimits 40 41 var memoryBytes uint64 42 var uVal int 43 var err error 44 45 // the json unmarshaller returns numbers as float64 while yaml returns int 46 for key, val := range mapData { 47 if key == "memory" { 48 switch val.(type) { 49 case string: 50 memoryBytes, err = parseMemoryLimit(val.(string)) 51 if err != nil { 52 return ContainerLimits{}, err 53 } 54 case *string: 55 if val.(*string) == nil { 56 c.Memory = nil 57 continue 58 } 59 memoryBytes, err = parseMemoryLimit(*val.(*string)) 60 if err != nil { 61 return ContainerLimits{}, err 62 } 63 case float64: 64 memoryBytes = uint64(int(val.(float64))) 65 case int: 66 memoryBytes = uint64(val.(int)) 67 } 68 c.Memory = &memoryBytes 69 70 } else if key == "cpu" { 71 switch val.(type) { 72 case float64: 73 uVal = int(val.(float64)) 74 case int: 75 uVal = val.(int) 76 case *int: 77 if val.(*int) == nil { 78 c.CPU = nil 79 continue 80 } 81 uVal = *val.(*int) 82 default: 83 return ContainerLimits{}, errors.New("cpu limit must be an integer") 84 } 85 helper := uint64(uVal) 86 c.CPU = &helper 87 88 } 89 } 90 91 return c, nil 92} 93 94func parseMemoryLimit(limit string) (uint64, error) { 95 limit = strings.ToUpper(limit) 96 var sizeRegex *regexp.Regexp = regexp.MustCompile(MemoryRegex) 97 matches := sizeRegex.FindStringSubmatch(limit) 98 99 if len(matches) > 3 || len(matches) < 1 { 100 return 0, errors.New("could not parse container memory limit") 101 } 102 103 value, err := strconv.ParseUint(matches[1], 10, 64) 104 if err != nil { 105 return 0, err 106 } 107 108 var power float64 109 var base float64 = 2 110 var unit string = matches[2] 111 switch unit { 112 case "KB": 113 power = 10 114 case "MB": 115 power = 20 116 case "GB": 117 power = 30 118 default: 119 power = 0 120 } 121 122 return value * uint64(math.Pow(base, power)), nil 123} 124