github.com/rudderlabs/rudder-go-kit@v0.30.0/throttling/memory_gcra.go (about)

     1  package throttling
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"sync/atomic"
     8  	"time"
     9  
    10  	"github.com/throttled/throttled/v2"
    11  
    12  	"github.com/rudderlabs/rudder-go-kit/cachettl"
    13  )
    14  
    15  const defaultMaxCASAttemptsLimit = 100
    16  
    17  type gcra struct {
    18  	mu    sync.Mutex
    19  	store *cachettl.Cache[string, *throttled.GCRARateLimiterCtx]
    20  }
    21  
    22  func (g *gcra) limit(ctx context.Context, key string, cost, burst, rate, period int64) (
    23  	bool, time.Duration, error,
    24  ) {
    25  	rl, err := g.getLimiter(key, burst, rate, period)
    26  	if err != nil {
    27  		return false, 0, err
    28  	}
    29  
    30  	limited, res, err := rl.RateLimitCtx(ctx, "key", int(cost))
    31  	if err != nil {
    32  		return false, 0, fmt.Errorf("could not rate limit: %w", err)
    33  	}
    34  
    35  	return !limited, res.RetryAfter, nil
    36  }
    37  
    38  func (g *gcra) getLimiter(key string, burst, rate, period int64) (*throttled.GCRARateLimiterCtx, error) {
    39  	g.mu.Lock()
    40  	defer g.mu.Unlock()
    41  
    42  	if g.store == nil {
    43  		g.store = cachettl.New[string, *throttled.GCRARateLimiterCtx]()
    44  	}
    45  
    46  	rl := g.store.Get(key)
    47  	if rl == nil {
    48  		var err error
    49  		rl, err = throttled.NewGCRARateLimiterCtx(newGCRAMemStore(), throttled.RateQuota{
    50  			MaxRate:  throttled.PerDuration(int(rate), time.Duration(period)*time.Second),
    51  			MaxBurst: int(burst),
    52  		})
    53  		if err != nil {
    54  			return nil, fmt.Errorf("could not create rate limiter: %w", err)
    55  		}
    56  		rl.SetMaxCASAttemptsLimit(defaultMaxCASAttemptsLimit)
    57  		// rate limiter should be cached for (burst/rate)*period seconds
    58  		// e.g. if burst is 100 and rate is 10/sec, then the rate limiter should be cached for 10 seconds
    59  		g.store.Put(key, rl, time.Duration((burst/rate)*period)*time.Second)
    60  	}
    61  
    62  	return rl, nil
    63  }
    64  
    65  func newGCRAMemStore() throttled.GCRAStoreCtx {
    66  	var v atomic.Int64
    67  	v.Store(-1)
    68  	return &gcraMemStore{
    69  		v:       &v,
    70  		timeNow: time.Now,
    71  	}
    72  }
    73  
    74  type gcraMemStore struct {
    75  	v       *atomic.Int64
    76  	timeNow func() time.Time
    77  }
    78  
    79  // GetWithTime returns the value of the key or -1 if it does not exist.
    80  // It also returns the current time at the Store.
    81  // The time must be representable as a positive int64 of nanoseconds since the epoch.
    82  func (ms *gcraMemStore) GetWithTime(_ context.Context, _ string) (int64, time.Time, error) {
    83  	return ms.v.Load(), ms.timeNow(), nil
    84  }
    85  
    86  // SetIfNotExistsWithTTL sets the value of key only if it is not ready set in the store (-1 == not exists)
    87  // it returns whether a new value was set.
    88  func (ms *gcraMemStore) SetIfNotExistsWithTTL(_ context.Context, _ string, value int64, _ time.Duration) (bool, error) {
    89  	swapped := ms.v.CompareAndSwap(-1, value)
    90  	return swapped, nil
    91  }
    92  
    93  // CompareAndSwapWithTTL atomically compares the value at key to the old value.
    94  // If it matches, it sets it to the new value and returns true, false otherwise.
    95  func (ms *gcraMemStore) CompareAndSwapWithTTL(_ context.Context, _ string, old, new int64, _ time.Duration) (bool, error) {
    96  	swapped := ms.v.CompareAndSwap(old, new)
    97  	return swapped, nil
    98  }