vitess.io/vitess@v0.16.2/go/vt/throttler/thread_throttler.go (about) 1 /* 2 Copyright 2019 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package throttler 18 19 import ( 20 "fmt" 21 "time" 22 23 "vitess.io/vitess/go/sync2" 24 25 "golang.org/x/time/rate" 26 ) 27 28 // threadThrottler implements the core logic which decides if a Throttle() call 29 // should be throttled (and for how long) or not. 30 type threadThrottler struct { 31 threadID int 32 33 // Fields below are unguarded because they must not be modified concurrently. 34 // currentSecond is the last time throttle() was called. The value is truncated to a 35 // second granularity. 36 currentSecond time.Time 37 // currentRate is the number of allowed requests since 'currentSecond'. 38 currentRate int64 39 40 // maxRate holds the last rate set by setMaxRate. It will be set 41 // in the limiter object in the next call to throttle(). 42 maxRate sync2.AtomicInt64 43 limiter *rate.Limiter 44 actualRateHistory *aggregatedIntervalHistory 45 } 46 47 func newThreadThrottler(threadID int, actualRateHistory *aggregatedIntervalHistory) *threadThrottler { 48 // Initially, the set rate is 0 and the throttler should deny all requests. We'd like the first 49 // throttle() call to be accepted after setMaxRate() has been called with a nonzero rate. 50 // Unfortunately, if we initialize the limiter rate to 0, the internal token buffer will be 51 // empty by the time the first throttle() call is executed and it will be denied. 52 // Instead, we initialize the limiter rate to 1. This way the token buffer will be full (assuming 53 // the 'now' parameter of the first throttle() call is at least 1 second) and the rate will 54 // be reset to 0 if setMaxRate() has not been called with a nonzero rate. 55 result := threadThrottler{ 56 threadID: threadID, 57 actualRateHistory: actualRateHistory, 58 limiter: rate.NewLimiter(1 /* limit */, 1 /* burst */), 59 } 60 return &result 61 } 62 63 var ( 64 oneSecond = time.Time{}.Add(1 * time.Second) 65 ) 66 67 func (t *threadThrottler) throttle(now time.Time) time.Duration { 68 if now.Before(oneSecond) { 69 panic(fmt.Sprintf( 70 "BUG: throttle() must not be called with a time of less than 1 second. now: %v", 71 now)) 72 } 73 74 // Pass the limit set by the last call to setMaxRate. Limiter.SetLimitAt 75 // is idempotent, so we can call it with the same value more than once without 76 // issues. 77 t.limiter.SetLimitAt(now, rate.Limit(t.maxRate.Get())) 78 79 // Initialize or advance the current second interval when necessary. 80 nowSecond := now.Truncate(time.Second) 81 if t.currentSecond != nowSecond { 82 // Report the number of successful (not-throttled) requests from the "last" second if this is 83 // not the first time 'throttle' is called. 84 if !t.currentSecond.IsZero() { 85 t.actualRateHistory.addPerThread(t.threadID, record{t.currentSecond, t.currentRate}) 86 } 87 t.currentRate = 0 88 t.currentSecond = nowSecond 89 } 90 91 if t.limiter.Limit() == 0 { 92 // If the limit is 0 this request won't be let through. However, the caller 93 // should poll again in the future in case the limit has changed. 94 return 1 * time.Second 95 } 96 97 // Figure out how long to backoff: We use the limiter.ReserveN() method to reserve an event. 98 // The returned reservation contains the backoff delay. We cancel the reservation if the 99 // delay is greater than 0, since the caller is expected to call throttle() again at that time 100 // rather than proceed. 101 reservation := t.limiter.ReserveN(now, 1) 102 if !reservation.OK() { 103 panic(fmt.Sprintf("BUG: limiter was unable to reserve an event. "+ 104 "threadThrottler: %v, reservation:%v", *t, *reservation)) 105 } 106 waitDuration := reservation.DelayFrom(now) 107 if waitDuration <= 0 { 108 t.currentRate++ 109 return NotThrottled 110 } 111 reservation.CancelAt(now) 112 return waitDuration 113 } 114 115 // setMaxRate sets the maximum rate for the next time throttle() is called. 116 // setMaxRate() can be called concurrently with other methods of this object. 117 func (t *threadThrottler) setMaxRate(newRate int64) { 118 t.maxRate.Set(newRate) 119 } 120 121 // maxRate returns the rate set by the last call to setMaxRate(). 122 // If setMaxRate() was not called, this method returns 0. 123 func (t *threadThrottler) getMaxRate() int64 { 124 return t.maxRate.Get() 125 }