Browse Source

update bandwidth_limit

fatedier 5 years ago
parent
commit
12cc53d699
4 changed files with 74 additions and 22 deletions
  1. 1 1
      client/proxy/proxy.go
  2. 4 4
      models/config/proxy.go
  3. 29 17
      models/config/types.go
  4. 40 0
      models/config/types_test.go

+ 1 - 1
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))
 	}

+ 4 - 4
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
 	}
 

+ 29 - 17
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
 }

+ 40 - 0
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))
+}