zt-closet-backend/internal/redisdao/quota.go

package redisdao

import (
	"context"
	"fmt"
	"time"

	"git.lzcwy.cn/zhengchiliu/zt-closet-backend/global"
	"git.lzcwy.cn/zhengchiliu/zt-closet-backend/internal/redismodel"
	"github.com/redis/go-redis/v9"
)

// quotaDailyTTL 每日计数 key 的过期时间（跨天，保证到期自然过期）
const quotaDailyTTL = 48 * time.Hour

// quotaDailyIncrScript 首次自增时设置过期时间，保证 INCR 与 EXPIRE 原子执行
var quotaDailyIncrScript = redis.NewScript(`
local v = redis.call('INCR', KEYS[1])
if v == 1 then
  redis.call('EXPIRE', KEYS[1], ARGV[1])
end
return v
`)

// quotaDailyIncr 每日计数 +1
func quotaDailyIncr(ctx context.Context, key string) (int64, error) {
	return quotaDailyIncrScript.Run(ctx, global.RedisEngine, []string{key}, int(quotaDailyTTL.Seconds())).Int64()
}

// quotaDailyValue 读取每日计数，key 不存在时返回 0
func quotaDailyValue(ctx context.Context, key string) (int64, error) {
	v, err := global.RedisEngine.Get(ctx, key).Int64()
	if err == redis.Nil {
		return 0, nil
	}
	return v, err
}

// QuotaDailyAIIncr 当日 AI 推荐次数 +1
func QuotaDailyAIIncr(ctx context.Context, userID int64, date int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.QuotaDailyAIIncr")
	defer span.End()

	val, err := quotaDailyIncr(ctx, fmt.Sprintf(redismodel.KeyQuotaDailyAI, userID, date))
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// QuotaDailyOutfitShareIncr 当日社区分享（搭配分享）次数 +1
func QuotaDailyOutfitShareIncr(ctx context.Context, userID int64, date int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.QuotaDailyOutfitShareIncr")
	defer span.End()

	val, err := quotaDailyIncr(ctx, fmt.Sprintf(redismodel.KeyQuotaDailyOutfitShare, userID, date))
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// QuotaDailyClosetShareIncr 当日衣柜分享次数 +1
func QuotaDailyClosetShareIncr(ctx context.Context, userID int64, date int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.QuotaDailyClosetShareIncr")
	defer span.End()

	val, err := quotaDailyIncr(ctx, fmt.Sprintf(redismodel.KeyQuotaDailyClosetShare, userID, date))
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// QuotaDailyAIValue 当日 AI 推荐次数
func QuotaDailyAIValue(ctx context.Context, userID int64, date int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.QuotaDailyAIValue")
	defer span.End()

	val, err := quotaDailyValue(ctx, fmt.Sprintf(redismodel.KeyQuotaDailyAI, userID, date))
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// QuotaDailyOutfitShareValue 当日社区分享（搭配分享）次数
func QuotaDailyOutfitShareValue(ctx context.Context, userID int64, date int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.QuotaDailyOutfitShareValue")
	defer span.End()

	val, err := quotaDailyValue(ctx, fmt.Sprintf(redismodel.KeyQuotaDailyOutfitShare, userID, date))
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// QuotaDailyClosetShareValue 当日衣柜分享次数
func QuotaDailyClosetShareValue(ctx context.Context, userID int64, date int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.QuotaDailyClosetShareValue")
	defer span.End()

	val, err := quotaDailyValue(ctx, fmt.Sprintf(redismodel.KeyQuotaDailyClosetShare, userID, date))
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// ---- 通用每日额度原子限流（带限额） ----

// dailyQuotaConsumeScript 带限额的原子自增：
//
//	KEYS[1]=当日计数 key；ARGV[1]=每日上限；ARGV[2]=消费数量；ARGV[3]=TTL 秒。
//
// 当前计数 + n 超过上限返回 -1；否则 INCRBY 并在首次写入时设置 TTL，返回消费后的当日计数。
var dailyQuotaConsumeScript = redis.NewScript(`
local v = tonumber(redis.call('GET', KEYS[1]) or '0')
local baseline = tonumber(ARGV[4])
if baseline > v then v = baseline; redis.call('SET', KEYS[1], v); redis.call('EXPIRE', KEYS[1], ARGV[3]) end
local limit = tonumber(ARGV[1])
local n = tonumber(ARGV[2])
if v + n > limit then
  return -1
end
local r = redis.call('INCRBY', KEYS[1], n)
if v == 0 then
  redis.call('EXPIRE', KEYS[1], ARGV[3])
end
return r
`)

// dailyQuotaRefundScript 回退当日计数（DECRBY 且不小于 0）。
// ARGV[1]=回退数量。
var dailyQuotaRefundScript = redis.NewScript(`
local v = tonumber(redis.call('GET', KEYS[1]) or '0')
local n = tonumber(ARGV[1])
if v <= 0 then
  return 0
end
if n > v then
  n = v
end
return redis.call('DECRBY', KEYS[1], n)
`)

// DailyQuotaConsume 原子消费当日额度。
// 返回 (消费后计数, nil)；额度不足返回 (-1, nil)；Redis 异常返回 (0, err)。
func DailyQuotaConsume(ctx context.Context, key string, limit int64, n int64, baseline int64) (int64, error) {
	ctx, span := global.Tracer.Start(ctx, "redisdao.DailyQuotaConsume")
	defer span.End()

	if limit <= 0 || n <= 0 {
		return -1, nil
	}
	val, err := dailyQuotaConsumeScript.Run(ctx, global.RedisEngine, []string{key}, limit, n, int(quotaDailyTTL.Seconds()), baseline).Int64()
	if err != nil {
		span.RecordError(err)
		return 0, err
	}
	return val, nil
}

// DailyQuotaRefund 回退当日计数（保证不小于 0）。失败仅记录，调用方按 best-effort 处理。
func DailyQuotaRefund(ctx context.Context, key string, n int64) error {
	ctx, span := global.Tracer.Start(ctx, "redisdao.DailyQuotaRefund")
	defer span.End()

	if n <= 0 {
		return nil
	}
	if _, err := dailyQuotaRefundScript.Run(ctx, global.RedisEngine, []string{key}, n).Int64(); err != nil {
		span.RecordError(err)
		return err
	}
	return nil
}

// DailyQuotaValue 读取当日计数，key 不存在时返回 0。
func DailyQuotaValue(ctx context.Context, key string) (int64, error) {
	return quotaDailyValue(ctx, key)
}

// DailyQuotaDelete 删除每日计数缓存；MySQL 是权威账本，后续读取会从数据库恢复。
func DailyQuotaDelete(ctx context.Context, key string) error {
	if key == "" || global.RedisEngine == nil {
		return nil
	}
	return global.RedisEngine.Del(ctx, key).Err()
}