github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/framework/ratelimit/redisrate/gcralimiter.go (about)

     1  package redisrate
     2  
     3  import (
     4  	"context"
     5  	logger "github.com/unionj-cloud/go-doudou/toolkit/zlogger"
     6  	"github.com/unionj-cloud/go-doudou/framework/ratelimit"
     7  	"strconv"
     8  	"time"
     9  
    10  	"github.com/go-redis/redis/v8"
    11  )
    12  
    13  const redisPrefix = "go-doudou:rate:"
    14  
    15  type Rediser interface {
    16  	Eval(ctx context.Context, script string, keys []string, args ...interface{}) *redis.Cmd
    17  	EvalSha(ctx context.Context, sha1 string, keys []string, args ...interface{}) *redis.Cmd
    18  	ScriptExists(ctx context.Context, hashes ...string) *redis.BoolSliceCmd
    19  	ScriptLoad(ctx context.Context, script string) *redis.StringCmd
    20  	Del(ctx context.Context, keys ...string) *redis.IntCmd
    21  }
    22  
    23  type LimitFn func(ctx context.Context) ratelimit.Limit
    24  
    25  type GcraLimiter struct {
    26  	rdb     Rediser
    27  	key     string
    28  	limit   ratelimit.Limit
    29  	limitFn LimitFn
    30  }
    31  
    32  func (gl *GcraLimiter) AllowCtx(ctx context.Context) bool {
    33  	allow, err := gl.AllowECtx(ctx)
    34  	if err != nil {
    35  		logger.Error().Err(err).Msg("")
    36  		return false
    37  	}
    38  	return allow
    39  }
    40  
    41  func (gl *GcraLimiter) Allow() bool {
    42  	allow, err := gl.AllowE()
    43  	if err != nil {
    44  		logger.Error().Err(err).Msg("")
    45  		return false
    46  	}
    47  	return allow
    48  }
    49  
    50  // Wait you'd better pass a timeout or cancelable context
    51  func (gl *GcraLimiter) Wait(ctx context.Context) error {
    52  	for {
    53  		retryAfter, allow, err := gl.ReserveECtx(ctx)
    54  		if err != nil {
    55  			return err
    56  		}
    57  		if allow {
    58  			return nil
    59  		}
    60  		time.Sleep(retryAfter)
    61  	}
    62  }
    63  
    64  func (gl *GcraLimiter) AllowE() (bool, error) {
    65  	allow, err := gl.AllowN(context.Background(), 1)
    66  	return allow.Allowed > 0, err
    67  }
    68  
    69  func (gl *GcraLimiter) ReserveE() (time.Duration, bool, error) {
    70  	allow, err := gl.AllowN(context.Background(), 1)
    71  	if err != nil {
    72  		return 0, false, err
    73  	}
    74  	return allow.RetryAfter, allow.Allowed > 0, nil
    75  }
    76  
    77  func (gl *GcraLimiter) AllowECtx(ctx context.Context) (bool, error) {
    78  	allow, err := gl.AllowN(ctx, 1)
    79  	return allow.Allowed > 0, err
    80  }
    81  
    82  func (gl *GcraLimiter) ReserveECtx(ctx context.Context) (time.Duration, bool, error) {
    83  	allow, err := gl.AllowN(ctx, 1)
    84  	if err != nil {
    85  		return 0, false, err
    86  	}
    87  	return allow.RetryAfter, allow.Allowed > 0, nil
    88  }
    89  
    90  // NewGcraLimiter returns a new Limiter.
    91  func NewGcraLimiter(rdb Rediser, key string, r float64, period time.Duration, b int) ratelimit.Limiter {
    92  	return &GcraLimiter{
    93  		rdb: rdb,
    94  		key: key,
    95  		limit: ratelimit.Limit{
    96  			Rate:   r,
    97  			Burst:  b,
    98  			Period: period,
    99  		},
   100  	}
   101  }
   102  
   103  // NewGcraLimiterLimit returns a new Limiter.
   104  func NewGcraLimiterLimit(rdb Rediser, key string, l ratelimit.Limit) ratelimit.Limiter {
   105  	return &GcraLimiter{
   106  		rdb:   rdb,
   107  		key:   key,
   108  		limit: l,
   109  	}
   110  }
   111  
   112  // NewGcraLimiterLimitFn returns a new Limiter.
   113  func NewGcraLimiterLimitFn(rdb Rediser, key string, fn LimitFn) ratelimit.Limiter {
   114  	return &GcraLimiter{
   115  		rdb:     rdb,
   116  		key:     key,
   117  		limitFn: fn,
   118  	}
   119  }
   120  
   121  // AllowN reports whether n events may happen at time now.
   122  func (gl *GcraLimiter) AllowN(ctx context.Context, n int) (res *Result, err error) {
   123  	limit := gl.limit
   124  	if gl.limitFn != nil {
   125  		limit = gl.limitFn(ctx)
   126  	}
   127  	values := []interface{}{limit.Burst, limit.Rate, limit.Period.Seconds(), n}
   128  	v, err := allowN.Run(ctx, gl.rdb, []string{redisPrefix + gl.key}, values...).Result()
   129  	if err != nil {
   130  		return nil, err
   131  	}
   132  
   133  	values = v.([]interface{})
   134  
   135  	retryAfter, err := strconv.ParseFloat(values[2].(string), 64)
   136  	if err != nil {
   137  		return nil, err
   138  	}
   139  
   140  	resetAfter, err := strconv.ParseFloat(values[3].(string), 64)
   141  	if err != nil {
   142  		return nil, err
   143  	}
   144  
   145  	res = &Result{
   146  		Limit:      limit,
   147  		Allowed:    int(values[0].(int64)),
   148  		Remaining:  int(values[1].(int64)),
   149  		RetryAfter: dur(retryAfter),
   150  		ResetAfter: dur(resetAfter),
   151  	}
   152  	return res, nil
   153  }
   154  
   155  // Reset gets a key and reset all limitations and previous usages
   156  func (gl *GcraLimiter) Reset(ctx context.Context) error {
   157  	return gl.rdb.Del(ctx, redisPrefix+gl.key).Err()
   158  }
   159  
   160  func dur(f float64) time.Duration {
   161  	if f == -1 {
   162  		return -1
   163  	}
   164  	return time.Duration(f * float64(time.Second))
   165  }
   166  
   167  type Result struct {
   168  	// Limit is the limit that was used to obtain this result.
   169  	Limit ratelimit.Limit
   170  
   171  	// Allowed is the number of events that may happen at time now.
   172  	Allowed int
   173  
   174  	// Remaining is the maximum number of requests that could be
   175  	// permitted instantaneously for this key given the current
   176  	// state. For example, if a rate limiter allows 10 requests per
   177  	// second and has already received 6 requests for this key this
   178  	// second, Remaining would be 4.
   179  	Remaining int
   180  
   181  	// RetryAfter is the time until the next request will be permitted.
   182  	// It should be -1 unless the rate limit has been exceeded.
   183  	RetryAfter time.Duration
   184  
   185  	// ResetAfter is the time until the RateLimiter returns to its
   186  	// initial state for a given key. For example, if a rate limiter
   187  	// manages requests per second and received one request 200ms ago,
   188  	// Reset would return 800ms. You can also think of this as the time
   189  	// until Limit and Remaining will be equal.
   190  	ResetAfter time.Duration
   191  }