From a513663a7254efa51efda567506c357b1e6c534d Mon Sep 17 00:00:00 2001 From: Qian Wang Date: Sun, 4 Aug 2019 17:16:18 +0100 Subject: [PATCH] Use a fork of ratelimit for unlimited bucket --- go.mod | 6 ++++-- go.sum | 28 -------------------------- internal/multiplex/qos.go | 41 ++++++++++++++++++++++++--------------- 3 files changed, 29 insertions(+), 46 deletions(-) delete mode 100644 go.sum diff --git a/go.mod b/go.mod index 4d3ca37..49e4444 100644 --- a/go.mod +++ b/go.mod @@ -1,11 +1,13 @@ module github.com/cbeuw/Cloak +go 1.12 + require ( github.com/boltdb/bolt v1.3.1 + github.com/cbeuw/ratelimit v1.1.0 github.com/gorilla/mux v1.7.3 - github.com/juju/ratelimit v1.0.1 github.com/kr/pretty v0.1.0 // indirect github.com/sirupsen/logrus v1.4.2 - golang.org/x/crypto v0.0.0-20190123085648-057139ce5d2b + golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect ) diff --git a/go.sum b/go.sum deleted file mode 100644 index a3e3b7a..0000000 --- a/go.sum +++ /dev/null @@ -1,28 +0,0 @@ -github.com/boltdb/bolt v1.3.1 h1:JQmyP4ZBrce+ZQu0dY660FMfatumYDLun9hBCUVIkF4= -github.com/boltdb/bolt v1.3.1/go.mod h1:clJnj/oiGkjum5o1McbSZDSLxVThjynRyGBgiAx27Ps= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= -github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/gorilla/mux v1.7.3 h1:gnP5JzjVOuiZD07fKKToCAOjS0yOpj/qPETTXCCS6hw= -github.com/gorilla/mux v1.7.3/go.mod h1:1lud6UwP+6orDFRuTfBEV8e9/aOM/c4fVVCaMa2zaAs= -github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY= -github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk= -github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk= -github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= -github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= -github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= -github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= -github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= -github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= -github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/sirupsen/logrus v1.4.2 h1:SPIRibHv4MatM3XXNO2BJeFLZwZ2LvZgfQ5+UNI2im4= -github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE= -github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= -github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= -golang.org/x/crypto v0.0.0-20190123085648-057139ce5d2b h1:Elez2XeF2p9uyVj0yEUDqQ56NFcDtcBNkYP7yv8YbUE= -golang.org/x/crypto v0.0.0-20190123085648-057139ce5d2b/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= -golang.org/x/sys v0.0.0-20190422165155-953cdadca894 h1:Cz4ceDQGXuKRnVBDTS23GTn/pU5OE2C0WrNTOYK1Uuc= -golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= -gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/multiplex/qos.go b/internal/multiplex/qos.go index f01dd39..6a96f78 100644 --- a/internal/multiplex/qos.go +++ b/internal/multiplex/qos.go @@ -3,7 +3,7 @@ package multiplex import ( "sync/atomic" - "github.com/juju/ratelimit" + "github.com/cbeuw/ratelimit" ) // Valve needs to be universal, across all sessions that belong to a user @@ -14,8 +14,8 @@ type Valve struct { // rx is from client to server, tx is from server to client // DO NOT use terms up or down as this is used in usermanager // for bandwidth limiting - rxtb atomic.Value // *ratelimit.Bucket - txtb atomic.Value // *ratelimit.Bucket + rxtb ratelimit.Bucket + txtb ratelimit.Bucket rx *int64 tx *int64 @@ -24,24 +24,33 @@ type Valve struct { func MakeValve(rxRate, txRate int64) *Valve { var rx, tx int64 v := &Valve{ - rx: &rx, - tx: &tx, + rxtb: ratelimit.NewLimitedBucketWithRate(float64(rxRate), rxRate), + txtb: ratelimit.NewLimitedBucketWithRate(float64(txRate), txRate), + rx: &rx, + tx: &tx, } - v.SetRxRate(rxRate) - v.SetTxRate(txRate) return v } -var UNLIMITED_VALVE = MakeValve(1<<63-1, 1<<63-1) +func MakeUnlimitedValve() *Valve { + var rx, tx int64 + v := &Valve{ + rxtb: ratelimit.NewUnlimitedBucket(), + txtb: ratelimit.NewUnlimitedBucket(), + rx: &rx, + tx: &tx, + } + return v +} -func (v *Valve) SetRxRate(rate int64) { v.rxtb.Store(ratelimit.NewBucketWithRate(float64(rate), rate)) } -func (v *Valve) SetTxRate(rate int64) { v.txtb.Store(ratelimit.NewBucketWithRate(float64(rate), rate)) } -func (v *Valve) rxWait(n int) { v.rxtb.Load().(*ratelimit.Bucket).Wait(int64(n)) } -func (v *Valve) txWait(n int) { v.txtb.Load().(*ratelimit.Bucket).Wait(int64(n)) } -func (v *Valve) AddRx(n int64) { atomic.AddInt64(v.rx, n) } -func (v *Valve) AddTx(n int64) { atomic.AddInt64(v.tx, n) } -func (v *Valve) GetRx() int64 { return atomic.LoadInt64(v.rx) } -func (v *Valve) GetTx() int64 { return atomic.LoadInt64(v.tx) } +var UNLIMITED_VALVE = MakeUnlimitedValve() + +func (v *Valve) rxWait(n int) { v.rxtb.Wait(int64(n)) } +func (v *Valve) txWait(n int) { v.txtb.Wait(int64(n)) } +func (v *Valve) AddRx(n int64) { atomic.AddInt64(v.rx, n) } +func (v *Valve) AddTx(n int64) { atomic.AddInt64(v.tx, n) } +func (v *Valve) GetRx() int64 { return atomic.LoadInt64(v.rx) } +func (v *Valve) GetTx() int64 { return atomic.LoadInt64(v.tx) } func (v *Valve) Nullify() (int64, int64) { rx := atomic.SwapInt64(v.rx, 0) tx := atomic.SwapInt64(v.tx, 0)