119 lines
3.9 KiB
Go
119 lines
3.9 KiB
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
|
|
semaphore chan struct{}
|
|
}
|
|
|
|
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{}, semaphore: make(chan struct{}, config.MaxConcurrent)}
|
|
}
|
|
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) {
|
|
select {
|
|
case c.semaphore <- struct{}{}:
|
|
return func() { <-c.semaphore }, nil
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
func (c *Coordinator) Publish(ctx context.Context, reportID string) {
|
|
if c.redis != nil {
|
|
_ = 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)
|
|
}
|
|
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)
|
|
}
|