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

     1  //go:build !race
     2  // +build !race
     3  
     4  /*
     5   * Copyright 2021 CloudWeGo Authors
     6   *
     7   * Licensed under the Apache License, Version 2.0 (the "License");
     8   * you may not use this file except in compliance with the License.
     9   * You may obtain a copy of the License at
    10   *
    11   *     http://www.apache.org/licenses/LICENSE-2.0
    12   *
    13   * Unless required by applicable law or agreed to in writing, software
    14   * distributed under the License is distributed on an "AS IS" BASIS,
    15   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    16   * See the License for the specific language governing permissions and
    17   * limitations under the License.
    18   */
    19  
    20  package limiter
    21  
    22  import (
    23  	"context"
    24  	"math"
    25  	"sync"
    26  	"sync/atomic"
    27  	"testing"
    28  	"time"
    29  
    30  	"github.com/cloudwego/kitex/internal/test"
    31  )
    32  
    33  func TestQPSLimiter(t *testing.T) {
    34  	qps := 1000
    35  	interval := time.Second / 100
    36  	ctx := context.Background()
    37  	limiter := NewQPSLimiter(interval, qps)
    38  
    39  	// case1: init
    40  	max, cur, itv := limiter.Status(ctx)
    41  	test.Assert(t, max == qps)
    42  	test.Assert(t, cur == int(float64(qps)/(time.Second.Seconds()/interval.Seconds())), cur)
    43  	test.Assert(t, itv == interval)
    44  
    45  	// case2: get tokens with parallel 3 and QPS is 1000
    46  	concurrent := 3
    47  	var wg sync.WaitGroup
    48  	wg.Add(concurrent)
    49  	var count, stopFlag int32
    50  	now := time.Now()
    51  	for i := 0; i < concurrent; i++ {
    52  		go func() {
    53  			for atomic.LoadInt32(&stopFlag) == 0 {
    54  				if limiter.Acquire(ctx) {
    55  					atomic.AddInt32(&count, 1)
    56  				}
    57  			}
    58  			wg.Done()
    59  		}()
    60  	}
    61  	time.AfterFunc(time.Second*2, func() {
    62  		atomic.StoreInt32(&stopFlag, 1)
    63  	})
    64  	wg.Wait()
    65  	actualQPS := count / int32(time.Since(now).Seconds())
    66  	delta := math.Abs(float64(actualQPS - int32(qps)))
    67  	// the diff range need be < 2%
    68  	test.Assert(t, delta < float64(qps)*0.02, delta)
    69  
    70  	// case3: UpdateQPSLimit and get Status
    71  	qps = 10000
    72  	count = 0
    73  	interval = time.Second / 50
    74  	limiter.(*qpsLimiter).UpdateQPSLimit(interval, qps)
    75  	max, _, itv = limiter.Status(ctx)
    76  	test.Assert(t, max == qps)
    77  	test.Assert(t, limiter.(*qpsLimiter).once == int32(float64(qps)/(time.Second.Seconds()/interval.Seconds())), limiter.(*qpsLimiter).once)
    78  	test.Assert(t, itv == interval)
    79  	wg.Add(concurrent)
    80  	stopFlag = 0
    81  	now = time.Now()
    82  	for i := 0; i < concurrent; i++ {
    83  		go func() {
    84  			for atomic.LoadInt32(&stopFlag) == 0 {
    85  				if limiter.Acquire(ctx) {
    86  					atomic.AddInt32(&count, 1)
    87  				}
    88  			}
    89  			wg.Done()
    90  		}()
    91  	}
    92  	time.AfterFunc(time.Second*2, func() {
    93  		atomic.StoreInt32(&stopFlag, 1)
    94  	})
    95  	wg.Wait()
    96  	actualQPS = count / int32(time.Since(now).Seconds())
    97  	delta = math.Abs(float64(actualQPS - int32(qps)))
    98  	// the diff range need be < 5%, the range is larger because the interval config is larger
    99  	test.Assert(t, delta < float64(qps)*0.05, delta, actualQPS)
   100  
   101  	// case4: UpdateLimit and get Status
   102  	qps = 300
   103  	count = 0
   104  	limiter.(*qpsLimiter).UpdateLimit(qps)
   105  	time.Sleep(interval)
   106  	max, _, itv = limiter.Status(ctx)
   107  	test.Assert(t, max == qps)
   108  	test.Assert(t, limiter.(*qpsLimiter).once == int32(float64(qps)/(time.Second.Seconds()/interval.Seconds())), limiter.(*qpsLimiter).once)
   109  	test.Assert(t, itv == interval)
   110  	wg.Add(concurrent)
   111  	stopFlag = 0
   112  	now = time.Now()
   113  	for i := 0; i < concurrent; i++ {
   114  		go func() {
   115  			for atomic.LoadInt32(&stopFlag) == 0 {
   116  				if limiter.Acquire(ctx) {
   117  					atomic.AddInt32(&count, 1)
   118  				}
   119  			}
   120  			wg.Done()
   121  		}()
   122  	}
   123  	time.AfterFunc(time.Second*2, func() {
   124  		atomic.StoreInt32(&stopFlag, 1)
   125  	})
   126  	wg.Wait()
   127  	actualQPS = count / int32(time.Since(now).Seconds())
   128  	delta = math.Abs(float64(actualQPS - int32(qps)))
   129  	// the diff range need be < 5%, the range is larger because the interval config is larger
   130  	test.Assert(t, delta < float64(qps)*0.05, delta, actualQPS)
   131  
   132  	// case5: UpdateLimit zero
   133  	qps = 0
   134  	count = 0
   135  	limiter.(*qpsLimiter).UpdateLimit(qps)
   136  	time.Sleep(interval)
   137  	max, _, itv = limiter.Status(ctx)
   138  	test.Assert(t, max == qps)
   139  	test.Assert(t, limiter.(*qpsLimiter).once == int32(float64(qps)/(time.Second.Seconds()/interval.Seconds())), limiter.(*qpsLimiter).once)
   140  	test.Assert(t, itv == interval)
   141  	wg.Add(concurrent)
   142  	var failedCount int32
   143  	stopFlag = 0
   144  	for i := 0; i < concurrent; i++ {
   145  		go func() {
   146  			for atomic.LoadInt32(&stopFlag) == 0 {
   147  				if !limiter.Acquire(ctx) {
   148  					atomic.AddInt32(&failedCount, 1)
   149  				}
   150  			}
   151  			wg.Done()
   152  		}()
   153  	}
   154  	time.AfterFunc(time.Second*2, func() {
   155  		atomic.StoreInt32(&stopFlag, 1)
   156  	})
   157  	wg.Wait()
   158  	test.Assert(t, failedCount == 0)
   159  }
   160  
   161  func TestCalcOnce(t *testing.T) {
   162  	ret := calcOnce(10*time.Millisecond, 100)
   163  	test.Assert(t, ret > 0)
   164  
   165  	// interval > time.Second
   166  	ret = calcOnce(time.Minute, 100)
   167  	test.Assert(t, ret == 100)
   168  
   169  	// limit < 0
   170  	ret = calcOnce(time.Second, -100)
   171  	test.Assert(t, ret == 0)
   172  }