diff --git a/internal/adapter/out/redis/quota.go b/internal/adapter/out/redis/quota.go index 35e51c0..8c38718 100644 --- a/internal/adapter/out/redis/quota.go +++ b/internal/adapter/out/redis/quota.go @@ -31,10 +31,36 @@ func (q *RoomQuota) TakePick(ctx context.Context, roomID string, burst int, refi } func (q *RoomQuota) take(ctx context.Context, roomID, bucket, what string, burst int, refill time.Duration) (time.Duration, error) { - res, err := q.limiter.Allow(ctx, "room:"+roomID+":"+bucket, redis_rate.Limit{Rate: 1, Burst: burst, Period: refill}) + wait, err := take(ctx, q.limiter, "room:"+roomID+":"+bucket, burst, refill) if err != nil { return 0, fmt.Errorf("取房间 %s 的%s额度:%w", roomID, what, err) } + return wait, nil +} + +var _ port.AccountQuota = (*AccountQuota)(nil) + +type AccountQuota struct { + limiter *redis_rate.Limiter +} + +func NewAccountQuota(client *goredis.Client) *AccountQuota { + return &AccountQuota{limiter: redis_rate.NewLimiter(client)} +} + +func (q *AccountQuota) Take(ctx context.Context, key string, burst int, refill time.Duration) (time.Duration, error) { + wait, err := take(ctx, q.limiter, "account:"+key, burst, refill) + if err != nil { + return 0, fmt.Errorf("取 %s 的额度:%w", key, err) + } + return wait, nil +} + +func take(ctx context.Context, limiter *redis_rate.Limiter, key string, burst int, refill time.Duration) (time.Duration, error) { + res, err := limiter.Allow(ctx, key, redis_rate.Limit{Rate: 1, Burst: burst, Period: refill}) + if err != nil { + return 0, err + } if res.Allowed > 0 { return 0, nil } diff --git a/internal/adapter/out/redis/quota_test.go b/internal/adapter/out/redis/quota_test.go index f9281a7..f517cef 100644 --- a/internal/adapter/out/redis/quota_test.go +++ b/internal/adapter/out/redis/quota_test.go @@ -43,3 +43,26 @@ func TestTakePick(t *testing.T) { t.Errorf("另一个房间 = %v, %v,点歌额度按房间计算", wait, err) } } + +func TestAccountQuota(t *testing.T) { + client := newClient(t) + quota := NewAccountQuota(client) + for i := range 2 { + if wait, err := quota.Take(t.Context(), "login:ip:203.0.113.7", 2, time.Minute); err != nil || wait != 0 { + t.Fatalf("第 %d 次 = %v, %v,突发额度之内应直接取到", i+1, wait, err) + } + } + wait, err := quota.Take(t.Context(), "login:ip:203.0.113.7", 2, time.Minute) + if err != nil { + t.Fatal(err) + } + if wait < 59*time.Second || wait > time.Minute { + t.Errorf("额度用完后要等 %v,期望约 1 分钟", wait) + } + if wait, err := quota.Take(t.Context(), "login:user:alice", 2, time.Minute); err != nil || wait != 0 { + t.Errorf("另一个键 = %v, %v,额度按键计算", wait, err) + } + if wait, err := NewRoomQuota(client).TakeSearch(t.Context(), "login:ip:203.0.113.7", 1, time.Minute); err != nil || wait != 0 { + t.Errorf("房间额度 = %v, %v,账号额度的键不能与房间额度的键重合", wait, err) + } +} diff --git a/internal/app/port/quota.go b/internal/app/port/quota.go index 124f52f..a18f1ea 100644 --- a/internal/app/port/quota.go +++ b/internal/app/port/quota.go @@ -13,3 +13,9 @@ type RoomQuota interface { // 与 TakeSearch 相同,取的是点歌额度。 TakePick(ctx context.Context, roomID string, burst int, refill time.Duration) (time.Duration, error) } + +// 注册、登录、修改密码的额度,按调用方拼出的键计数,所有实例共用。 +type AccountQuota interface { + // 与 RoomQuota.TakeSearch 相同:取到时返回 0,额度用完时返回还要等待的时长。 + Take(ctx context.Context, key string, burst int, refill time.Duration) (time.Duration, error) +}