zarinpal-analytics-backend/internal/growthreport/coordinator.go

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, 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 {
	if c.config.DailyRequests < 1 || c.config.DailyTokens < 1 {
		return ErrBudgetExhausted
	}
	if c.redis == nil {
		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 {
		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)
}