65 lines
1.9 KiB
Go
65 lines
1.9 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"time"
|
||
|
||
"github.com/agentmail/gateway/internal/db"
|
||
)
|
||
|
||
// RateLimitCheckAndRecord 原子地检查 bucket 在 window 内的事件数是否超过 limit。
|
||
// 未超过则同时记录本次事件（判断与写入在同一个事务里，防并发刷穿）。
|
||
// DB 不可用时放行（宁可放开限速也不能让用户完全无法使用）。
|
||
func RateLimitCheckAndRecord(ctx context.Context, bucket string, window time.Duration, limit int) (allowed bool, retryAfter int) {
|
||
now := time.Now()
|
||
cutoff := now.Add(-window)
|
||
|
||
// 用 IMMEDIATE 事务：SQLite 的 IMMEDIATE 会在开始时获取 RESERVED 锁，
|
||
// 防止其他写事务同时进入 COMMIT 阶段。这是 SQLite 并发写的正确方式。
|
||
tx, err := db.DB.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return true, 0 // DB 不可用 → 放行
|
||
}
|
||
defer tx.Rollback() // Commit 成功后 Rollback 是 no-op
|
||
|
||
// 清理过期记录
|
||
tx.ExecContext(ctx,
|
||
`DELETE FROM rate_limits WHERE bucket = $1 AND ts < $2`, bucket, cutoff)
|
||
|
||
// 统计当前窗口内事件数
|
||
var count int
|
||
err = tx.QueryRowContext(ctx,
|
||
`SELECT COUNT(*) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
|
||
bucket, cutoff).Scan(&count)
|
||
if err != nil {
|
||
return true, 0
|
||
}
|
||
|
||
if count >= limit {
|
||
var earliest time.Time
|
||
err = tx.QueryRowContext(ctx,
|
||
`SELECT MIN(ts) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
|
||
bucket, cutoff).Scan(&earliest)
|
||
if err == nil && !earliest.IsZero() {
|
||
retry := int(earliest.Add(window).Sub(now).Seconds()) + 1
|
||
if retry < 1 {
|
||
retry = 1
|
||
}
|
||
return false, retry
|
||
}
|
||
return false, 60
|
||
}
|
||
|
||
// 记账
|
||
tx.ExecContext(ctx,
|
||
`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`, bucket, now)
|
||
tx.Commit()
|
||
return true, 0
|
||
}
|
||
|
||
// RateLimitReset 清除指定 bucket 的所有记录（登录成功后调用）。
|
||
func RateLimitReset(ctx context.Context, bucket string) {
|
||
_, _ = db.DB.ExecContext(ctx,
|
||
`DELETE FROM rate_limits WHERE bucket = $1`, bucket)
|
||
}
|