From 12cc53d69938860daedfbce89d4d29923555270f Mon Sep 17 00:00:00 2001 From: fatedier Date: Sat, 9 Nov 2019 01:13:30 +0800 Subject: [PATCH] update bandwidth_limit --- client/proxy/proxy.go | 2 +- models/config/proxy.go | 8 +++---- models/config/types.go | 46 +++++++++++++++++++++++-------------- models/config/types_test.go | 40 ++++++++++++++++++++++++++++++++ 4 files changed, 74 insertions(+), 22 deletions(-) create mode 100644 models/config/types_test.go diff --git a/client/proxy/proxy.go b/client/proxy/proxy.go index aca1b94..268b317 100644 --- a/client/proxy/proxy.go +++ b/client/proxy/proxy.go @@ -54,7 +54,7 @@ type Proxy interface { func NewProxy(ctx context.Context, pxyConf config.ProxyConf, clientCfg config.ClientCommonConf, serverUDPPort int) (pxy Proxy) { var limiter *rate.Limiter - limitBytes := pxyConf.GetBaseInfo().BandwithLimit.Bytes() + limitBytes := pxyConf.GetBaseInfo().BandwidthLimit.Bytes() if limitBytes > 0 { limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes)) } diff --git a/models/config/proxy.go b/models/config/proxy.go index 2b129cf..85c353e 100644 --- a/models/config/proxy.go +++ b/models/config/proxy.go @@ -126,9 +126,9 @@ type BaseProxyConf struct { // version will be automatically selected. By default, this value is "". ProxyProtocolVersion string `json:"proxy_protocol_version"` - // BandwithLimit limit the proxy bandwith + // BandwidthLimit limit the proxy bandwidth // 0 means no limit - BandwithLimit BandwithQuantity `json:"bandwith_limit"` + BandwidthLimit BandwidthQuantity `json:"bandwidth_limit"` LocalSvrConf HealthCheckConf @@ -146,7 +146,7 @@ func (cfg *BaseProxyConf) compare(cmp *BaseProxyConf) bool { cfg.Group != cmp.Group || cfg.GroupKey != cmp.GroupKey || cfg.ProxyProtocolVersion != cmp.ProxyProtocolVersion || - cfg.BandwithLimit.Equal(&cmp.BandwithLimit) { + cfg.BandwidthLimit.Equal(&cmp.BandwidthLimit) { return false } if !cfg.LocalSvrConf.compare(&cmp.LocalSvrConf) { @@ -190,7 +190,7 @@ func (cfg *BaseProxyConf) UnmarshalFromIni(prefix string, name string, section i cfg.GroupKey = section["group_key"] cfg.ProxyProtocolVersion = section["proxy_protocol_version"] - if cfg.BandwithLimit, err = NewBandwithQuantity(section["bandwidth_limit"]); err != nil { + if cfg.BandwidthLimit, err = NewBandwidthQuantity(section["bandwidth_limit"]); err != nil { return err } diff --git a/models/config/types.go b/models/config/types.go index 9f3de66..87c240d 100644 --- a/models/config/types.go +++ b/models/config/types.go @@ -15,6 +15,7 @@ package config import ( + "encoding/json" "errors" "strconv" "strings" @@ -25,14 +26,14 @@ const ( KB = 1024 ) -type BandwithQuantity struct { +type BandwidthQuantity struct { s string // MB or KB i int64 // bytes } -func NewBandwithQuantity(s string) (BandwithQuantity, error) { - q := BandwithQuantity{} +func NewBandwidthQuantity(s string) (BandwidthQuantity, error) { + q := BandwidthQuantity{} err := q.UnmarshalString(s) if err != nil { return q, err @@ -40,7 +41,7 @@ func NewBandwithQuantity(s string) (BandwithQuantity, error) { return q, nil } -func (q *BandwithQuantity) Equal(u *BandwithQuantity) bool { +func (q *BandwidthQuantity) Equal(u *BandwidthQuantity) bool { if q == nil && u == nil { return true } @@ -50,13 +51,13 @@ func (q *BandwithQuantity) Equal(u *BandwithQuantity) bool { return false } -func (q *BandwithQuantity) String() string { +func (q *BandwidthQuantity) String() string { return q.s } -func (q *BandwithQuantity) UnmarshalString(s string) error { - q.s = strings.TrimSpace(s) - if q.s == "" { +func (q *BandwidthQuantity) UnmarshalString(s string) error { + s = strings.TrimSpace(s) + if s == "" { return nil } @@ -67,15 +68,15 @@ func (q *BandwithQuantity) UnmarshalString(s string) error { ) if strings.HasSuffix(s, "MB") { base = MB - s = strings.TrimSuffix(s, "MB") - f, err = strconv.ParseFloat(s, 64) + fstr := strings.TrimSuffix(s, "MB") + f, err = strconv.ParseFloat(fstr, 64) if err != nil { return err } } else if strings.HasSuffix(s, "KB") { base = KB - s = strings.TrimSuffix(s, "KB") - f, err = strconv.ParseFloat(s, 64) + fstr := strings.TrimSuffix(s, "KB") + f, err = strconv.ParseFloat(fstr, 64) if err != nil { return err } @@ -83,18 +84,29 @@ func (q *BandwithQuantity) UnmarshalString(s string) error { return errors.New("unit not support") } + q.s = s q.i = int64(f * float64(base)) return nil } -func (q *BandwithQuantity) UnmarshalJSON(b []byte) error { - return q.UnmarshalString(string(b)) +func (q *BandwidthQuantity) UnmarshalJSON(b []byte) error { + if len(b) == 4 && string(b) == "null" { + return nil + } + + var str string + err := json.Unmarshal(b, &str) + if err != nil { + return err + } + + return q.UnmarshalString(str) } -func (q *BandwithQuantity) MarshalJSON() ([]byte, error) { - return []byte(q.s), nil +func (q *BandwidthQuantity) MarshalJSON() ([]byte, error) { + return []byte("\"" + q.s + "\""), nil } -func (q *BandwithQuantity) Bytes() int64 { +func (q *BandwidthQuantity) Bytes() int64 { return q.i } diff --git a/models/config/types_test.go b/models/config/types_test.go new file mode 100644 index 0000000..ab03dfd --- /dev/null +++ b/models/config/types_test.go @@ -0,0 +1,40 @@ +// Copyright 2019 fatedier, fatedier@gmail.com +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package config + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" +) + +type Wrap struct { + B BandwidthQuantity `json:"b"` + Int int `json:"int"` +} + +func TestBandwidthQuantity(t *testing.T) { + assert := assert.New(t) + + var w Wrap + err := json.Unmarshal([]byte(`{"b":"1KB","int":5}`), &w) + assert.NoError(err) + assert.EqualValues(1*KB, w.B.Bytes()) + + buf, err := json.Marshal(&w) + assert.NoError(err) + assert.Equal(`{"b":"1KB","int":5}`, string(buf)) +}