github.com/cloudwego/kitex@v0.9.0/pkg/limiter/qps_limiter.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo 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 limiter
    18  
    19  import (
    20  	"context"
    21  	"sync/atomic"
    22  	"time"
    23  )
    24  
    25  var fixedWindowTime = time.Second
    26  
    27  // qpsLimiter implements the RateLimiter interface.
    28  type qpsLimiter struct {
    29  	limit      int32
    30  	tokens     int32
    31  	interval   time.Duration
    32  	once       int32
    33  	ticker     *time.Ticker
    34  	tickerDone chan bool
    35  }
    36  
    37  // NewQPSLimiter creates qpsLimiter.
    38  func NewQPSLimiter(interval time.Duration, limit int) RateLimiter {
    39  	once := calcOnce(interval, limit)
    40  	l := &qpsLimiter{
    41  		limit:    int32(limit),
    42  		interval: interval,
    43  		tokens:   once,
    44  		once:     once,
    45  	}
    46  	go l.startTicker(interval)
    47  	return l
    48  }
    49  
    50  // UpdateLimit update limitation of QPS. It is **not** concurrent-safe.
    51  func (l *qpsLimiter) UpdateLimit(limit int) {
    52  	once := calcOnce(l.interval, limit)
    53  	atomic.StoreInt32(&l.limit, int32(limit))
    54  	atomic.StoreInt32(&l.once, once)
    55  	l.resetTokens(once)
    56  }
    57  
    58  // UpdateQPSLimit update the interval and limit. It is **not** concurrent-safe.
    59  func (l *qpsLimiter) UpdateQPSLimit(interval time.Duration, limit int) {
    60  	once := calcOnce(interval, limit)
    61  	atomic.StoreInt32(&l.limit, int32(limit))
    62  	atomic.StoreInt32(&l.once, once)
    63  	l.resetTokens(once)
    64  	if interval != l.interval {
    65  		l.interval = interval
    66  		l.stopTicker()
    67  		go l.startTicker(interval)
    68  	}
    69  }
    70  
    71  // Acquire one token.
    72  func (l *qpsLimiter) Acquire(ctx context.Context) bool {
    73  	if atomic.LoadInt32(&l.limit) <= 0 {
    74  		return true
    75  	}
    76  	if atomic.LoadInt32(&l.tokens) <= 0 {
    77  		return false
    78  	}
    79  	return atomic.AddInt32(&l.tokens, -1) >= 0
    80  }
    81  
    82  // Status returns the current status.
    83  func (l *qpsLimiter) Status(ctx context.Context) (max, cur int, interval time.Duration) {
    84  	max = int(atomic.LoadInt32(&l.limit))
    85  	cur = int(atomic.LoadInt32(&l.tokens))
    86  	interval = l.interval
    87  	return
    88  }
    89  
    90  func (l *qpsLimiter) startTicker(interval time.Duration) {
    91  	l.ticker = time.NewTicker(interval)
    92  	defer l.ticker.Stop()
    93  	l.tickerDone = make(chan bool, 1)
    94  	tc := l.ticker.C
    95  	td := l.tickerDone
    96  	// ticker and tickerDone can be reset, cannot use l.ticker or l.tickerDone directly
    97  	for {
    98  		select {
    99  		case <-tc:
   100  			l.updateToken()
   101  		case <-td:
   102  			return
   103  		}
   104  	}
   105  }
   106  
   107  func (l *qpsLimiter) stopTicker() {
   108  	if l.tickerDone == nil {
   109  		return
   110  	}
   111  	select {
   112  	case l.tickerDone <- true:
   113  	default:
   114  	}
   115  }
   116  
   117  // Some deviation is allowed here to gain better performance.
   118  func (l *qpsLimiter) updateToken() {
   119  	if atomic.LoadInt32(&l.limit) < atomic.LoadInt32(&l.tokens) {
   120  		return
   121  	}
   122  
   123  	once := atomic.LoadInt32(&l.once)
   124  
   125  	delta := atomic.LoadInt32(&l.limit) - atomic.LoadInt32(&l.tokens)
   126  
   127  	if delta > once || delta < 0 {
   128  		delta = once
   129  	}
   130  
   131  	newTokens := atomic.AddInt32(&l.tokens, delta)
   132  	if newTokens < once {
   133  		atomic.StoreInt32(&l.tokens, once)
   134  	}
   135  }
   136  
   137  func calcOnce(interval time.Duration, limit int) int32 {
   138  	if interval > time.Second {
   139  		interval = time.Second
   140  	}
   141  	once := int32(float64(limit) / (fixedWindowTime.Seconds() / interval.Seconds()))
   142  	if once < 0 {
   143  		once = 0
   144  	}
   145  	return once
   146  }
   147  
   148  func (l *qpsLimiter) resetTokens(once int32) {
   149  	if atomic.LoadInt32(&l.tokens) > once {
   150  		atomic.StoreInt32(&l.tokens, once)
   151  	}
   152  }