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

     1  package ratelimite
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/juju/ratelimit"
     9  	"github.com/songzhibin97/gkit/restrictor"
    10  )
    11  
    12  var ErrTimeOut = errors.New("restrictor/ratelimit: 超时")
    13  
    14  // package ratelimite: https://pkg.go.dev/github.com/juju/ratelimit 实现 limiter 接口
    15  
    16  func NewRateLimit(bucket *ratelimit.Bucket) (restrictor.AllowFunc, restrictor.WaitFunc) {
    17  	return func(now time.Time, n int) bool {
    18  			return bucket.TakeAvailable(int64(n)) >= int64(n)
    19  		},
    20  		func(ctx context.Context, n int) error {
    21  			// 获取超时时间
    22  			if d, ok := ctx.Deadline(); ok {
    23  				if !bucket.WaitMaxDuration(int64(n), time.Until(d)) {
    24  					return ErrTimeOut
    25  				}
    26  				return nil
    27  			}
    28  			// 表示context没有设置超时时间
    29  			if bucket.WaitMaxDuration(int64(n), 100*time.Millisecond) {
    30  				return ErrTimeOut
    31  			}
    32  			return nil
    33  		}
    34  }