Jeffrey Smith
3d4228f868
All checks were successful
Co-authored-by: Jeffrey Smith <jasafpro@gmail.com> Co-committed-by: Jeffrey Smith <jasafpro@gmail.com>
36 lines
1.0 KiB
Go
36 lines
1.0 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// RateLimitStore manages distributed rate limit counters in Postgres.
|
|
type RateLimitStore struct{}
|
|
|
|
func NewRateLimitStore() *RateLimitStore { return &RateLimitStore{} }
|
|
|
|
func (s *RateLimitStore) Allow(ctx context.Context, key string, rate float64, burst int) (bool, error) {
|
|
// Upsert the current-second window counter and return the new token count.
|
|
var tokens float64
|
|
err := DB.QueryRowContext(ctx, `
|
|
INSERT INTO rate_limit_counters (key, window_start, tokens)
|
|
VALUES ($1, date_trunc('second', NOW()), 1)
|
|
ON CONFLICT (key, window_start)
|
|
DO UPDATE SET tokens = rate_limit_counters.tokens + 1
|
|
RETURNING tokens`, key).Scan(&tokens)
|
|
if err != nil {
|
|
return false, fmt.Errorf("rate limit upsert: %w", err)
|
|
}
|
|
return tokens <= float64(burst), nil
|
|
}
|
|
|
|
func (s *RateLimitStore) Cleanup(ctx context.Context, maxAge time.Duration) error {
|
|
_, err := DB.ExecContext(ctx, `
|
|
DELETE FROM rate_limit_counters
|
|
WHERE window_start < NOW() - $1 * INTERVAL '1 second'`,
|
|
int(maxAge.Seconds()))
|
|
return err
|
|
}
|