wip

ratelimiter/ratelimiter.go

@@ -0,0 +1,86 @@
+package ratelimiter
+
+import (
+	"errors"
+	"sync"
+	"time"
+)
+
+var ErrBackoff = errors.New("Backoff")
+
+type Config struct {
+	BurstLimit   int64         // Number of requests to allow to burst.
+	FillPeriod   time.Duration // Add one per period.
+	MaxWaitCount int64         // Max number of waiting requests. 0 disables.
+}
+
+type Limiter struct {
+	lock sync.Mutex
+
+	fillPeriod  time.Duration
+	minWaitTime time.Duration
+	maxWaitTime time.Duration
+
+	waitTime    time.Duration // If waitTime < 0, no waiting occurs.
+	lastRequest time.Time
+}
+
+func New(conf Config) *Limiter {
+	if conf.BurstLimit < 0 {
+		panic(conf.BurstLimit)
+	}
+	if conf.FillPeriod <= 0 {
+		panic(conf.FillPeriod)
+	}
+	if conf.MaxWaitCount < 0 {
+		panic(conf.MaxWaitCount)
+	}
+
+	lim := &Limiter{
+		lastRequest: time.Now(),
+		fillPeriod:  conf.FillPeriod,
+		waitTime:    -conf.FillPeriod * time.Duration(conf.BurstLimit),
+		minWaitTime: -conf.FillPeriod * time.Duration(conf.BurstLimit),
+		maxWaitTime: conf.FillPeriod * time.Duration(conf.MaxWaitCount),
+	}
+
+	lim.waitTime = lim.minWaitTime
+
+	return lim
+}
+
+func (lim *Limiter) limit(count int64) (time.Duration, error) {
+	lim.lock.Lock()
+	defer lim.lock.Unlock()
+
+	dt := time.Since(lim.lastRequest)
+
+	waitTime := lim.waitTime - dt + time.Duration(count)*lim.fillPeriod
+
+	if waitTime < lim.minWaitTime {
+		waitTime = lim.minWaitTime
+	} else if waitTime > lim.maxWaitTime {
+		return 0, ErrBackoff
+	}
+
+	lim.waitTime = waitTime
+	lim.lastRequest = lim.lastRequest.Add(dt)
+
+	return lim.waitTime, nil
+}
+
+// Apply the limiter to the calling thread. The function may sleep for up to
+// maxWaitTime before returning. If the timeout would need to be more than
+// maxWaitTime to enforce the rate limit, ErrBackoff is returned.
+func (lim *Limiter) Limit() error {
+	dt, err := lim.limit(1)
+	time.Sleep(dt) // Will return immediately for dt <= 0.
+	return err
+}
+
+// Apply the limiter for multiple items at once.
+func (lim *Limiter) LimitMultiple(count int64) error {
+	dt, err := lim.limit(count)
+	time.Sleep(dt)
+	return err
+}