throttle.go 1.82 KB
Newer Older
Aaron Raddon's avatar
Aaron Raddon committed
1 2 3 4 5 6 7 8
package gou

import (
	"time"
)

type Throttler struct {

9 10 11
	// Limit to this events/per
	maxPer float64
	per    float64
12
	count  int32
Aaron Raddon's avatar
Aaron Raddon committed
13 14 15 16 17 18 19 20 21 22

	// Last Event
	last time.Time

	// How many events are allowed left to happen?
	// Starts at limit, decrements down
	allowance float64
}

// new Throttler that will tell you to limit or not based
Aaron Raddon's avatar
Aaron Raddon committed
23
// on given @max events @per duration
24
func NewThrottler(max int, per time.Duration) *Throttler {
Aaron Raddon's avatar
Aaron Raddon committed
25
	return &Throttler{
26 27
		maxPer:    float64(max),
		allowance: float64(max),
28
		count:     int32(0),
Aaron Raddon's avatar
Aaron Raddon committed
29
		last:      time.Now(),
30
		per:       per.Seconds(),
Aaron Raddon's avatar
Aaron Raddon committed
31 32 33 34
	}
}

// Should we limit this because we are above rate?
35 36
// Returns a bool of whether to throttle the message, and a count
// of previous log messages throttled since last log message.
Aaron Raddon's avatar
Aaron Raddon committed
37
func (r *Throttler) ThrottleAdd(ct int32) (bool, int32) {
Aaron Raddon's avatar
Aaron Raddon committed
38

39
	if r.maxPer == 0 {
40
		return false, 0
Aaron Raddon's avatar
Aaron Raddon committed
41 42 43 44
	}

	// http://stackoverflow.com/questions/667508/whats-a-good-rate-limiting-algorithm
	now := time.Now()
45
	elapsed := float64(now.Sub(r.last).Nanoseconds()) / 1e9 // seconds
Aaron Raddon's avatar
Aaron Raddon committed
46
	r.last = now
47
	r.allowance += elapsed * (r.maxPer / r.per)
Aaron Raddon's avatar
Aaron Raddon committed
48

49 50 51
	//Infof("maxRate: %v  cur: %v elapsed:%-6.6f  incr: %v", r.maxPer, int(r.allowance), elapsed, elapsed*float64(r.maxPer))
	if r.allowance > r.maxPer {
		r.allowance = r.maxPer
Aaron Raddon's avatar
Aaron Raddon committed
52 53
	}

54
	if r.allowance < 1.0 {
Aaron Raddon's avatar
Aaron Raddon committed
55
		r.count += ct        // increment throttled log count
56
		return true, r.count // do throttle/limit
Aaron Raddon's avatar
Aaron Raddon committed
57 58
	}

59 60 61
	tmpCount := r.count
	r.count = 0 // reset count

Aaron Raddon's avatar
Aaron Raddon committed
62
	r.allowance -= 1.0
63 64 65
	return false, tmpCount // dont throttle, return previous throttle count
}

Aaron Raddon's avatar
Aaron Raddon committed
66 67 68 69 70 71 72
// Should we limit this because we are above rate?
// Returns a bool of whether to throttle the message, and a count
// of previous log messages throttled since last log message.
func (r *Throttler) Throttle() (bool, int32) {
	return r.ThrottleAdd(1)
}

73 74
func (r *Throttler) ThrottleCount() int32 {
	return r.count
Aaron Raddon's avatar
Aaron Raddon committed
75
}