Skip to content

Commit 777be05

Browse files
author
QTom
committed
feat: add RPMCache interface and Redis implementation with Lua scripts
1 parent 0bb3e4a commit 777be05

2 files changed

Lines changed: 105 additions & 0 deletions

File tree

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package repository
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"github.com/redis/go-redis/v9"
8+
)
9+
10+
const rpmKeyPrefix = "rpm:"
11+
12+
// Lua scripts use Redis TIME for server-side minute key calculation
13+
var rpmIncrScript = redis.NewScript(`
14+
local timeResult = redis.call('TIME')
15+
local minuteKey = math.floor(tonumber(timeResult[1]) / 60)
16+
local key = ARGV[1] .. ':' .. minuteKey
17+
local count = redis.call('INCR', key)
18+
if count == 1 then
19+
redis.call('EXPIRE', key, 120)
20+
end
21+
return count
22+
`)
23+
24+
var rpmGetScript = redis.NewScript(`
25+
local timeResult = redis.call('TIME')
26+
local minuteKey = math.floor(tonumber(timeResult[1]) / 60)
27+
local key = ARGV[1] .. ':' .. minuteKey
28+
local count = redis.call('GET', key)
29+
if count == false then
30+
return 0
31+
end
32+
return tonumber(count)
33+
`)
34+
35+
type RPMCacheImpl struct {
36+
rdb *redis.Client
37+
}
38+
39+
func NewRPMCache(rdb *redis.Client) *RPMCacheImpl {
40+
return &RPMCacheImpl{rdb: rdb}
41+
}
42+
43+
func rpmKeyBase(accountID int64) string {
44+
return fmt.Sprintf("%s%d", rpmKeyPrefix, accountID)
45+
}
46+
47+
func (c *RPMCacheImpl) IncrementRPM(ctx context.Context, accountID int64) (int, error) {
48+
result, err := rpmIncrScript.Run(ctx, c.rdb, nil, rpmKeyBase(accountID)).Int()
49+
if err != nil {
50+
return 0, fmt.Errorf("rpm increment: %w", err)
51+
}
52+
return result, nil
53+
}
54+
55+
func (c *RPMCacheImpl) GetRPM(ctx context.Context, accountID int64) (int, error) {
56+
result, err := rpmGetScript.Run(ctx, c.rdb, nil, rpmKeyBase(accountID)).Int()
57+
if err != nil {
58+
return 0, fmt.Errorf("rpm get: %w", err)
59+
}
60+
return result, nil
61+
}
62+
63+
func (c *RPMCacheImpl) GetRPMBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error) {
64+
if len(accountIDs) == 0 {
65+
return map[int64]int{}, nil
66+
}
67+
68+
pipe := c.rdb.Pipeline()
69+
cmds := make(map[int64]*redis.Cmd, len(accountIDs))
70+
for _, id := range accountIDs {
71+
cmds[id] = rpmGetScript.Run(ctx, pipe, nil, rpmKeyBase(id))
72+
}
73+
74+
_, err := pipe.Exec(ctx)
75+
if err != nil && err != redis.Nil {
76+
return nil, fmt.Errorf("rpm batch get: %w", err)
77+
}
78+
79+
result := make(map[int64]int, len(accountIDs))
80+
for id, cmd := range cmds {
81+
if val, err := cmd.Int(); err == nil {
82+
result[id] = val
83+
} else {
84+
result[id] = 0
85+
}
86+
}
87+
return result, nil
88+
}
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package service
2+
3+
import "context"
4+
5+
// RPMCache RPM 计数器缓存接口
6+
// 用于 Anthropic OAuth/SetupToken 账号的每分钟请求数限制
7+
type RPMCache interface {
8+
// IncrementRPM 原子递增并返回当前分钟的计数
9+
// 使用 Redis 服务器时间确定 minute key,避免多实例时钟偏差
10+
IncrementRPM(ctx context.Context, accountID int64) (count int, err error)
11+
12+
// GetRPM 获取当前分钟的 RPM 计数
13+
GetRPM(ctx context.Context, accountID int64) (count int, err error)
14+
15+
// GetRPMBatch 批量获取多个账号的 RPM 计数(使用 Pipeline)
16+
GetRPMBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error)
17+
}

0 commit comments

Comments
 (0)