36 lines
680 B
Go
36 lines
680 B
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"tm/pkg/redis"
|
|
|
|
redisv9 "github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
// Allow returns whether the request is within limit for the given key.
|
|
func Allow(ctx context.Context, client redis.Client, key string, limit int64, window time.Duration) (bool, error) {
|
|
if client == nil {
|
|
return true, nil
|
|
}
|
|
|
|
count, err := client.Incr(ctx, key)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
if count == 1 {
|
|
if err := client.Expire(ctx, key, window); err != nil {
|
|
return false, err
|
|
}
|
|
}
|
|
|
|
return count <= limit, nil
|
|
}
|
|
|
|
// IsNotFound reports whether err indicates a missing Redis key.
|
|
func IsNotFound(err error) bool {
|
|
return err == redisv9.Nil
|
|
}
|