package repository import ( "context" "errors" "fmt" "log" "strconv" "time" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" ) const ( billingBalanceKeyPrefix = "billing:balance:" billingSubKeyPrefix = "billing:sub:" billingCacheTTL = 5 * time.Minute ) // 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" ) 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 `) ) 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, billingCacheTTL).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(billingCacheTTL.Seconds())).Result() if err != nil && !errors.Is(err, redis.Nil) { log.Printf("Warning: deduct balance cache failed for user %d: %v", userID, 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, billingCacheTTL) _, 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(billingCacheTTL.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 nil } func (c *billingCache) InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error { key := billingSubKey(userID, groupID) return c.rdb.Del(ctx, key).Err() }