1// Package ratelimiter implements the Leaky Bucket ratelimiting algorithm with memcached and in-memory backends.
 2package ratelimiter
 3
 4import (
 5	"time"
 6)
 7
 8type LeakyBucket struct {
 9	Size         uint16
10	Fill         float64
11	LeakInterval time.Duration // time.Duration for 1 unit of size to leak
12	Lastupdate   time.Time
13	Now          func() time.Time
14}
15
16func NewLeakyBucket(size uint16, leakInterval time.Duration) *LeakyBucket {
17	bucket := LeakyBucket{
18		Size:         size,
19		Fill:         0,
20		LeakInterval: leakInterval,
21		Now:          time.Now,
22		Lastupdate:   time.Now(),
23	}
24
25	return &bucket
26}
27
28func (b *LeakyBucket) updateFill() {
29	now := b.Now()
30	if b.Fill > 0 {
31		elapsed := now.Sub(b.Lastupdate)
32
33		b.Fill -= float64(elapsed) / float64(b.LeakInterval)
34		if b.Fill < 0 {
35			b.Fill = 0
36		}
37	}
38	b.Lastupdate = now
39}
40
41func (b *LeakyBucket) Pour(amount uint16) bool {
42	b.updateFill()
43
44	var newfill float64 = b.Fill + float64(amount)
45
46	if newfill > float64(b.Size) {
47		return false
48	}
49
50	b.Fill = newfill
51
52	return true
53}
54
55// The time at which this bucket will be completely drained
56func (b *LeakyBucket) DrainedAt() time.Time {
57	return b.Lastupdate.Add(time.Duration(b.Fill * float64(b.LeakInterval)))
58}
59
60// The duration until this bucket is completely drained
61func (b *LeakyBucket) TimeToDrain() time.Duration {
62	return b.DrainedAt().Sub(b.Now())
63}
64
65func (b *LeakyBucket) TimeSinceLastUpdate() time.Duration {
66	return b.Now().Sub(b.Lastupdate)
67}
68
69type LeakyBucketSer struct {
70	Size         uint16
71	Fill         float64
72	LeakInterval time.Duration // time.Duration for 1 unit of size to leak
73	Lastupdate   time.Time
74}
75
76func (b *LeakyBucket) Serialise() *LeakyBucketSer {
77	bucket := LeakyBucketSer{
78		Size:         b.Size,
79		Fill:         b.Fill,
80		LeakInterval: b.LeakInterval,
81		Lastupdate:   b.Lastupdate,
82	}
83
84	return &bucket
85}
86
87func (b *LeakyBucketSer) DeSerialise() *LeakyBucket {
88	bucket := LeakyBucket{
89		Size:         b.Size,
90		Fill:         b.Fill,
91		LeakInterval: b.LeakInterval,
92		Lastupdate:   b.Lastupdate,
93		Now:          time.Now,
94	}
95
96	return &bucket
97}