github.com/songzhibin97/gkit@v1.2.13/restrictor/rate/rate_test.go (about)

     1  package rate
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"golang.org/x/time/rate"
     9  )
    10  
    11  func TestRate(t *testing.T) {
    12  	// 第一个参数是 r Limit。代表每秒可以向 Token 桶中产生多少 token。Limit 实际上是 float64 的别名
    13  	// 第二个参数是 b int。b 代表 Token 桶的容量大小。
    14  	// limit := Every(100 * time.Millisecond);
    15  	// limiter := rate.NewLimiter(limit, 4)
    16  	// 以上就表示每 100ms 往桶中放一个 Token。本质上也就是一秒钟产生 10 个。
    17  	limiter := rate.NewLimiter(2, 4)
    18  	af, wf := NewRate(limiter)
    19  	// 暂停3秒,等待桶满
    20  	time.Sleep(3 * time.Second)
    21  	for i := 0; i < 10; i++ {
    22  		t.Log("i:", i, af.Allow())
    23  	}
    24  
    25  	// 暂停3秒,等待桶满
    26  	time.Sleep(3 * time.Second)
    27  	for i := 0; i < 5; i++ {
    28  		t.Log("i:", i, af.AllowN(time.Now(), 2))
    29  	}
    30  
    31  	// 暂停3秒,等待桶满
    32  	time.Sleep(3 * time.Second)
    33  	for i := 0; i < 10; i++ {
    34  		func(i int) {
    35  			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    36  			defer cancel()
    37  			t.Log("i:", i, wf.Wait(ctx))
    38  		}(i)
    39  	}
    40  	// 暂停3秒,等待桶满
    41  	time.Sleep(3 * time.Second)
    42  	for i := 0; i < 5; i++ {
    43  		func(i int) {
    44  			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    45  			defer cancel()
    46  			t.Log("i:", i, wf.WaitN(ctx, 2))
    47  		}(i)
    48  	}
    49  }