|
| 1 | +package limiter |
| 2 | + |
| 3 | +import ( |
| 4 | + "sync" |
| 5 | + "time" |
| 6 | + |
| 7 | + "github.com/golang/glog" |
| 8 | + utilruntime "k8s.io/apimachinery/pkg/util/runtime" |
| 9 | +) |
| 10 | + |
| 11 | +// HandlerFunc defines function signature for a CoalescingSerializingRateLimiter. |
| 12 | +type HandlerFunc func() error |
| 13 | + |
| 14 | +// CoalescingSerializingRateLimiter guarantees that calls will not happen to the given function |
| 15 | +// more frequently than the given interval, and it guarantees that only one call will happen at a time. |
| 16 | +// The calls are not queued, i.e. if you make 5 calls to RegisterChange(), it does not guarantee that the |
| 17 | +// handler will be invoked 5 times, it merely guarantees it will be invoked once, and no more often than |
| 18 | +// the rate. |
| 19 | +// The calls to the handler will happen in the background and are expected to do their own locking if needed. |
| 20 | +type CoalescingSerializingRateLimiter struct { |
| 21 | + // handlerFunc is the function to rate limit and seriaize calls to. |
| 22 | + handlerFunc HandlerFunc |
| 23 | + |
| 24 | + // callInterval is the minimum time between the starts of handler calls. |
| 25 | + callInterval time.Duration |
| 26 | + |
| 27 | + // lastStart is the time the last run of the handler started. |
| 28 | + lastStart time.Time |
| 29 | + |
| 30 | + // changeReqTime is nil if no change has been registered since the last handler run completed, otherwise it is the |
| 31 | + // time last change was registered. |
| 32 | + changeReqTime *time.Time |
| 33 | + |
| 34 | + // handlerRunning indicates whether the Handler is actively running. |
| 35 | + handlerRunning bool |
| 36 | + |
| 37 | + // lock protects the CoalescingSerializingRateLimiter structure from multiple threads manipulating it at once. |
| 38 | + lock sync.Mutex |
| 39 | + |
| 40 | + // callbackTimer is the timer we use to make callbacks to re-run the function to decide if we need to do work. |
| 41 | + callbackTimer *time.Timer |
| 42 | +} |
| 43 | + |
| 44 | +func NewCoalescingSerializingRateLimiter(interval time.Duration, handlerFunc HandlerFunc) *CoalescingSerializingRateLimiter { |
| 45 | + limiter := &CoalescingSerializingRateLimiter{ |
| 46 | + handlerFunc: handlerFunc, |
| 47 | + callInterval: interval, |
| 48 | + lastStart: time.Time{}, |
| 49 | + changeReqTime: nil, |
| 50 | + handlerRunning: false, |
| 51 | + } |
| 52 | + |
| 53 | + return limiter |
| 54 | +} |
| 55 | + |
| 56 | +// RegisterChange() indicates that the rate limited function should be called. It may not immediately run it, but it will cause it to run within |
| 57 | +// the ReloadInterval. It will always immediately return, the function will be run in the background. Not every call to RegisterChange() will |
| 58 | +// result in the function getting called. If it is called repeatedly while it is still within the ReloadInterval since the last run, it will |
| 59 | +// only run once when the time allows it. |
| 60 | +func (csrl *CoalescingSerializingRateLimiter) RegisterChange() { |
| 61 | + glog.V(8).Infof("RegisterChange called") |
| 62 | + |
| 63 | + csrl.changeWorker(true) |
| 64 | +} |
| 65 | + |
| 66 | +func (csrl *CoalescingSerializingRateLimiter) changeWorker(userChanged bool) { |
| 67 | + csrl.lock.Lock() |
| 68 | + defer csrl.lock.Unlock() |
| 69 | + |
| 70 | + glog.V(8).Infof("changeWorker called") |
| 71 | + |
| 72 | + if userChanged && csrl.changeReqTime == nil { |
| 73 | + // They just registered a change manually (and we aren't in the middle of a change) |
| 74 | + now := time.Now() |
| 75 | + csrl.changeReqTime = &now |
| 76 | + } |
| 77 | + |
| 78 | + if csrl.handlerRunning { |
| 79 | + // We don't need to do anything else... there's a run in progress, and when it is done it will re-call this function at which point the work will then happen |
| 80 | + glog.V(8).Infof("The handler was already running (%v) started at %s, returning from the worker", csrl.handlerRunning, csrl.lastStart.String()) |
| 81 | + return |
| 82 | + } |
| 83 | + |
| 84 | + if csrl.changeReqTime == nil { |
| 85 | + // There's no work queued so we have nothing to do. We should only get here when |
| 86 | + // the function is re-called after a reload |
| 87 | + glog.V(8).Infof("No invoke requested time, so there's no queued work. Nothing to do.") |
| 88 | + return |
| 89 | + } |
| 90 | + |
| 91 | + // There is no handler running, let's see if we should run yet, or schedule a callback |
| 92 | + now := time.Now() |
| 93 | + sinceLastRun := now.Sub(csrl.lastStart) |
| 94 | + untilNextCallback := csrl.callInterval - sinceLastRun |
| 95 | + glog.V(8).Infof("Checking reload; now: %v, lastStart: %v, sinceLast %v, limit %v, remaining %v", now, csrl.lastStart, sinceLastRun, csrl.callInterval, untilNextCallback) |
| 96 | + |
| 97 | + if untilNextCallback > 0 { |
| 98 | + // We want to reload... but can't yet because some window is not satisfied |
| 99 | + if csrl.callbackTimer == nil { |
| 100 | + csrl.callbackTimer = time.AfterFunc(untilNextCallback, func() { csrl.changeWorker(false) }) |
| 101 | + } else { |
| 102 | + // While we are resetting the timer, it should have fired and be stopped. |
| 103 | + // The first time the worker is called it will know the precise duration |
| 104 | + // until when a run would be valid and has scheduled a timer for that point |
| 105 | + csrl.callbackTimer.Reset(untilNextCallback) |
| 106 | + } |
| 107 | + |
| 108 | + glog.V(8).Infof("Can't invoke the handler yet, need to delay %s, callback scheduled", untilNextCallback.String()) |
| 109 | + |
| 110 | + return |
| 111 | + } |
| 112 | + |
| 113 | + // Otherwise we can reload immediately... let's do it! |
| 114 | + glog.V(8).Infof("Calling the handler function (for invoke time %v)", csrl.changeReqTime) |
| 115 | + csrl.handlerRunning = true |
| 116 | + csrl.changeReqTime = nil |
| 117 | + csrl.lastStart = now |
| 118 | + |
| 119 | + // Go run the handler so we don't block the caller |
| 120 | + go csrl.runHandler() |
| 121 | + |
| 122 | + return |
| 123 | +} |
| 124 | + |
| 125 | +func (csrl *CoalescingSerializingRateLimiter) runHandler() { |
| 126 | + // Call the handler, but do it in its own function so we can cleanup in case the handler panics |
| 127 | + runHandler := func() error { |
| 128 | + defer func() { |
| 129 | + csrl.lock.Lock() |
| 130 | + csrl.handlerRunning = false |
| 131 | + csrl.lock.Unlock() |
| 132 | + }() |
| 133 | + |
| 134 | + return csrl.handlerFunc() |
| 135 | + } |
| 136 | + if err := runHandler(); err != nil { |
| 137 | + utilruntime.HandleError(err) |
| 138 | + } |
| 139 | + |
| 140 | + // Re-call the commit in case there is work waiting that came in while we were working |
| 141 | + // we want to call the top level commit in case the state has not changed |
| 142 | + glog.V(8).Infof("Re-Calling the worker after a reload in case work came in") |
| 143 | + csrl.changeWorker(false) |
| 144 | +} |
0 commit comments