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

     1  package restrictor
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  )
     7  
     8  type Allow interface {
     9  	// Allow AllowN(time.Now(),1)
    10  	Allow() bool
    11  	// AllowN 截止到某一时刻，目前桶中数目是否至少为 n 个，满足则返回 true，同时从桶中消费 n 个 token
    12  	AllowN(now time.Time, n int) bool
    13  }
    14  
    15  type Wait interface {
    16  	// Wait WaitN(ctx,1)
    17  	Wait(ctx context.Context) error
    18  	// WaitN 如果此时桶内 Token 数组不足 (小于 N)，那么 Wait 方法将会阻塞一段时间，直至 Token 满足条件。如果充足则直接返回
    19  	// 我们可以设置 context 的 Deadline 或者 Timeout，来决定此次 Wait 的最长时间。
    20  	WaitN(ctx context.Context, n int) error
    21  }
    22  
    23  // AllowFunc 实现 Allow 接口
    24  type AllowFunc func(now time.Time, n int) bool
    25  
    26  func (a AllowFunc) Allow() bool {
    27  	return a.AllowN(time.Now(), 1)
    28  }
    29  
    30  func (a AllowFunc) AllowN(now time.Time, n int) bool {
    31  	return a(now, n)
    32  }
    33  
    34  // WaitFunc 实现 Wait 接口
    35  type WaitFunc func(ctx context.Context, n int) error
    36  
    37  func (w WaitFunc) Wait(ctx context.Context) error {
    38  	return w.WaitN(ctx, 1)
    39  }
    40  
    41  func (w WaitFunc) WaitN(ctx context.Context, n int) error {
    42  	return w(ctx, n)
    43  }