1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798 |
- // Package ratelimiter implements the Leaky Bucket ratelimiting algorithm with memcached and in-memory backends.
- package ratelimiter
- import (
- "time"
- )
- type LeakyBucket struct {
- Size uint16
- Fill float64
- LeakInterval time.Duration // time.Duration for 1 unit of size to leak
- Lastupdate time.Time
- Now func() time.Time
- }
- func NewLeakyBucket(size uint16, leakInterval time.Duration) *LeakyBucket {
- bucket := LeakyBucket{
- Size: size,
- Fill: 0,
- LeakInterval: leakInterval,
- Now: time.Now,
- Lastupdate: time.Now(),
- }
- return &bucket
- }
- func (b *LeakyBucket) updateFill() {
- now := b.Now()
- if b.Fill > 0 {
- elapsed := now.Sub(b.Lastupdate)
- b.Fill -= float64(elapsed) / float64(b.LeakInterval)
- if b.Fill < 0 {
- b.Fill = 0
- }
- }
- b.Lastupdate = now
- }
- func (b *LeakyBucket) Pour(amount uint16) bool {
- b.updateFill()
- var newfill float64 = b.Fill + float64(amount)
- if newfill > float64(b.Size) {
- return false
- }
- b.Fill = newfill
- return true
- }
- // The time at which this bucket will be completely drained
- func (b *LeakyBucket) DrainedAt() time.Time {
- return b.Lastupdate.Add(time.Duration(b.Fill * float64(b.LeakInterval)))
- }
- // The duration until this bucket is completely drained
- func (b *LeakyBucket) TimeToDrain() time.Duration {
- return b.DrainedAt().Sub(b.Now())
- }
- func (b *LeakyBucket) TimeSinceLastUpdate() time.Duration {
- return b.Now().Sub(b.Lastupdate)
- }
- type LeakyBucketSer struct {
- Size uint16
- Fill float64
- LeakInterval time.Duration // time.Duration for 1 unit of size to leak
- Lastupdate time.Time
- }
- func (b *LeakyBucket) Serialise() *LeakyBucketSer {
- bucket := LeakyBucketSer{
- Size: b.Size,
- Fill: b.Fill,
- LeakInterval: b.LeakInterval,
- Lastupdate: b.Lastupdate,
- }
- return &bucket
- }
- func (b *LeakyBucketSer) DeSerialise() *LeakyBucket {
- bucket := LeakyBucket{
- Size: b.Size,
- Fill: b.Fill,
- LeakInterval: b.LeakInterval,
- Lastupdate: b.Lastupdate,
- Now: time.Now,
- }
- return &bucket
- }
|