Use a fork of ratelimit for unlimited bucket

pull/71/head
Qian Wang 5 years ago
parent 640f81461c
commit a513663a72

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

@ -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=

@ -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
}
var UNLIMITED_VALVE = MakeUnlimitedValve()
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) }
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)

Loading…
Cancel
Save