update bandwidth_limit

This commit is contained in:
fatedier 2019-11-09 01:13:30 +08:00
parent 2ab832bb89
commit 12cc53d699
4 changed files with 74 additions and 22 deletions

View File

@ -54,7 +54,7 @@ type Proxy interface {
func NewProxy(ctx context.Context, pxyConf config.ProxyConf, clientCfg config.ClientCommonConf, serverUDPPort int) (pxy Proxy) { func NewProxy(ctx context.Context, pxyConf config.ProxyConf, clientCfg config.ClientCommonConf, serverUDPPort int) (pxy Proxy) {
var limiter *rate.Limiter var limiter *rate.Limiter
limitBytes := pxyConf.GetBaseInfo().BandwithLimit.Bytes() limitBytes := pxyConf.GetBaseInfo().BandwidthLimit.Bytes()
if limitBytes > 0 { if limitBytes > 0 {
limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes)) limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
} }

View File

@ -126,9 +126,9 @@ type BaseProxyConf struct {
// version will be automatically selected. By default, this value is "". // version will be automatically selected. By default, this value is "".
ProxyProtocolVersion string `json:"proxy_protocol_version"` ProxyProtocolVersion string `json:"proxy_protocol_version"`
// BandwithLimit limit the proxy bandwith // BandwidthLimit limit the proxy bandwidth
// 0 means no limit // 0 means no limit
BandwithLimit BandwithQuantity `json:"bandwith_limit"` BandwidthLimit BandwidthQuantity `json:"bandwidth_limit"`
LocalSvrConf LocalSvrConf
HealthCheckConf HealthCheckConf
@ -146,7 +146,7 @@ func (cfg *BaseProxyConf) compare(cmp *BaseProxyConf) bool {
cfg.Group != cmp.Group || cfg.Group != cmp.Group ||
cfg.GroupKey != cmp.GroupKey || cfg.GroupKey != cmp.GroupKey ||
cfg.ProxyProtocolVersion != cmp.ProxyProtocolVersion || cfg.ProxyProtocolVersion != cmp.ProxyProtocolVersion ||
cfg.BandwithLimit.Equal(&cmp.BandwithLimit) { cfg.BandwidthLimit.Equal(&cmp.BandwidthLimit) {
return false return false
} }
if !cfg.LocalSvrConf.compare(&cmp.LocalSvrConf) { 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.GroupKey = section["group_key"]
cfg.ProxyProtocolVersion = section["proxy_protocol_version"] 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 return err
} }

View File

@ -15,6 +15,7 @@
package config package config
import ( import (
"encoding/json"
"errors" "errors"
"strconv" "strconv"
"strings" "strings"
@ -25,14 +26,14 @@ const (
KB = 1024 KB = 1024
) )
type BandwithQuantity struct { type BandwidthQuantity struct {
s string // MB or KB s string // MB or KB
i int64 // bytes i int64 // bytes
} }
func NewBandwithQuantity(s string) (BandwithQuantity, error) { func NewBandwidthQuantity(s string) (BandwidthQuantity, error) {
q := BandwithQuantity{} q := BandwidthQuantity{}
err := q.UnmarshalString(s) err := q.UnmarshalString(s)
if err != nil { if err != nil {
return q, err return q, err
@ -40,7 +41,7 @@ func NewBandwithQuantity(s string) (BandwithQuantity, error) {
return q, nil return q, nil
} }
func (q *BandwithQuantity) Equal(u *BandwithQuantity) bool { func (q *BandwidthQuantity) Equal(u *BandwidthQuantity) bool {
if q == nil && u == nil { if q == nil && u == nil {
return true return true
} }
@ -50,13 +51,13 @@ func (q *BandwithQuantity) Equal(u *BandwithQuantity) bool {
return false return false
} }
func (q *BandwithQuantity) String() string { func (q *BandwidthQuantity) String() string {
return q.s return q.s
} }
func (q *BandwithQuantity) UnmarshalString(s string) error { func (q *BandwidthQuantity) UnmarshalString(s string) error {
q.s = strings.TrimSpace(s) s = strings.TrimSpace(s)
if q.s == "" { if s == "" {
return nil return nil
} }
@ -67,15 +68,15 @@ func (q *BandwithQuantity) UnmarshalString(s string) error {
) )
if strings.HasSuffix(s, "MB") { if strings.HasSuffix(s, "MB") {
base = MB base = MB
s = strings.TrimSuffix(s, "MB") fstr := strings.TrimSuffix(s, "MB")
f, err = strconv.ParseFloat(s, 64) f, err = strconv.ParseFloat(fstr, 64)
if err != nil { if err != nil {
return err return err
} }
} else if strings.HasSuffix(s, "KB") { } else if strings.HasSuffix(s, "KB") {
base = KB base = KB
s = strings.TrimSuffix(s, "KB") fstr := strings.TrimSuffix(s, "KB")
f, err = strconv.ParseFloat(s, 64) f, err = strconv.ParseFloat(fstr, 64)
if err != nil { if err != nil {
return err return err
} }
@ -83,18 +84,29 @@ func (q *BandwithQuantity) UnmarshalString(s string) error {
return errors.New("unit not support") return errors.New("unit not support")
} }
q.s = s
q.i = int64(f * float64(base)) q.i = int64(f * float64(base))
return nil return nil
} }
func (q *BandwithQuantity) UnmarshalJSON(b []byte) error { func (q *BandwidthQuantity) UnmarshalJSON(b []byte) error {
return q.UnmarshalString(string(b)) 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) { func (q *BandwidthQuantity) MarshalJSON() ([]byte, error) {
return []byte(q.s), nil return []byte("\"" + q.s + "\""), nil
} }
func (q *BandwithQuantity) Bytes() int64 { func (q *BandwidthQuantity) Bytes() int64 {
return q.i return q.i
} }

View File

@ -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))
}