rate limit: make context aware (#731)

* rate limits: Make context aware

* binance: rate limit allow for cancellation of reservation when deadline is exceeded

* request: add context.done() before initiating any bulk work.

* binance: update error return for rate limiting

* request: updated dealine check to remove after time.Now procedure as this will obfuscate a deadline which will be limited by the context check on every attempt, so no need to sleep with delay.

exchanges/request/limit.go

@@ -1,6 +1,7 @@
 package request
 
 import (
+	"context"
 	"errors"
 	"sync/atomic"
 	"time"
@@ -22,9 +23,8 @@ type BasicLimit struct {
 }
 
 // Limit executes a single rate limit set by NewRateLimit
-func (b *BasicLimit) Limit(_ EndpointLimit) error {
-	time.Sleep(b.r.Reserve().Delay())
-	return nil
+func (b *BasicLimit) Limit(ctx context.Context, _ EndpointLimit) error {
+	return b.r.Wait(ctx)
 }
 
 // EndpointLimit defines individual endpoint rate limits that are set when
@@ -35,7 +35,7 @@ type EndpointLimit int
 // wrapper for extended rate limiting configuration i.e. Shells of rate
 // limits with a global rate for sub rates.
 type Limiter interface {
-	Limit(EndpointLimit) error
+	Limit(context.Context, EndpointLimit) error
 }
 
 // NewRateLimit creates a new RateLimit based of time interval and how many
@@ -59,13 +59,13 @@ func NewBasicRateLimit(interval time.Duration, actions int) Limiter {
 }
 
 // InitiateRateLimit sleeps for designated end point rate limits
-func (r *Requester) InitiateRateLimit(e EndpointLimit) error {
+func (r *Requester) InitiateRateLimit(ctx context.Context, e EndpointLimit) error {
 	if atomic.LoadInt32(&r.disableRateLimiter) == 1 {
 		return nil
 	}
 
 	if r.limiter != nil {
-		return r.limiter.Limit(e)
+		return r.limiter.Limit(ctx, e)
 	}
 
 	return nil

exchanges/request/request.go

@@ -112,10 +112,17 @@ func (i *Item) validateRequest(ctx context.Context, r *Requester) (*http.Request
 // DoRequest performs a HTTP/HTTPS request with the supplied params
 func (r *Requester) doRequest(ctx context.Context, endpoint EndpointLimit, newRequest Generate) error {
 	for attempt := 1; ; attempt++ {
+		// Check if context has finished before executing new attempt.
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		// Initiate a rate limit reservation and sleep on requested endpoint
-		err := r.InitiateRateLimit(endpoint)
+		err := r.InitiateRateLimit(ctx, endpoint)
 		if err != nil {
-			return err
+			return fmt.Errorf("failed to rate limit HTTP request: %w", err)
 		}
 
 		p, err := newRequest()
@@ -162,7 +169,7 @@ func (r *Requester) doRequest(ctx context.Context, endpoint EndpointLimit, newRe
 				delay = after
 			}
 
-			if d, ok := req.Context().Deadline(); ok && d.After(time.Now()) && time.Now().Add(delay).After(d) {
+			if dl, ok := req.Context().Deadline(); ok && dl.Before(time.Now().Add(delay)) {
 				if err != nil {
 					return fmt.Errorf("deadline would be exceeded by retry, err: %v", err)
 				}

exchanges/request/request_test.go

@@ -180,20 +180,18 @@ type GlobalLimitTest struct {
 
 var errEndpointLimitNotFound = errors.New("endpoint limit not found")
 
-func (g *GlobalLimitTest) Limit(e EndpointLimit) error {
+func (g *GlobalLimitTest) Limit(ctx context.Context, e EndpointLimit) error {
 	switch e {
 	case Auth:
 		if g.Auth == nil {
 			return errors.New("auth rate not set")
 		}
-		time.Sleep(g.Auth.Reserve().Delay())
-		return nil
+		return g.Auth.Wait(ctx)
 	case UnAuth:
 		if g.UnAuth == nil {
 			return errors.New("unauth rate not set")
 		}
-		time.Sleep(g.UnAuth.Reserve().Delay())
-		return nil
+		return g.UnAuth.Wait(ctx)
 	default:
 		return fmt.Errorf("cannot execute functionality: %d %w",
 			e,
@@ -527,11 +525,24 @@ func TestBasicLimiter(t *testing.T) {
 	ctx := context.Background()
 
 	tn := time.Now()
-	_ = r.SendPayload(ctx, Unset, func() (*Item, error) { return &i, nil })
-	_ = r.SendPayload(ctx, Unset, func() (*Item, error) { return &i, nil })
+	err := r.SendPayload(ctx, Unset, func() (*Item, error) { return &i, nil })
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = r.SendPayload(ctx, Unset, func() (*Item, error) { return &i, nil })
+	if err != nil {
+		t.Fatal(err)
+	}
 	if time.Since(tn) < time.Second {
 		t.Error("rate limit issues")
 	}
+
+	ctx, cancel := context.WithDeadline(ctx, tn.Add(time.Nanosecond))
+	defer cancel()
+	err = r.SendPayload(ctx, Unset, func() (*Item, error) { return &i, nil })
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("receieved: %v but expected: %v", err, context.DeadlineExceeded)
+	}
 }
 
 func TestEnableDisableRateLimit(t *testing.T) {