68 lines
2.0 KiB
Go
68 lines
2.0 KiB
Go
package ozon
|
|
|
|
import (
|
|
"fmt"
|
|
"github.com/redis/rueidis"
|
|
"net/http"
|
|
"sipro-mps/internal/redis"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
windowSize = time.Second
|
|
rps = 50 // requests per second
|
|
)
|
|
|
|
var (
|
|
rateLimiterScript = rueidis.NewLuaScript(`
|
|
local key = KEYS[1]
|
|
local now = tonumber(ARGV[1])
|
|
local window = tonumber(ARGV[2])
|
|
local limit = tonumber(ARGV[3])
|
|
|
|
-- Удаляем старые записи вне окна времени
|
|
redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window)
|
|
local count = redis.call('ZCARD', key)
|
|
|
|
if count < limit then
|
|
-- Лимит не превышен, добавляем новую метку и устанавливаем TTL
|
|
redis.call('ZADD', key, now, now)
|
|
redis.call('EXPIRE', key, math.ceil(window / 1000000000))
|
|
return 0 -- Можно выполнять запрос сразу
|
|
else
|
|
-- Лимит превышен, находим самую старую метку
|
|
local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')[2]
|
|
-- Возвращаем время, которое нужно подождать до освобождения слота
|
|
return (tonumber(oldest) + window) - now
|
|
end
|
|
`)
|
|
)
|
|
|
|
type RateLimitTransport struct {
|
|
http.RoundTripper
|
|
}
|
|
|
|
func (t *RateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
|
ctx := req.Context()
|
|
clientId := req.Header.Get("Client-Id")
|
|
now := time.Now().UnixNano()
|
|
|
|
waitTime, err := rateLimiterScript.Exec(ctx, *redis.Client, []string{clientId}, []string{
|
|
fmt.Sprintf("%d", now),
|
|
fmt.Sprintf("%d", int64(windowSize)),
|
|
fmt.Sprintf("%d", rps),
|
|
}).ToInt64()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to execute rate limit script: %w", err)
|
|
}
|
|
if waitTime > 0 {
|
|
fmt.Printf("Rate limit exceeded for client %s, waiting for %d nanoseconds\n", clientId, waitTime)
|
|
time.Sleep(time.Duration(waitTime))
|
|
}
|
|
return t.RoundTripper.RoundTrip(req)
|
|
}
|
|
func NewRateLimitTransport() *RateLimitTransport {
|
|
|
|
return &RateLimitTransport{RoundTripper: http.DefaultTransport}
|
|
}
|