vitess.io/vitess@v0.16.2/go/timer/rate_limiter.go (about)

     1  /*
     2  Copyright 2022 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 timer
    18  
    19  import (
    20  	"context"
    21  	"math"
    22  	"sync"
    23  	"sync/atomic"
    24  	"time"
    25  )
    26  
    27  // RateLimiter runs given tasks, at no more than one per defined duration.
    28  // For example, we can create a RateLimiter of 1second. Then, we can ask it, over time, to run many
    29  // tasks. It will only ever run a single task in any 1 second time frame. The rest are ignored.
    30  type RateLimiter struct {
    31  	tickerValue int64
    32  	lastDoValue int64
    33  
    34  	mu     sync.Mutex
    35  	cancel context.CancelFunc
    36  }
    37  
    38  // NewRateLimiter creates a new limiter with given duration. It is immediately ready to run tasks.
    39  func NewRateLimiter(d time.Duration) *RateLimiter {
    40  	r := &RateLimiter{tickerValue: 1}
    41  	ctx, cancel := context.WithCancel(context.Background())
    42  	r.cancel = cancel
    43  	go func() {
    44  		ticker := time.NewTicker(d)
    45  		defer ticker.Stop()
    46  		for {
    47  			select {
    48  			case <-ctx.Done():
    49  				return
    50  			case <-ticker.C:
    51  				atomic.StoreInt64(&r.tickerValue, r.tickerValue+1)
    52  			}
    53  		}
    54  	}()
    55  	return r
    56  }
    57  
    58  // Do runs a given func assuming rate limiting allows. This function is thread safe.
    59  // f may be nil, in which case it is not invoked.
    60  func (r *RateLimiter) Do(f func() error) (err error) {
    61  	r.mu.Lock()
    62  	defer r.mu.Unlock()
    63  
    64  	if r.lastDoValue >= atomic.LoadInt64(&r.tickerValue) {
    65  		return nil // rate limited. Skipped.
    66  	}
    67  	if f != nil {
    68  		err = f()
    69  	}
    70  	r.lastDoValue = atomic.LoadInt64(&r.tickerValue)
    71  	return err
    72  }
    73  
    74  // Stop terminates rate limiter's operation and will not allow any more Do() executions.
    75  func (r *RateLimiter) Stop() {
    76  	r.cancel()
    77  
    78  	r.mu.Lock()
    79  	defer r.mu.Unlock()
    80  
    81  	r.lastDoValue = math.MaxInt64
    82  }