go.temporal.io/server@v1.23.0/common/quotas/rate_limiter_impl.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package quotas
    26  
    27  import (
    28  	"sync"
    29  	"time"
    30  
    31  	"go.temporal.io/server/common/clock"
    32  	"golang.org/x/time/rate"
    33  )
    34  
    35  type (
    36  	// RateLimiterImpl is a wrapper around the golang rate limiter
    37  	RateLimiterImpl struct {
    38  		sync.RWMutex
    39  		rps        float64
    40  		burst      int
    41  		timeSource clock.TimeSource
    42  		ClockedRateLimiter
    43  	}
    44  )
    45  
    46  var _ RateLimiter = (*RateLimiterImpl)(nil)
    47  
    48  // NewRateLimiter returns a new rate limiter that can handle dynamic
    49  // configuration updates
    50  func NewRateLimiter(newRPS float64, newBurst int) *RateLimiterImpl {
    51  	limiter := rate.NewLimiter(rate.Limit(newRPS), newBurst)
    52  	ts := clock.NewRealTimeSource()
    53  	rl := &RateLimiterImpl{
    54  		rps:                newRPS,
    55  		burst:              newBurst,
    56  		timeSource:         ts,
    57  		ClockedRateLimiter: NewClockedRateLimiter(limiter, ts),
    58  	}
    59  
    60  	return rl
    61  }
    62  
    63  // SetRate set the rate of the rate limiter
    64  func (rl *RateLimiterImpl) SetRPS(rps float64) {
    65  	rl.refreshInternalRateLimiterImpl(&rps, nil)
    66  }
    67  
    68  // SetBurst set the burst of the rate limiter
    69  func (rl *RateLimiterImpl) SetBurst(burst int) {
    70  	rl.refreshInternalRateLimiterImpl(nil, &burst)
    71  }
    72  
    73  func (rl *RateLimiterImpl) Reserve() Reservation {
    74  	return rl.ClockedRateLimiter.Reserve()
    75  }
    76  
    77  func (rl *RateLimiterImpl) ReserveN(now time.Time, n int) Reservation {
    78  	return rl.ClockedRateLimiter.ReserveN(now, n)
    79  }
    80  
    81  // SetRateBurst set the rps & burst of the rate limiter
    82  func (rl *RateLimiterImpl) SetRateBurst(rps float64, burst int) {
    83  	rl.refreshInternalRateLimiterImpl(&rps, &burst)
    84  }
    85  
    86  // Rate returns the rps for this rate limiter
    87  func (rl *RateLimiterImpl) Rate() float64 {
    88  	rl.Lock()
    89  	defer rl.Unlock()
    90  
    91  	return rl.rps
    92  }
    93  
    94  // Burst returns the burst for this rate limiter
    95  func (rl *RateLimiterImpl) Burst() int {
    96  	rl.Lock()
    97  	defer rl.Unlock()
    98  
    99  	return rl.burst
   100  }
   101  
   102  func (rl *RateLimiterImpl) refreshInternalRateLimiterImpl(
   103  	newRate *float64,
   104  	newBurst *int,
   105  ) {
   106  	rl.Lock()
   107  	defer rl.Unlock()
   108  
   109  	refresh := false
   110  
   111  	if newRate != nil && rl.rps != *newRate {
   112  		rl.rps = *newRate
   113  		refresh = true
   114  	}
   115  
   116  	if newBurst != nil && rl.burst != *newBurst {
   117  		rl.burst = *newBurst
   118  		refresh = true
   119  	}
   120  
   121  	if refresh {
   122  		now := rl.timeSource.Now()
   123  		rl.SetLimitAt(now, rate.Limit(rl.rps))
   124  		rl.SetBurstAt(now, rl.burst)
   125  	}
   126  }