Merge pull request #18 from veksh/ratelim-window

change rate limiter to window

Author: veksh

internal/client/client.go

@@ -23,9 +23,13 @@ import (
 
 const (
 	HTTP_TIMEOUT = 10
-	HTTP_RPS     = 1
-	HTTP_BURST   = 60
-	DOMAINS_URL  = "/v1/domains/"
+	// burst RL: not currently used
+	HTTP_RPS   = 1
+	HTTP_BURST = 60
+	// window RL: window size, max requests per window
+	HTTP_RATE_WINDOW = time.Duration(60) * time.Second
+	HTTP_RATE_RPW    = 60
+	DOMAINS_URL      = "/v1/domains/"
 )
 
 var _ model.DNSApiClient = Client{}
@@ -45,17 +49,14 @@ func NewClient(apiURL string, key string, secret string) (*Client, error) {
 			Timeout: HTTP_TIMEOUT * time.Second}).DialContext,
 		TLSHandshakeTimeout:   HTTP_TIMEOUT * time.Second,
 		ResponseHeaderTimeout: HTTP_TIMEOUT * time.Second,
-		MaxIdleConns:          10,
-		MaxIdleConnsPerHost:   10,
-		MaxConnsPerHost:       10,
-		IdleConnTimeout:       60,
 	}
-	rateLimiter, err := ratelimiter.New(HTTP_RPS, HTTP_BURST)
+	// TODO: mb make it pluggable as a parameter
+	// rateLimiter, err := ratelimiter.NewBucketRL(HTTP_RPS, HTTP_BURST)
+	rateLimiter, err := ratelimiter.NewWindowRL(HTTP_RATE_WINDOW, HTTP_RATE_RPW)
 	if err != nil {
 		return nil, errors.Wrap(err, "cannot create rate limiter")
 	}
 	httpClient := http.Client{
-		Timeout: HTTP_TIMEOUT * time.Second,
 		Transport: &rateLimitedHTTPTransport{
 			limiter: rateLimiter,
 			next:    httpTransport,
@@ -93,9 +94,6 @@ func (c Client) makeRecordsRequest(ctx context.Context, path string, method stri
 
 	requestURL, _ := url.JoinPath(c.apiURL, DOMAINS_URL, path)
 
-	ctx, fnCancel := context.WithTimeout(ctx, HTTP_TIMEOUT*time.Second)
-	defer fnCancel()
-
 	req, err := http.NewRequestWithContext(ctx, method, requestURL, body)
 	if err != nil {
 		return nil, errors.Wrap(err, "cannot create request")

internal/client/ratelimittransport.go

@@ -7,11 +7,13 @@ import (
 )
 
 type rateLimitedHTTPTransport struct {
-	limiter *ratelimiter.RateLimiter
+	limiter ratelimiter.Limiter
 	next    http.RoundTripper
 }
 
 func (t *rateLimitedHTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	t.limiter.Wait()
+	if err := t.limiter.WaitCtx(req.Context()); err != nil {
+		return nil, err
+	}
 	return t.next.RoundTrip(req)
 }

internal/provider/record_test.go

@@ -52,7 +52,7 @@ func TestUnitALifecycle(t *testing.T) {
 	// 2nd step: update
 	mRecsToUpd := makeTestRecSet(model.REC_A, []model.DNSRecordData{"2.2.2.2", "4.4.4.4"})
 	mRecsToUpdTgt := makeTestRecSet(model.REC_A, []model.DNSRecordData{"1.1.1.1", "2.2.2.2", "4.4.4.4"})
-	// mRecsTgt4 := makeTestRecSet(model.REC_A, []model.DNSRecordData{"1.1.1.1", "4.4.4.4"})
+	mRecsTgt4 := makeTestRecSet(model.REC_A, []model.DNSRecordData{"1.1.1.1", "4.4.4.4"})
 
 	mType, mName := model.REC_A, mRecsPre.DNSRecName
 
@@ -67,11 +67,27 @@ func TestUnitALifecycle(t *testing.T) {
 	mockClientUpd := model.NewMockDNSApiClient(t)
 	mockClientUpd.EXPECT().GetRecords(mCtx, mDom, mType, mName).Return(mRecsTgt.Records, nil).Times(3).Run(traceMarker("get 1 (mRegsTgt)*3"))
 	mockClientUpd.EXPECT().SetRecords(mCtx, mDom, mType, mName, mRecsToUpdTgt.UpdRecords).Return(nil).Once().Run(traceMarker("set 1 (mRecsToUpdTgt)"))
-	// again, clean-up order is not deterministic, so those are a bit fragile
 	mockClientUpd.EXPECT().GetRecords(mCtx, mDom, mType, mName).Return(mRecsToUpdTgt.Records, nil).Times(3).Run(traceMarker("get 2/1 (mRecsToUpdTgt)*3"))
-	mockClientUpd.EXPECT().GetRecords(mCtx, mDom, mType, mName).Return(mRecsTgt1.Records, nil).Once().Run(traceMarker("get 2/2 (mRecsTgt1)"))
 
-	mockClientUpd.EXPECT().SetRecords(mCtx, mDom, mType, mName, mRecsTgt1.UpdRecords).Return(nil).Once().Run(traceMarker("set 2/1 (mRecsTgt1)"))
+	// again, clean-up order is not deterministic, either could be called first
+	// lets stash the expected result and return it later
+	var res1or4 []model.DNSRecord = make([]model.DNSRecord, 2)
+	mockClientUpd.EXPECT().SetRecords(mCtx, mDom, mType, mName, mRecsTgt1.UpdRecords).Return(nil).Maybe().
+		Run(func(args mock.Arguments) {
+			fmt.Println("set 2/1 (mRecsTgt1)", args[4])
+			copy(res1or4, mRecsTgt1.Records)
+		})
+	mockClientUpd.EXPECT().SetRecords(mCtx, mDom, mType, mName, mRecsTgt4.UpdRecords).Return(nil).Maybe().
+		Run(func(args mock.Arguments) {
+			fmt.Println("set 2/1 (mRecsTgt4)", args[4])
+			copy(res1or4, mRecsTgt4.Records)
+		})
+	// now need to return right results depending on set above
+	mockClientUpd.EXPECT().GetRecords(mCtx, mDom, mType, mName).Return(res1or4, nil).Once().
+		Run(func(args mock.Arguments) {
+			fmt.Println("get 2/2 (mRecsTgt1/4 = " + res1or4[1].Data + ")")
+		})
+	// and, finally, the common cleanup
 	mockClientUpd.EXPECT().SetRecords(mCtx, mDom, mType, mName, mRecsPre.UpdRecords).Return(nil).Once().Run(traceMarker("set 2/2 (mRecsPre)"))
 
 	resource.UnitTest(t, resource.TestCase{

libs/ratelimiter/ratelimiter.go renamed to libs/ratelimiter/bucket.go

@@ -14,7 +14,7 @@ import (
 //   - check the time since a token was added, mb add some more
 //   - spend one if now it is available
 //   - if none, wait until next one will be available (while holding mutex)
-type RateLimiter struct {
+type BucketRateLimiter struct {
 	mu sync.Mutex
 	// interval between tokens, s
 	period time.Duration
@@ -27,11 +27,11 @@ type RateLimiter struct {
 }
 
 // context to cancel, rate per second, burst (bucket) size
-func New(rate, burst int) (*RateLimiter, error) {
+func NewBucketRL(rate, burst int) (Limiter, error) {
 	if rate <= 0 || burst <= 0 {
 		return nil, errors.New("limiter rate and burst must be positive")
 	}
-	return &RateLimiter{
+	return &BucketRateLimiter{
 		period:        time.Second / time.Duration(rate),
 		bucketSize:    burst,
 		numTokens:     burst,
@@ -40,7 +40,7 @@ func New(rate, burst int) (*RateLimiter, error) {
 }
 
 // block until token becomes available
-func (s *RateLimiter) Wait() {
+func (s *BucketRateLimiter) Wait() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	// if bucket is empty: add tokens that are due, update last refill time
@@ -67,7 +67,7 @@ func (s *RateLimiter) Wait() {
 }
 
 // block until token becomes available, with cancellable context
-func (s *RateLimiter) WaitCtx(ctx context.Context) error {
+func (s *BucketRateLimiter) WaitCtx(ctx context.Context) error {
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}

libs/ratelimiter/limiter.go

@@ -0,0 +1,11 @@
+package ratelimiter
+
+import "context"
+
+// GoDaddy current rate limiter: 60 requests in 60 seconds window per API endpoint
+// https://developer.godaddy.com/getstarted#terms
+// mb have separate limits per endpoint (get/set/del): need to see a request
+type Limiter interface {
+	Wait()
+	WaitCtx(ctx context.Context) error
+}

libs/ratelimiter/window.go

@@ -0,0 +1,94 @@
+package ratelimiter
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+)
+
+// fixed winow rate limiter: limit number of request allowed per time interval
+// - keeps track of interval start and last call time
+// - on acquire:
+//   - if last call was > interval length ago, refill bucket
+//   - if there are tokens available, spend one and return immediately
+//   - if not, wait until the end of interval, refill and spend one
+//   - somewhat unpleasant side effect will be the burst of `bucketSize` requests
+//     followed by the large wait for next: alternative is 1 RPS which is way slower
+//     for shorter batches
+type WindowRateLimiter struct {
+	mu sync.Mutex
+	// interval between bucket refills
+	period time.Duration
+	// total bucket size: num of requests per period
+	bucketSize int
+	// number of requests left for current period
+	numTokens int
+	// the time last bucket was refilled
+	lastRefillTime time.Time
+}
+
+// period length and num of requests per period
+func NewWindowRL(period time.Duration, RPP int) (Limiter, error) {
+	if period <= 0 || RPP <= 0 {
+		return nil, errors.New("limiter period and num of requests must be positive")
+	}
+	return &WindowRateLimiter{
+		period:         period,
+		bucketSize:     RPP,
+		numTokens:      RPP,
+		lastRefillTime: time.Now(),
+	}, nil
+}
+
+// block until token becomes available
+func (s *WindowRateLimiter) Wait() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	// if last refilled happened long ago: refill now
+	if time.Since(s.lastRefillTime) > s.period {
+		s.lastRefillTime = time.Now()
+		s.numTokens = s.bucketSize
+	}
+	// if bucket is empty: wait for next refill
+	if s.numTokens == 0 {
+		nextTokenTime := s.lastRefillTime.Add(s.period)
+		waitDuration := time.Until(nextTokenTime)
+		time.Sleep(waitDuration)
+
+		s.lastRefillTime = time.Now()
+		s.numTokens = s.bucketSize
+	}
+	s.numTokens--
+}
+
+// block until token becomes available, with cancellable context
+func (s *WindowRateLimiter) WaitCtx(ctx context.Context) error {
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// if last refilled happened long ago: refill now
+	if time.Since(s.lastRefillTime) > s.period {
+		s.lastRefillTime = time.Now()
+		s.numTokens = s.bucketSize
+	}
+	// if bucket is empty: wait for next refill
+	if s.numTokens == 0 {
+		nextTokenTime := s.lastRefillTime.Add(s.period)
+		waitDuration := time.Until(nextTokenTime)
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(waitDuration):
+		}
+
+		s.lastRefillTime = time.Now()
+		s.numTokens = s.bucketSize
+	}
+	s.numTokens--
+	return nil
+}