package growthreport import ( "context" "errors" "fmt" "sync" "time" "github.com/redis/go-redis/v9" ) var ErrRateLimited = errors.New("growth report: rate limited") var ErrBudgetExhausted = errors.New("growth report: budget exhausted") type CoordinatorConfig struct { CreatePerMinute, UncachedPerDay, DailyRequests, DailyTokens int64 MaxConcurrent int } type Coordinator struct { redis *redis.Client config CoordinatorConfig mu sync.Mutex counters map[string]int64 expiries map[string]time.Time } func NewCoordinator(client *redis.Client, config CoordinatorConfig) *Coordinator { if config.MaxConcurrent < 1 { config.MaxConcurrent = 4 } return &Coordinator{redis: client, config: config, counters: map[string]int64{}, expiries: map[string]time.Time{}} } func (c *Coordinator) AllowCreate(ctx context.Context, userID int64, now time.Time) error { return c.increment(ctx, fmt.Sprintf("growth:create:%d:%s", userID, now.UTC().Format("200601021504")), c.config.CreatePerMinute, time.Minute, now) } func (c *Coordinator) AllowUncached(ctx context.Context, merchantID int64, terminalID string, now time.Time) error { return c.increment(ctx, fmt.Sprintf("growth:scope:%d:%s:%s", merchantID, terminalID, tehranDate(now)), c.config.UncachedPerDay, 36*time.Hour, now) } 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 || requests < 1 || tokens < 1 || requests > c.config.DailyRequests || tokens > c.config.DailyTokens { return ErrBudgetExhausted } 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 } 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) { if c.redis == nil { return nil, ErrBudgetExhausted } acquire := redis.NewScript(`local v=redis.call('INCR',KEYS[1]);if v==1 then redis.call('PEXPIRE',KEYS[1],ARGV[2]) end;if v>tonumber(ARGV[1]) then redis.call('DECR',KEYS[1]);return 0 end;return 1`) for { ok, err := acquire.Run(ctx, c.redis, []string{"growth:ai:concurrency"}, c.config.MaxConcurrent, (3 * time.Minute).Milliseconds()).Int() if err != nil { return nil, ErrBudgetExhausted } if ok == 1 { return func() { releaseCtx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() script := redis.NewScript(`local v=tonumber(redis.call('GET',KEYS[1]) or '0');if v>0 then return redis.call('DECR',KEYS[1]) end return 0`) _ = script.Run(releaseCtx, c.redis, []string{"growth:ai:concurrency"}).Err() }, nil } select { case <-ctx.Done(): return nil, ctx.Err() case <-time.After(100 * time.Millisecond): } } } 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) } func (c *Coordinator) incrementBy(ctx context.Context, key string, amount, limit int64, ttl time.Duration, now time.Time) error { if limit < 1 { return ErrRateLimited } if c.redis != nil { 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 { if value == 0 { return ErrRateLimited } return nil } } c.mu.Lock() defer c.mu.Unlock() if expiry := c.expiries[key]; !expiry.After(now) { c.counters[key] = 0 c.expiries[key] = now.Add(ttl) } next := c.counters[key] + amount if next > limit { return ErrRateLimited } c.counters[key] = next return nil } func tehranDate(now time.Time) string { location, err := time.LoadLocation("Asia/Tehran") if err != nil { location = time.FixedZone("Asia/Tehran", 12600) } return now.In(location).Format(time.DateOnly) }