feat: harden growth report orchestration
This commit is contained in:
@@ -42,12 +42,27 @@ func (c *Coordinator) ReserveBudget(ctx context.Context, tokens int64, now time.
|
||||
if c.config.DailyRequests < 1 || c.config.DailyTokens < 1 {
|
||||
return ErrBudgetExhausted
|
||||
}
|
||||
if err := c.increment(ctx, "growth:budget:req:"+tehranDate(now), c.config.DailyRequests, 36*time.Hour, now); err != nil {
|
||||
if c.redis == nil {
|
||||
return ErrBudgetExhausted
|
||||
}
|
||||
if err := c.incrementBy(ctx, "growth:budget:tok:"+tehranDate(now), tokens, c.config.DailyTokens, 36*time.Hour, now); err != nil {
|
||||
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 {
|
||||
return ErrBudgetExhausted
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Coordinator) incrementRedis(ctx context.Context, key string, amount, limit int64, ttl time.Duration) error {
|
||||
script := redis.NewScript(`local v=redis.call('INCRBY',KEYS[1],ARGV[1]);if v==tonumber(ARGV[1]) then redis.call('PEXPIRE',KEYS[1],ARGV[2]) end;if v>tonumber(ARGV[3]) then return 0 end;return 1`)
|
||||
value, err := script.Run(ctx, c.redis, []string{key}, amount, ttl.Milliseconds(), limit).Int()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if value == 0 {
|
||||
return ErrRateLimited
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (c *Coordinator) Enter(ctx context.Context) (func(), error) {
|
||||
@@ -63,6 +78,7 @@ func (c *Coordinator) Publish(ctx context.Context, reportID string) {
|
||||
_ = c.redis.LPush(ctx, "growth:jobs", reportID).Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Coordinator) increment(ctx context.Context, key string, limit int64, ttl time.Duration, now time.Time) error {
|
||||
return c.incrementBy(ctx, key, 1, limit, ttl, now)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user