@@ -2,78 +2,130 @@ package repository
22
33import (
44 "context"
5+ "errors"
56 "fmt"
7+ "strconv"
8+ "time"
69
710 "github.com/Wei-Shaw/sub2api/internal/service"
811 "github.com/redis/go-redis/v9"
912)
1013
11- const rpmKeyPrefix = "rpm:"
12-
13- // Lua scripts use Redis TIME for server-side minute key calculation
14- var rpmIncrScript = redis .NewScript (`
15- local timeResult = redis.call('TIME')
16- local minuteKey = math.floor(tonumber(timeResult[1]) / 60)
17- local key = ARGV[1] .. ':' .. minuteKey
18- local count = redis.call('INCR', key)
19- if count == 1 then
20- redis.call('EXPIRE', key, 120)
21- end
22- return count
23- ` )
24-
25- var rpmGetScript = redis .NewScript (`
26- local timeResult = redis.call('TIME')
27- local minuteKey = math.floor(tonumber(timeResult[1]) / 60)
28- local key = ARGV[1] .. ':' .. minuteKey
29- local count = redis.call('GET', key)
30- if count == false then
31- return 0
32- end
33- return tonumber(count)
34- ` )
14+ // RPM 计数器缓存常量定义
15+ //
16+ // 设计说明：
17+ // 使用 Redis 简单计数器跟踪每个账号每分钟的请求数：
18+ // - Key: rpm:{accountID}:{minuteTimestamp}
19+ // - Value: 当前分钟内的请求计数
20+ // - TTL: 120 秒（覆盖当前分钟 + 一定冗余）
21+ //
22+ // 使用 TxPipeline（MULTI/EXEC）执行 INCR + EXPIRE，保证原子性且兼容 Redis Cluster。
23+ // 通过 rdb.Time() 获取服务端时间，避免多实例时钟不同步。
24+ //
25+ // 设计决策：
26+ // - TxPipeline vs Pipeline：Pipeline 仅合并发送但不保证原子，TxPipeline 使用 MULTI/EXEC 事务保证原子执行。
27+ // - rdb.Time() 单独调用：Pipeline/TxPipeline 中无法引用前一命令的结果，因此 TIME 必须单独调用（2 RTT）。
28+ // Lua 脚本可以做到 1 RTT，但在 Redis Cluster 中动态拼接 key 存在 CROSSSLOT 风险，选择安全性优先。
29+ const (
30+ // RPM 计数器键前缀
31+ // 格式: rpm:{accountID}:{minuteTimestamp}
32+ rpmKeyPrefix = "rpm:"
3533
34+ // RPM 计数器 TTL（120 秒，覆盖当前分钟窗口 + 冗余）
35+ rpmKeyTTL = 120 * time .Second
36+ )
37+
38+ // RPMCacheImpl RPM 计数器缓存 Redis 实现
3639type RPMCacheImpl struct {
3740 rdb * redis.Client
3841}
3942
43+ // NewRPMCache 创建 RPM 计数器缓存
4044func NewRPMCache (rdb * redis.Client ) service.RPMCache {
4145 return & RPMCacheImpl {rdb : rdb }
4246}
4347
44- func rpmKeyBase (accountID int64 ) string {
45- return fmt .Sprintf ("%s%d" , rpmKeyPrefix , accountID )
48+ // currentMinuteKey 获取当前分钟的完整 Redis key
49+ // 使用 rdb.Time() 获取 Redis 服务端时间，避免多实例时钟偏差
50+ func (c * RPMCacheImpl ) currentMinuteKey (ctx context.Context , accountID int64 ) (string , error ) {
51+ serverTime , err := c .rdb .Time (ctx ).Result ()
52+ if err != nil {
53+ return "" , fmt .Errorf ("redis TIME: %w" , err )
54+ }
55+ minuteTS := serverTime .Unix () / 60
56+ return fmt .Sprintf ("%s%d:%d" , rpmKeyPrefix , accountID , minuteTS ), nil
57+ }
58+
59+ // currentMinuteSuffix 获取当前分钟时间戳后缀（供批量操作使用）
60+ // 使用 rdb.Time() 获取 Redis 服务端时间
61+ func (c * RPMCacheImpl ) currentMinuteSuffix (ctx context.Context ) (string , error ) {
62+ serverTime , err := c .rdb .Time (ctx ).Result ()
63+ if err != nil {
64+ return "" , fmt .Errorf ("redis TIME: %w" , err )
65+ }
66+ minuteTS := serverTime .Unix () / 60
67+ return strconv .FormatInt (minuteTS , 10 ), nil
4668}
4769
70+ // IncrementRPM 原子递增并返回当前分钟的计数
71+ // 使用 TxPipeline (MULTI/EXEC) 执行 INCR + EXPIRE，保证原子性且兼容 Redis Cluster
4872func (c * RPMCacheImpl ) IncrementRPM (ctx context.Context , accountID int64 ) (int , error ) {
49- result , err := rpmIncrScript . Run (ctx , c . rdb , nil , rpmKeyBase ( accountID )). Int ( )
73+ key , err := c . currentMinuteKey (ctx , accountID )
5074 if err != nil {
5175 return 0 , fmt .Errorf ("rpm increment: %w" , err )
5276 }
53- return result , nil
77+
78+ // 使用 TxPipeline (MULTI/EXEC) 保证 INCR + EXPIRE 原子执行
79+ // EXPIRE 幂等，每次都设置不影响正确性
80+ pipe := c .rdb .TxPipeline ()
81+ incrCmd := pipe .Incr (ctx , key )
82+ pipe .Expire (ctx , key , rpmKeyTTL )
83+
84+ if _ , err := pipe .Exec (ctx ); err != nil {
85+ return 0 , fmt .Errorf ("rpm increment: %w" , err )
86+ }
87+
88+ return int (incrCmd .Val ()), nil
5489}
5590
91+ // GetRPM 获取当前分钟的 RPM 计数
5692func (c * RPMCacheImpl ) GetRPM (ctx context.Context , accountID int64 ) (int , error ) {
57- result , err := rpmGetScript . Run (ctx , c . rdb , nil , rpmKeyBase ( accountID )). Int ( )
93+ key , err := c . currentMinuteKey (ctx , accountID )
5894 if err != nil {
5995 return 0 , fmt .Errorf ("rpm get: %w" , err )
6096 }
61- return result , nil
97+
98+ val , err := c .rdb .Get (ctx , key ).Int ()
99+ if errors .Is (err , redis .Nil ) {
100+ return 0 , nil // 当前分钟无记录
101+ }
102+ if err != nil {
103+ return 0 , fmt .Errorf ("rpm get: %w" , err )
104+ }
105+ return val , nil
62106}
63107
108+ // GetRPMBatch 批量获取多个账号的 RPM 计数（使用 Pipeline）
64109func (c * RPMCacheImpl ) GetRPMBatch (ctx context.Context , accountIDs []int64 ) (map [int64 ]int , error ) {
65110 if len (accountIDs ) == 0 {
66111 return map [int64 ]int {}, nil
67112 }
68113
114+ // 获取当前分钟后缀
115+ minuteSuffix , err := c .currentMinuteSuffix (ctx )
116+ if err != nil {
117+ return nil , fmt .Errorf ("rpm batch get: %w" , err )
118+ }
119+
120+ // 使用 Pipeline 批量 GET
69121 pipe := c .rdb .Pipeline ()
70- cmds := make (map [int64 ]* redis.Cmd , len (accountIDs ))
122+ cmds := make (map [int64 ]* redis.StringCmd , len (accountIDs ))
71123 for _ , id := range accountIDs {
72- cmds [id ] = rpmGetScript .Run (ctx , pipe , nil , rpmKeyBase (id ))
124+ key := fmt .Sprintf ("%s%d:%s" , rpmKeyPrefix , id , minuteSuffix )
125+ cmds [id ] = pipe .Get (ctx , key )
73126 }
74127
75- _ , err := pipe .Exec (ctx )
76- if err != nil && err != redis .Nil {
128+ if _ , err := pipe .Exec (ctx ); err != nil && ! errors .Is (err , redis .Nil ) {
77129 return nil , fmt .Errorf ("rpm batch get: %w" , err )
78130 }
79131
0 commit comments