1 package limiter 2 3 import "time" 4 5 // A limiter can be used to rate limit and/or coalesce a series of 6 // time-based events. This interface captures the logic of deciding 7 // what action should be taken when an event occurs at a specific time 8 // T. The possible actions are act upon the event, do nothing, or 9 // check back after a specific delay. 10 type Limiter interface { 11 // The Limit() method works kinda like an 8 ball. You pass it 12 // the time of an event, and it returns one of Yes, No, or 13 // "Later". A zero return value means "Yes", the event should 14 // be acted upon right away. A negative return value means 15 // "No", do nothing, and a positive return value means to 16 // check back after the returned delay. 17 // 18 // The order that this is invoked is important, and it is 19 // expected that this is invoked with a monotonically 20 // increasing set of timestamps. Typically this will be 21 // invoked when an event is generated and will be the result 22 // of time.Now(), but for testing or other cases, this could 23 // be invoked with any set of historic or future timestamps so 24 // long as they are invoked in monotonically increasing order. 25 // 26 // The result of this (when positive) is always relative to 27 // the passed in time. In other words to compute a deadline 28 // rather than a delay, you should take the result of Limit() 29 // and add it to whatever value you passed in: 30 // 31 // deadline = now + Limit(now). 32 // 33 Limit(now time.Time) time.Duration 34 } 35 36 type limiter struct { 37 // we should never fire more than one event within this timespan 38 interval time.Duration 39 // records the last event we fired 40 lastAction time.Time 41 // records the point in the future to which we delayed an event 42 deadline time.Time 43 } 44 45 // Constructs a new limiter that will coalesce any events occurring 46 // within the specified interval. 47 func NewInterval(interval time.Duration) Limiter { 48 return &limiter{ 49 interval: interval, 50 } 51 } 52 53 func (l *limiter) Limit(now time.Time) time.Duration { 54 since := now.Sub(l.lastAction) 55 switch { 56 case since >= l.interval: 57 // we haven't triggered any event recently, so we 58 // trigger this event and record that we did so 59 l.lastAction = now 60 return 0 61 case l.deadline.After(now): 62 // we are still waiting for an event that we delayed 63 // into the future, so we can drop this one 64 return -1 65 default: 66 // we have multiple events within the same interval, 67 // and there are no delayed events, so we delay this 68 // one and record the point in the future to which we 69 // delayed it 70 delay := l.interval - since 71 l.deadline = now.Add(delay) 72 return delay 73 } 74 } 75 76 type composite struct { 77 first Limiter 78 second Limiter 79 delay time.Duration 80 started bool 81 deadline time.Time 82 } 83 84 func NewComposite(first, second Limiter, delay time.Duration) Limiter { 85 return &composite{ 86 first: first, 87 second: second, 88 delay: delay, 89 } 90 } 91 92 func (c *composite) Limit(now time.Time) time.Duration { 93 if !c.started { 94 c.started = true 95 c.deadline = now.Add(c.delay) 96 } 97 98 if now.After(c.deadline) { 99 return c.second.Limit(now) 100 } else { 101 return c.first.Limit(now) 102 } 103 } 104 105 type unlimited struct{} 106 107 func NewUnlimited() Limiter { 108 return &unlimited{} 109 } 110 111 func (u *unlimited) Limit(now time.Time) time.Duration { return 0 } 112