fix(growth): harden AI jobs and cost controls

This commit is contained in:
2026-08-21 01:48:47 +03:30
parent 86e9d7c5f9
commit a15c04cac1
13 changed files with 366 additions and 78 deletions

View File

@@ -37,17 +37,17 @@ func (c *Coordinator) AllowCreate(ctx context.Context, userID int64, now time.Ti
func (c *Coordinator) AllowUncached(ctx context.Context, merchantID int64, now time.Time) error {
return c.increment(ctx, fmt.Sprintf("growth:scope:%d:%s", merchantID, tehranDate(now)), c.config.UncachedPerDay, 36*time.Hour, now)
}
func (c *Coordinator) ReserveBudget(ctx context.Context, tokens int64, now time.Time) error {
func (c *Coordinator) ReserveBudget(ctx context.Context, requests, tokens int64, now time.Time) error {
if c.config.DailyRequests < 1 || c.config.DailyTokens < 1 {
return ErrBudgetExhausted
}
if c.redis == nil {
if c.redis == nil || requests < 1 || tokens < 1 || requests > c.config.DailyRequests || tokens > c.config.DailyTokens {
return ErrBudgetExhausted
}
if err := c.incrementRedis(ctx, "growth:budget:req:"+tehranDate(now), 1, c.config.DailyRequests, 36*time.Hour); err != nil {
return ErrBudgetExhausted
}
if err := c.incrementRedis(ctx, "growth:budget:tok:"+tehranDate(now), tokens, c.config.DailyTokens, 36*time.Hour); err != nil {
script := redis.NewScript(`local req=tonumber(redis.call('GET',KEYS[1]) or '0');local tok=tonumber(redis.call('GET',KEYS[2]) or '0');local addReq=tonumber(ARGV[1]);local addTok=tonumber(ARGV[2]);if req+addReq>tonumber(ARGV[4]) or tok+addTok>tonumber(ARGV[5]) then return 0 end;redis.call('INCRBY',KEYS[1],addReq);redis.call('INCRBY',KEYS[2],addTok);redis.call('PEXPIRE',KEYS[1],ARGV[3]);redis.call('PEXPIRE',KEYS[2],ARGV[3]);return 1`)
date := tehranDate(now)
allowed, err := script.Run(ctx, c.redis, []string{"growth:budget:req:" + date, "growth:budget:tok:" + date}, requests, tokens, (36 * time.Hour).Milliseconds(), c.config.DailyRequests, c.config.DailyTokens).Int()
if err != nil || allowed != 1 {
return ErrBudgetExhausted
}
return nil