李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
644 lines
21 KiB
Go
644 lines
21 KiB
Go
package repository
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"log"
|
||
"math/rand/v2"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
const (
|
||
billingBalanceKeyPrefix = "billing:balance:"
|
||
billingSubKeyPrefix = "billing:sub:"
|
||
billingRateLimitKeyPrefix = "apikey:rate:"
|
||
subCacheInvalidateChannel = "subscription:cache:invalidate"
|
||
billingCacheTTL = 5 * time.Minute
|
||
billingCacheJitter = 30 * time.Second
|
||
rateLimitCacheTTL = 7 * 24 * time.Hour // 7 days matches the longest window
|
||
|
||
// Rate limit window durations — must match service.RateLimitWindow* constants.
|
||
rateLimitWindow5h = 5 * time.Hour
|
||
rateLimitWindow1d = 24 * time.Hour
|
||
rateLimitWindow7d = 7 * 24 * time.Hour
|
||
)
|
||
|
||
// jitteredTTL 返回带随机抖动的 TTL，防止缓存雪崩
|
||
func jitteredTTL() time.Duration {
|
||
// 只做“减法抖动”，确保实际 TTL 不会超过 billingCacheTTL（避免上界预期被打破）。
|
||
if billingCacheJitter <= 0 {
|
||
return billingCacheTTL
|
||
}
|
||
jitter := time.Duration(rand.IntN(int(billingCacheJitter)))
|
||
return billingCacheTTL - jitter
|
||
}
|
||
|
||
// billingBalanceKey generates the Redis key for user balance cache.
|
||
func billingBalanceKey(userID int64) string {
|
||
return fmt.Sprintf("%s%d", billingBalanceKeyPrefix, userID)
|
||
}
|
||
|
||
// billingSubKey generates the Redis key for subscription cache.
|
||
func billingSubKey(userID, groupID int64) string {
|
||
return fmt.Sprintf("%s%d:%d", billingSubKeyPrefix, userID, groupID)
|
||
}
|
||
|
||
const (
|
||
subFieldStatus = "status"
|
||
subFieldExpiresAt = "expires_at"
|
||
subFieldDailyUsage = "daily_usage"
|
||
subFieldWeeklyUsage = "weekly_usage"
|
||
subFieldMonthlyUsage = "monthly_usage"
|
||
subFieldVersion = "version"
|
||
)
|
||
|
||
// billingRateLimitKey generates the Redis key for API key rate limit cache.
|
||
func billingRateLimitKey(keyID int64) string {
|
||
return fmt.Sprintf("%s%d", billingRateLimitKeyPrefix, keyID)
|
||
}
|
||
|
||
const (
|
||
rateLimitFieldUsage5h = "usage_5h"
|
||
rateLimitFieldUsage1d = "usage_1d"
|
||
rateLimitFieldUsage7d = "usage_7d"
|
||
rateLimitFieldWindow5h = "window_5h"
|
||
rateLimitFieldWindow1d = "window_1d"
|
||
rateLimitFieldWindow7d = "window_7d"
|
||
)
|
||
|
||
var (
|
||
deductBalanceScript = redis.NewScript(`
|
||
local current = redis.call('GET', KEYS[1])
|
||
if current == false then
|
||
return 0
|
||
end
|
||
local newVal = tonumber(current) - tonumber(ARGV[1])
|
||
redis.call('SET', KEYS[1], newVal)
|
||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||
return 1
|
||
`)
|
||
|
||
updateSubUsageScript = redis.NewScript(`
|
||
local exists = redis.call('EXISTS', KEYS[1])
|
||
if exists == 0 then
|
||
return 0
|
||
end
|
||
local cost = tonumber(ARGV[1])
|
||
redis.call('HINCRBYFLOAT', KEYS[1], 'daily_usage', cost)
|
||
redis.call('HINCRBYFLOAT', KEYS[1], 'weekly_usage', cost)
|
||
redis.call('HINCRBYFLOAT', KEYS[1], 'monthly_usage', cost)
|
||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||
return 1
|
||
`)
|
||
|
||
// updateRateLimitUsageScript atomically increments all three rate limit usage counters
|
||
// with window expiration checking. If a window has expired, its usage is reset to cost
|
||
// (instead of accumulated) and the window timestamp is updated, matching the DB-side
|
||
// IncrementRateLimitUsage semantics.
|
||
//
|
||
// ARGV: [1]=cost, [2]=ttl_seconds, [3]=now_unix, [4]=window_5h_seconds, [5]=window_1d_seconds, [6]=window_7d_seconds
|
||
updateRateLimitUsageScript = redis.NewScript(`
|
||
local exists = redis.call('EXISTS', KEYS[1])
|
||
if exists == 0 then
|
||
return 0
|
||
end
|
||
local cost = tonumber(ARGV[1])
|
||
local now = tonumber(ARGV[3])
|
||
local win5h = tonumber(ARGV[4])
|
||
local win1d = tonumber(ARGV[5])
|
||
local win7d = tonumber(ARGV[6])
|
||
|
||
-- Helper: check if window is expired and update usage + window accordingly
|
||
-- Returns nothing, modifies the hash in-place.
|
||
local function update_window(usage_field, window_field, window_duration)
|
||
local w = tonumber(redis.call('HGET', KEYS[1], window_field) or 0)
|
||
if w == 0 or (now - w) >= window_duration then
|
||
-- Window expired or never started: reset usage to cost, start new window
|
||
redis.call('HSET', KEYS[1], usage_field, tostring(cost))
|
||
redis.call('HSET', KEYS[1], window_field, tostring(now))
|
||
else
|
||
-- Window still valid: accumulate
|
||
redis.call('HINCRBYFLOAT', KEYS[1], usage_field, cost)
|
||
end
|
||
end
|
||
|
||
update_window('usage_5h', 'window_5h', win5h)
|
||
update_window('usage_1d', 'window_1d', win1d)
|
||
update_window('usage_7d', 'window_7d', win7d)
|
||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||
return 1
|
||
`)
|
||
)
|
||
|
||
type billingCache struct {
|
||
rdb *redis.Client
|
||
}
|
||
|
||
func NewBillingCache(rdb *redis.Client) service.BillingCache {
|
||
return &billingCache{rdb: rdb}
|
||
}
|
||
|
||
func (c *billingCache) GetUserBalance(ctx context.Context, userID int64) (float64, error) {
|
||
key := billingBalanceKey(userID)
|
||
val, err := c.rdb.Get(ctx, key).Result()
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
return strconv.ParseFloat(val, 64)
|
||
}
|
||
|
||
func (c *billingCache) SetUserBalance(ctx context.Context, userID int64, balance float64) error {
|
||
key := billingBalanceKey(userID)
|
||
return c.rdb.Set(ctx, key, balance, jitteredTTL()).Err()
|
||
}
|
||
|
||
func (c *billingCache) DeductUserBalance(ctx context.Context, userID int64, amount float64) error {
|
||
key := billingBalanceKey(userID)
|
||
_, err := deductBalanceScript.Run(ctx, c.rdb, []string{key}, amount, int(jitteredTTL().Seconds())).Result()
|
||
if err != nil && !errors.Is(err, redis.Nil) {
|
||
log.Printf("Warning: deduct balance cache failed for user %d: %v", userID, err)
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (c *billingCache) InvalidateUserBalance(ctx context.Context, userID int64) error {
|
||
key := billingBalanceKey(userID)
|
||
return c.rdb.Del(ctx, key).Err()
|
||
}
|
||
|
||
func (c *billingCache) GetSubscriptionCache(ctx context.Context, userID, groupID int64) (*service.SubscriptionCacheData, error) {
|
||
key := billingSubKey(userID, groupID)
|
||
result, err := c.rdb.HGetAll(ctx, key).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(result) == 0 {
|
||
return nil, redis.Nil
|
||
}
|
||
return c.parseSubscriptionCache(result)
|
||
}
|
||
|
||
func (c *billingCache) parseSubscriptionCache(data map[string]string) (*service.SubscriptionCacheData, error) {
|
||
result := &service.SubscriptionCacheData{}
|
||
|
||
result.Status = data[subFieldStatus]
|
||
if result.Status == "" {
|
||
return nil, errors.New("invalid cache: missing status")
|
||
}
|
||
|
||
if expiresStr, ok := data[subFieldExpiresAt]; ok {
|
||
expiresAt, err := strconv.ParseInt(expiresStr, 10, 64)
|
||
if err == nil {
|
||
result.ExpiresAt = time.Unix(expiresAt, 0)
|
||
}
|
||
}
|
||
|
||
if dailyStr, ok := data[subFieldDailyUsage]; ok {
|
||
result.DailyUsage, _ = strconv.ParseFloat(dailyStr, 64)
|
||
}
|
||
|
||
if weeklyStr, ok := data[subFieldWeeklyUsage]; ok {
|
||
result.WeeklyUsage, _ = strconv.ParseFloat(weeklyStr, 64)
|
||
}
|
||
|
||
if monthlyStr, ok := data[subFieldMonthlyUsage]; ok {
|
||
result.MonthlyUsage, _ = strconv.ParseFloat(monthlyStr, 64)
|
||
}
|
||
|
||
if versionStr, ok := data[subFieldVersion]; ok {
|
||
result.Version, _ = strconv.ParseInt(versionStr, 10, 64)
|
||
}
|
||
|
||
return result, nil
|
||
}
|
||
|
||
func (c *billingCache) SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *service.SubscriptionCacheData) error {
|
||
if data == nil {
|
||
return nil
|
||
}
|
||
|
||
key := billingSubKey(userID, groupID)
|
||
|
||
fields := map[string]any{
|
||
subFieldStatus: data.Status,
|
||
subFieldExpiresAt: data.ExpiresAt.Unix(),
|
||
subFieldDailyUsage: data.DailyUsage,
|
||
subFieldWeeklyUsage: data.WeeklyUsage,
|
||
subFieldMonthlyUsage: data.MonthlyUsage,
|
||
subFieldVersion: data.Version,
|
||
}
|
||
|
||
pipe := c.rdb.Pipeline()
|
||
pipe.HSet(ctx, key, fields)
|
||
pipe.Expire(ctx, key, jitteredTTL())
|
||
_, err := pipe.Exec(ctx)
|
||
return err
|
||
}
|
||
|
||
func (c *billingCache) UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error {
|
||
key := billingSubKey(userID, groupID)
|
||
_, err := updateSubUsageScript.Run(ctx, c.rdb, []string{key}, cost, int(jitteredTTL().Seconds())).Result()
|
||
if err != nil && !errors.Is(err, redis.Nil) {
|
||
log.Printf("Warning: update subscription usage cache failed for user %d group %d: %v", userID, groupID, err)
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (c *billingCache) InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error {
|
||
key := billingSubKey(userID, groupID)
|
||
return c.rdb.Del(ctx, key).Err()
|
||
}
|
||
|
||
func (c *billingCache) PublishSubscriptionCacheInvalidation(ctx context.Context, cacheKey string) error {
|
||
return c.rdb.Publish(ctx, subCacheInvalidateChannel, cacheKey).Err()
|
||
}
|
||
|
||
func (c *billingCache) SubscribeSubscriptionCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error {
|
||
pubsub := c.rdb.Subscribe(ctx, subCacheInvalidateChannel)
|
||
|
||
_, err := pubsub.Receive(ctx)
|
||
if err != nil {
|
||
_ = pubsub.Close()
|
||
return fmt.Errorf("subscribe to subscription cache invalidation: %w", err)
|
||
}
|
||
|
||
go func() {
|
||
defer func() {
|
||
if err := pubsub.Close(); err != nil {
|
||
log.Printf("Warning: failed to close subscription cache invalidation pubsub: %v", err)
|
||
}
|
||
}()
|
||
|
||
ch := pubsub.Channel()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case msg, ok := <-ch:
|
||
if !ok {
|
||
return
|
||
}
|
||
if msg != nil {
|
||
handler(msg.Payload)
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
|
||
return nil
|
||
}
|
||
|
||
func (c *billingCache) GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*service.APIKeyRateLimitCacheData, error) {
|
||
key := billingRateLimitKey(keyID)
|
||
result, err := c.rdb.HGetAll(ctx, key).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(result) == 0 {
|
||
return nil, redis.Nil
|
||
}
|
||
data := &service.APIKeyRateLimitCacheData{}
|
||
if v, ok := result[rateLimitFieldUsage5h]; ok {
|
||
data.Usage5h, _ = strconv.ParseFloat(v, 64)
|
||
}
|
||
if v, ok := result[rateLimitFieldUsage1d]; ok {
|
||
data.Usage1d, _ = strconv.ParseFloat(v, 64)
|
||
}
|
||
if v, ok := result[rateLimitFieldUsage7d]; ok {
|
||
data.Usage7d, _ = strconv.ParseFloat(v, 64)
|
||
}
|
||
if v, ok := result[rateLimitFieldWindow5h]; ok {
|
||
data.Window5h, _ = strconv.ParseInt(v, 10, 64)
|
||
}
|
||
if v, ok := result[rateLimitFieldWindow1d]; ok {
|
||
data.Window1d, _ = strconv.ParseInt(v, 10, 64)
|
||
}
|
||
if v, ok := result[rateLimitFieldWindow7d]; ok {
|
||
data.Window7d, _ = strconv.ParseInt(v, 10, 64)
|
||
}
|
||
return data, nil
|
||
}
|
||
|
||
func (c *billingCache) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *service.APIKeyRateLimitCacheData) error {
|
||
if data == nil {
|
||
return nil
|
||
}
|
||
key := billingRateLimitKey(keyID)
|
||
fields := map[string]any{
|
||
rateLimitFieldUsage5h: data.Usage5h,
|
||
rateLimitFieldUsage1d: data.Usage1d,
|
||
rateLimitFieldUsage7d: data.Usage7d,
|
||
rateLimitFieldWindow5h: data.Window5h,
|
||
rateLimitFieldWindow1d: data.Window1d,
|
||
rateLimitFieldWindow7d: data.Window7d,
|
||
}
|
||
pipe := c.rdb.Pipeline()
|
||
pipe.HSet(ctx, key, fields)
|
||
pipe.Expire(ctx, key, rateLimitCacheTTL)
|
||
_, err := pipe.Exec(ctx)
|
||
return err
|
||
}
|
||
|
||
func (c *billingCache) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
|
||
key := billingRateLimitKey(keyID)
|
||
now := time.Now().Unix()
|
||
_, err := updateRateLimitUsageScript.Run(ctx, c.rdb, []string{key},
|
||
cost,
|
||
int(rateLimitCacheTTL.Seconds()),
|
||
now,
|
||
int(rateLimitWindow5h.Seconds()),
|
||
int(rateLimitWindow1d.Seconds()),
|
||
int(rateLimitWindow7d.Seconds()),
|
||
).Result()
|
||
if err != nil && !errors.Is(err, redis.Nil) {
|
||
log.Printf("Warning: update rate limit usage cache failed for api key %d: %v", keyID, err)
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (c *billingCache) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
|
||
key := billingRateLimitKey(keyID)
|
||
return c.rdb.Del(ctx, key).Err()
|
||
}
|
||
|
||
// ============================================
|
||
// user × platform quota 缓存
|
||
// ============================================
|
||
|
||
// userPlatformQuotaCacheKey 构造 Redis key
|
||
func userPlatformQuotaCacheKey(userID int64, platform string) string {
|
||
return fmt.Sprintf("billing:user_platform_quota:%d:%s", userID, platform)
|
||
}
|
||
|
||
// parseUserPlatformQuotaHash 将 Redis HGETALL 返回的 map[string]string 反序列化为
|
||
// *service.UserPlatformQuotaCacheEntry。空 map（key 不存在）返回 nil。
|
||
// GetUserPlatformQuotaCache 和 BatchGetUserPlatformQuotaCache 共用此函数，确保解析逻辑一致。
|
||
func parseUserPlatformQuotaHash(m map[string]string) *service.UserPlatformQuotaCacheEntry {
|
||
if len(m) == 0 {
|
||
return nil
|
||
}
|
||
parseFloat := func(s string) float64 {
|
||
if s == "" {
|
||
return 0
|
||
}
|
||
f, err := strconv.ParseFloat(s, 64)
|
||
if err != nil {
|
||
log.Printf("billing_cache: corrupt quota usage field %q (using 0): %v", s, err)
|
||
return 0
|
||
}
|
||
return f
|
||
}
|
||
parseFloatPtr := func(s string) *float64 {
|
||
if s == "" {
|
||
return nil
|
||
}
|
||
f, err := strconv.ParseFloat(s, 64)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
return &f
|
||
}
|
||
parseTimePtr := func(s string) *time.Time {
|
||
if s == "" {
|
||
return nil
|
||
}
|
||
n, err := strconv.ParseInt(s, 10, 64)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
t := time.Unix(n, 0).UTC()
|
||
return &t
|
||
}
|
||
parseInt64 := func(s string) int64 {
|
||
n, _ := strconv.ParseInt(s, 10, 64)
|
||
return n
|
||
}
|
||
return &service.UserPlatformQuotaCacheEntry{
|
||
DailyUsageUSD: parseFloat(m["daily_usage"]),
|
||
WeeklyUsageUSD: parseFloat(m["weekly_usage"]),
|
||
MonthlyUsageUSD: parseFloat(m["monthly_usage"]),
|
||
Version: parseInt64(m["version"]),
|
||
SchemaVersion: parseInt64(m["schema_version"]),
|
||
DailyLimitUSD: parseFloatPtr(m["daily_limit"]),
|
||
WeeklyLimitUSD: parseFloatPtr(m["weekly_limit"]),
|
||
MonthlyLimitUSD: parseFloatPtr(m["monthly_limit"]),
|
||
DailyWindowStart: parseTimePtr(m["daily_window_start"]),
|
||
WeeklyWindowStart: parseTimePtr(m["weekly_window_start"]),
|
||
MonthlyWindowStart: parseTimePtr(m["monthly_window_start"]),
|
||
}
|
||
}
|
||
|
||
func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) {
|
||
key := userPlatformQuotaCacheKey(userID, platform)
|
||
m, err := c.rdb.HGetAll(ctx, key).Result()
|
||
if err != nil {
|
||
return nil, false, err
|
||
}
|
||
entry := parseUserPlatformQuotaHash(m)
|
||
if entry == nil {
|
||
// 空 map → key 不存在 → MISS
|
||
return nil, false, nil
|
||
}
|
||
return entry, true, nil
|
||
}
|
||
|
||
func (c *billingCache) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error {
|
||
if entry == nil {
|
||
return nil
|
||
}
|
||
key := userPlatformQuotaCacheKey(userID, platform)
|
||
pipe := c.rdb.TxPipeline()
|
||
|
||
// 浮点可空字段：nil → 空字符串（读取时 parseFloatPtr 返回 nil，表示无限额）
|
||
fmtFloatPtr := func(p *float64) string {
|
||
if p == nil {
|
||
return ""
|
||
}
|
||
return strconv.FormatFloat(*p, 'f', -1, 64)
|
||
}
|
||
// time.Time 可空字段：nil → 空字符串；有值 → unix 秒
|
||
fmtTimePtr := func(p *time.Time) string {
|
||
if p == nil {
|
||
return ""
|
||
}
|
||
return strconv.FormatInt(p.Unix(), 10)
|
||
}
|
||
|
||
pipe.HSet(ctx, key,
|
||
"daily_usage", entry.DailyUsageUSD,
|
||
"weekly_usage", entry.WeeklyUsageUSD,
|
||
"monthly_usage", entry.MonthlyUsageUSD,
|
||
"version", entry.Version,
|
||
"schema_version", entry.SchemaVersion,
|
||
"daily_limit", fmtFloatPtr(entry.DailyLimitUSD),
|
||
"weekly_limit", fmtFloatPtr(entry.WeeklyLimitUSD),
|
||
"monthly_limit", fmtFloatPtr(entry.MonthlyLimitUSD),
|
||
"daily_window_start", fmtTimePtr(entry.DailyWindowStart),
|
||
"weekly_window_start", fmtTimePtr(entry.WeeklyWindowStart),
|
||
"monthly_window_start", fmtTimePtr(entry.MonthlyWindowStart),
|
||
)
|
||
pipe.Expire(ctx, key, ttl)
|
||
_, err := pipe.Exec(ctx)
|
||
return err
|
||
}
|
||
|
||
func (c *billingCache) DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error {
|
||
return c.rdb.Del(ctx, userPlatformQuotaCacheKey(userID, platform)).Err()
|
||
}
|
||
|
||
// updateUserPlatformQuotaUsageScript 缓存累加：EXISTS + schema_version 双重守卫。
|
||
// 旧版 entry（schema_version != ARGV[3]，包括缺字段的 0 值）不参与累加，由上层走 DB fallback 后
|
||
// SetCache 重建为新版 entry —— 若此处仍累加，上层覆盖时会丢失这部分增量，导致 Redis usage 比真实偏小。
|
||
// key 不存在同样跳过（由下次 SetCache 重建）。
|
||
// KEYS[1] = hash key
|
||
// KEYS[2] = 脏集 key（dirty set）
|
||
// ARGV[1] = cost (string float)
|
||
// ARGV[2] = ttl seconds
|
||
// ARGV[3] = expected schema_version (Go 侧 UserPlatformQuotaCacheSchemaV1)
|
||
// ARGV[4] = dirty set member（空串则不 SADD）
|
||
// ARGV[5] = 脏集兜底 TTL 秒
|
||
const updateUserPlatformQuotaUsageScript = `
|
||
if redis.call("EXISTS", KEYS[1]) == 0 then
|
||
return 0
|
||
end
|
||
local ver = redis.call("HGET", KEYS[1], "schema_version")
|
||
if ver == false or tonumber(ver) ~= tonumber(ARGV[3]) then
|
||
return 0
|
||
end
|
||
redis.call("HINCRBYFLOAT", KEYS[1], "daily_usage", ARGV[1])
|
||
redis.call("HINCRBYFLOAT", KEYS[1], "weekly_usage", ARGV[1])
|
||
redis.call("HINCRBYFLOAT", KEYS[1], "monthly_usage", ARGV[1])
|
||
redis.call("HINCRBY", KEYS[1], "version", 1)
|
||
redis.call("EXPIRE", KEYS[1], ARGV[2])
|
||
if ARGV[4] ~= "" then
|
||
redis.call("SADD", KEYS[2], ARGV[4])
|
||
redis.call("EXPIRE", KEYS[2], ARGV[5])
|
||
end
|
||
return 1
|
||
`
|
||
|
||
// userPlatformQuotaDirtySetKey 返回脏集（dirty set）的 Redis key。
|
||
// 使用与 userPlatformQuotaCacheKey 相同的前缀 "billing:"。
|
||
func userPlatformQuotaDirtySetKey() string { return "billing:" + "upq:dirty" }
|
||
|
||
// userPlatformQuotaDirtyTTLSeconds 脏集兜底 TTL（秒）：初始 SADD（Lua）与 Readd 共用，
|
||
// 确保 flusher 长期停摆时脏集最终过期；正常运行因持续 SADD 不断续期。
|
||
const userPlatformQuotaDirtyTTLSeconds = 86400
|
||
|
||
// userPlatformQuotaDirtyMember 构造脏集成员字符串 "userID:platform"。
|
||
func userPlatformQuotaDirtyMember(userID int64, platform string) string {
|
||
return strconv.FormatInt(userID, 10) + ":" + platform
|
||
}
|
||
|
||
func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error {
|
||
member := ""
|
||
if markDirty {
|
||
member = userPlatformQuotaDirtyMember(userID, platform)
|
||
}
|
||
_, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript,
|
||
[]string{userPlatformQuotaCacheKey(userID, platform), userPlatformQuotaDirtySetKey()},
|
||
strconv.FormatFloat(cost, 'f', -1, 64),
|
||
int(ttl.Seconds()),
|
||
service.UserPlatformQuotaCacheSchemaV1,
|
||
member,
|
||
userPlatformQuotaDirtyTTLSeconds,
|
||
).Result()
|
||
if err != nil && !errors.Is(err, redis.Nil) {
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// parseUserPlatformQuotaDirtyMember 将脏集成员字符串 "userID:platform" 解析为
|
||
// service.UserPlatformQuotaKey。解析失败返回 ok=false。
|
||
func parseUserPlatformQuotaDirtyMember(m string) (service.UserPlatformQuotaKey, bool) {
|
||
parts := strings.SplitN(m, ":", 2)
|
||
if len(parts) != 2 {
|
||
return service.UserPlatformQuotaKey{}, false
|
||
}
|
||
uid, err := strconv.ParseInt(parts[0], 10, 64)
|
||
if err != nil {
|
||
return service.UserPlatformQuotaKey{}, false
|
||
}
|
||
return service.UserPlatformQuotaKey{UserID: uid, Platform: parts[1]}, true
|
||
}
|
||
|
||
// PopDirtyUserPlatformQuotaKeys 从脏集随机弹出最多 n 个 key。
|
||
// 脏集为空时返回 (nil, nil)。
|
||
func (c *billingCache) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]service.UserPlatformQuotaKey, error) {
|
||
members, err := c.rdb.SPopN(ctx, userPlatformQuotaDirtySetKey(), int64(n)).Result()
|
||
if err != nil {
|
||
if errors.Is(err, redis.Nil) {
|
||
return nil, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
keys := make([]service.UserPlatformQuotaKey, 0, len(members))
|
||
for _, m := range members {
|
||
k, ok := parseUserPlatformQuotaDirtyMember(m)
|
||
if !ok {
|
||
log.Printf("billing_cache: skipping invalid dirty member %q", m)
|
||
continue
|
||
}
|
||
keys = append(keys, k)
|
||
}
|
||
return keys, nil
|
||
}
|
||
|
||
// ReaddDirtyUserPlatformQuotaKeys 将 keys 重新加入脏集（flush 失败时回填）。
|
||
// 通过 pipeline 同时执行 SAdd + Expire，确保 Readd 后脏集具有兜底 TTL。
|
||
// 空切片时直接返回 nil。
|
||
func (c *billingCache) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []service.UserPlatformQuotaKey) error {
|
||
if len(keys) == 0 {
|
||
return nil
|
||
}
|
||
dirtyKey := userPlatformQuotaDirtySetKey()
|
||
members := make([]any, len(keys))
|
||
for i, k := range keys {
|
||
members[i] = userPlatformQuotaDirtyMember(k.UserID, k.Platform)
|
||
}
|
||
pipe := c.rdb.Pipeline()
|
||
pipe.SAdd(ctx, dirtyKey, members...)
|
||
pipe.Expire(ctx, dirtyKey, userPlatformQuotaDirtyTTLSeconds*time.Second)
|
||
_, err := pipe.Exec(ctx)
|
||
return err
|
||
}
|
||
|
||
// BatchGetUserPlatformQuotaCache 通过 Pipeline 批量 HGETALL 获取多个 user×platform 的
|
||
// quota cache。返回切片与 keys 顺序、长度对齐；MISS 或解析失败位置返回 nil。
|
||
func (c *billingCache) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []service.UserPlatformQuotaKey) ([]*service.UserPlatformQuotaCacheEntry, error) {
|
||
if len(keys) == 0 {
|
||
return nil, nil
|
||
}
|
||
pipe := c.rdb.Pipeline()
|
||
cmds := make([]*redis.MapStringStringCmd, len(keys))
|
||
for i, k := range keys {
|
||
cmds[i] = pipe.HGetAll(ctx, userPlatformQuotaCacheKey(k.UserID, k.Platform))
|
||
}
|
||
if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
|
||
return nil, err
|
||
}
|
||
results := make([]*service.UserPlatformQuotaCacheEntry, len(keys))
|
||
for i, cmd := range cmds {
|
||
m, err := cmd.Result()
|
||
if err != nil {
|
||
if !errors.Is(err, redis.Nil) {
|
||
log.Printf("billing_cache: BatchGet HGETALL cmd[%d] failed: %v (skip, self-heal)", i, err)
|
||
}
|
||
// 单个命令失败 → 对应位置 nil，继续
|
||
continue
|
||
}
|
||
results[i] = parseUserPlatformQuotaHash(m)
|
||
}
|
||
return results, nil
|
||
}
|