Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package util
18

19
import (
20
	"encoding/json"
21
	"errors"
22
	"fmt"
23
	"time"
24

25
	"github.com/docker/go-units"
26
	"golang.org/x/time/rate"
27
	"gopkg.in/yaml.v3"
28
)
29

30
// RateLimit is a wrapper for rate.Limit, support json and yaml unmarshal function
31
// yaml example 1:
32
// rate_limit: 2097152 # 2MiB
33
// yaml example 2:
34
// rate_limit: 2MiB
35
type RateLimit struct {
36
	rate.Limit
37
}
38

39
type Size = RateLimit
40

41
func (r *RateLimit) UnmarshalJSON(b []byte) error {
42
	return r.unmarshal(json.Unmarshal, b)
43
}
44

45
func (r *RateLimit) UnmarshalYAML(node *yaml.Node) error {
46
	return r.unmarshal(yaml.Unmarshal, []byte(node.Value))
47
}
48

49
func (r *RateLimit) unmarshal(unmarshal func(in []byte, out any) (err error), b []byte) error {
50
	var v any
51
	if err := unmarshal(b, &v); err != nil {
52
		return err
53
	}
54
	switch value := v.(type) {
55
	case float64:
56
		r.Limit = rate.Limit(value)
57
		return nil
58
	case int:
59
		r.Limit = rate.Limit(value)
60
		return nil
61
	case string:
62
		limit, err := units.RAMInBytes(value)
63
		if err != nil {
64
			return fmt.Errorf("invalid rate limit: %w", err)
65
		}
66
		r.Limit = rate.Limit(limit)
67
		return nil
68
	default:
69
		return errors.New("invalid rate limit")
70
	}
71
}
72

73
type Duration struct {
74
	time.Duration
75
}
76

77
func (d *Duration) UnmarshalJSON(b []byte) error {
78
	var v any
79
	if err := json.Unmarshal(b, &v); err != nil {
80
		return err
81
	}
82
	return d.unmarshal(v)
83
}
84

85
func (d *Duration) UnmarshalYAML(node *yaml.Node) error {
86
	var v any
87
	switch node.Kind {
88
	case yaml.ScalarNode:
89
		switch node.Tag {
90
		case "!!int":
91
			var i int
92
			if err := node.Decode(&i); err != nil {
93
				return err
94
			}
95
			v = i
96
		case "!!str":
97
			var i string
98
			if err := node.Decode(&i); err != nil {
99
				return err
100
			}
101
			v = i
102
		default:
103
			return errors.New("invalid duration")
104
		}
105
	default:
106
		return errors.New("invalid duration")
107
	}
108
	return d.unmarshal(v)
109
}
110

111
func (d *Duration) unmarshal(v any) error {
112
	switch value := v.(type) {
113
	case float64:
114
		d.Duration = time.Duration(value)
115
		return nil
116
	case int:
117
		d.Duration = time.Duration(value)
118
		return nil
119
	case string:
120
		var err error
121
		d.Duration, err = time.ParseDuration(value)
122
		if err != nil {
123
			return err
124
		}
125
		return nil
126
	default:
127
		return errors.New("invalid duration")
128
	}
129
}
130