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  }