github.com/songzhibin97/gkit@v1.2.13/restrictor/ratelimite/ratelimite_test.go (about)

     1  package ratelimite
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/juju/ratelimit"
     9  )
    10  
    11  func TestRateLimit(t *testing.T) {
    12  	// 创建指定填充速率和容量大小的令牌桶
    13  	// func NewBucket(fillInterval time.Duration, capacity int64) *Bucket
    14  	// 创建指定填充速率、容量大小和每次填充的令牌数的令牌桶
    15  	// func NewBucketWithQuantum(fillInterval time.Duration, capacity, quantum int64) *Bucket
    16  	// 创建填充速度为指定速率和容量大小的令牌桶
    17  	// NewBucketWithRate(0.1, 200) 表示每秒填充20个令牌
    18  	// func NewBucketWithRate(rate float64, capacity int64) *Bucket
    19  
    20  	bucket := ratelimit.NewBucket(time.Second/2, 4)
    21  	af, wf := NewRateLimit(bucket)
    22  	// 暂停3秒,等待桶满
    23  	time.Sleep(3 * time.Second)
    24  	for i := 0; i < 10; i++ {
    25  		t.Log("i:", i, af.Allow())
    26  	}
    27  
    28  	// 暂停3秒,等待桶满
    29  	time.Sleep(3 * time.Second)
    30  	for i := 0; i < 5; i++ {
    31  		t.Log("i:", i, af.AllowN(time.Now(), 2))
    32  	}
    33  
    34  	// 暂停3秒,等待桶满
    35  	time.Sleep(3 * time.Second)
    36  	for i := 0; i < 10; i++ {
    37  		func(i int) {
    38  			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    39  			defer cancel()
    40  			t.Log("i:", i, wf.Wait(ctx))
    41  		}(i)
    42  	}
    43  	// 暂停3秒,等待桶满
    44  	time.Sleep(3 * time.Second)
    45  	for i := 0; i < 5; i++ {
    46  		func(i int) {
    47  			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    48  			defer cancel()
    49  			t.Log("i:", i, wf.WaitN(ctx, 2))
    50  		}(i)
    51  	}
    52  }