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 }