diff --git a/cuttle.yml b/cuttle.yml index 57c1574..58b7c51 100644 --- a/cuttle.yml +++ b/cuttle.yml @@ -14,8 +14,9 @@ addr: :3128 zones: - host: "*" # Apply to requests forwarded to all domains. shared: true # The rate limit is shared by all domains. - control: rps # Use request-per-second rate limit control. - rate: 2 # At most 2 requests per second in the entire zone. + control: rpns # Use request-per-second rate limit control. + rate: 1 # At most 1 requests per N second in the entire zone. + nseconds: 3 # N second = 3 # Example - Only throttle *.github.com and no rate limit on other domains. # diff --git a/cuttle/limitcontrol.go b/cuttle/limitcontrol.go index c0c4a57..b0dd4eb 100644 --- a/cuttle/limitcontrol.go +++ b/cuttle/limitcontrol.go @@ -186,3 +186,66 @@ func (c *RPMControl) Acquire() bool { return true } + +// RPNSControl provides Rate requests per N seconds rate limit control. +type RPNSControl struct { + // Label of this control. + Label string + // Rate holds the number of requests per N seconds. + Rate int + // Nseconds holds the number N of seconds within which limiting occurs + Nseconds int + + pendingChan chan uint + readyChan chan uint + seen *list.List +} + +// NewRPNSControl return a new RPNSControl with the given label and rate. +func NewRPNSControl(label string, rate int, nseconds int) *RPNSControl { + return &RPNSControl{label, rate, nseconds, make(chan uint), make(chan uint), list.New()} +} + +// Start running RPNSControl. +// A goroutine is launched to govern the rate limit of Acquire(). +func (c *RPNSControl) Start() { + go func() { + log.Debugf("RPNSControl[%s]: Activated.", c.Label) + + for { + <-c.pendingChan + + log.Debugf("RPNSControl[%s]: Limited at %dreq/%ds.", c.Label, c.Rate, c.Nseconds) + if c.seen.Len() == c.Rate { + front := c.seen.Front() + nanoElapsed := time.Now().UnixNano() - front.Value.(int64) + milliElapsed := nanoElapsed / int64(time.Millisecond) + secondElapsed := milliElapsed / 1000 + log.Debugf("RPNSControl[%s]: Elapsed %ds since first request.", c.Label, secondElapsed) + + if waitTime := int64(c.Nseconds) - secondElapsed; waitTime > 0 { + log.Infof("RPNSControl[%s]: Waiting for %ds.", c.Label, waitTime) + time.Sleep(time.Duration(waitTime) * time.Second) + } + + c.seen.Remove(front) + } + c.seen.PushBack(time.Now().UnixNano()) + + c.readyChan <- 1 + } + + log.Debugf("RPNSControl[%s]: Deactivated.", c.Label) + }() +} + +// Acquire permission from RPNSControl. +// Permission is granted at a rate of Rate requests per N seconds. +func (c *RPNSControl) Acquire() bool { + log.Debugf("RPNSControl[%s]: Seeking permission.", c.Label) + c.pendingChan <- 1 + <-c.readyChan + log.Debugf("RPNSControl[%s]: Granted permission.", c.Label) + + return true +} diff --git a/cuttle/zone.go b/cuttle/zone.go index 9c2c77c..a8d5cdc 100644 --- a/cuttle/zone.go +++ b/cuttle/zone.go @@ -29,14 +29,14 @@ type Zone struct { Control string // Rate specifies the rate of the rate limit controller. Rate int - + Nseconds int controllers map[string]LimitController } // NewZone returns a new Zone given the configurations. -func NewZone(host string, path string, limitby string, shared bool, control string, rate int) *Zone { +func NewZone(host string, path string, limitby string, shared bool, control string, rate int, nseconds int) *Zone { return &Zone{ - host, path, limitby, shared, control, rate, + host, path, limitby, shared, control, rate, nseconds, make(map[string]LimitController), } } @@ -102,6 +102,8 @@ func (z *Zone) GetController(host string, path string) LimitController { controller = NewRPSControl(key, z.Rate) case "rpm": controller = NewRPMControl(key, z.Rate) + case "rpns": + controller = NewRPNSControl(key, z.Rate, z.Nseconds) case "noop": controller = NewNoopControl(key) case "ban": diff --git a/main.go b/main.go index 9d7d459..aef13be 100644 --- a/main.go +++ b/main.go @@ -55,11 +55,15 @@ func main() { if c.LimitBy == "" { c.LimitBy = "host" } + + if c.Nseconds == 0 { + c.Nseconds = 1 + } - log.Debugf("ZoneConfig: host - %s, path - %s, limitby - %s, shared - %t, control - %s, rate - %d", - c.Host, c.Path, c.LimitBy, c.Shared, c.Control, c.Rate) + log.Debugf("ZoneConfig: host - %s, path - %s, limitby - %s, shared - %t, control - %s, rate - %d, nseconds - %d", + c.Host, c.Path, c.LimitBy, c.Shared, c.Control, c.Rate, c.Nseconds) - zones[i] = *cuttle.NewZone(c.Host, c.Path, c.LimitBy, c.Shared, c.Control, c.Rate) + zones[i] = *cuttle.NewZone(c.Host, c.Path, c.LimitBy, c.Shared, c.Control, c.Rate, c.Nseconds) } // Config CA Cert. @@ -130,10 +134,11 @@ type Config struct { } type ZoneConfig struct { - Host string - Path string // Optional, default "/" - LimitBy string // Optional, default "host" - Shared bool // Optional, default "false" - Control string - Rate int + Host string + Path string // Optional, default "/" + LimitBy string // Optional, default "host" + Shared bool // Optional, default "false" + Control string + Rate int + Nseconds int // Optional, default "1" }