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/binance/ratelimit.go

@@ -1,6 +1,8 @@
 package binance
 
 import (
+	"context"
+	"fmt"
 	"time"
 
 	"github.com/thrasher-corp/gocryptotrader/exchanges/request"
@@ -103,7 +105,7 @@ type RateLimit struct {
 }
 
 // Limit executes rate limiting functionality for Binance
-func (r *RateLimit) Limit(f request.EndpointLimit) error {
+func (r *RateLimit) Limit(ctx context.Context, f request.EndpointLimit) error {
 	var limiter *rate.Limiter
 	var tokens int
 	switch f {
@@ -214,11 +216,25 @@ func (r *RateLimit) Limit(f request.EndpointLimit) error {
 	}
 
 	var finalDelay time.Duration
+	var reserves = make([]*rate.Reservation, tokens)
 	for i := 0; i < tokens; i++ {
 		// Consume tokens 1 at a time as this avoids needing burst capacity in the limiter,
 		// which would otherwise allow the rate limit to be exceeded over short periods
-		finalDelay = limiter.Reserve().Delay()
+		reserves[i] = limiter.Reserve()
+		finalDelay = reserves[i].Delay()
 	}
+
+	if dl, ok := ctx.Deadline(); ok && dl.Before(time.Now().Add(finalDelay)) {
+		// Cancel all potential reservations to free up rate limiter if deadline
+		// is exceeded.
+		for x := range reserves {
+			reserves[x].Cancel()
+		}
+		return fmt.Errorf("rate limit delay of %s will exceed deadline: %w",
+			finalDelay,
+			context.DeadlineExceeded)
+	}
+
 	time.Sleep(finalDelay)
 	return nil
 }

exchanges/binance/ratelimit_test.go

@@ -1,7 +1,10 @@
 package binance
 
 import (
+	"context"
+	"errors"
 	"testing"
+	"time"
 
 	"github.com/thrasher-corp/gocryptotrader/exchanges/request"
 )
@@ -12,6 +15,7 @@ func TestRateLimit_Limit(t *testing.T) {
 	testTable := map[string]struct {
 		Expected request.EndpointLimit
 		Limit    request.EndpointLimit
+		Deadline time.Time
 	}{
 		"All Orderbooks Ticker": {Expected: spotOrderbookTickerAllRate, Limit: bestPriceLimit("")},
 		"Orderbook Ticker":      {Expected: spotDefaultRate, Limit: bestPriceLimit(symbol)},
@@ -24,6 +28,7 @@ func TestRateLimit_Limit(t *testing.T) {
 		"Orderbook Depth 500":   {Expected: spotOrderbookDepth500Rate, Limit: orderbookLimit(500)},
 		"Orderbook Depth 1000":  {Expected: spotOrderbookDepth1000Rate, Limit: orderbookLimit(1000)},
 		"Orderbook Depth 5000":  {Expected: spotOrderbookDepth5000Rate, Limit: orderbookLimit(5000)},
+		"Exceeds deadline":      {Expected: spotOrderbookDepth5000Rate, Limit: orderbookLimit(5000), Deadline: time.Now().Add(time.Nanosecond)},
 	}
 	for name, tt := range testTable {
 		tt := tt
@@ -35,8 +40,15 @@ func TestRateLimit_Limit(t *testing.T) {
 				t.Fatalf("incorrect limit applied.\nexp: %v\ngot: %v", exp, got)
 			}
 
+			ctx := context.Background()
+			if !tt.Deadline.IsZero() {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithDeadline(ctx, tt.Deadline)
+				defer cancel()
+			}
+
 			l := SetRateLimit()
-			if err := l.Limit(tt.Limit); err != nil {
+			if err := l.Limit(ctx, tt.Limit); err != nil && !errors.Is(err, context.DeadlineExceeded) {
 				t.Fatalf("error applying rate limit: %v", err)
 			}
 		})
@@ -56,7 +68,7 @@ func TestRateLimit_LimitStatic(t *testing.T) {
 			t.Parallel()
 
 			l := SetRateLimit()
-			if err := l.Limit(tt); err != nil {
+			if err := l.Limit(context.Background(), tt); err != nil {
 				t.Fatalf("error applying rate limit: %v", err)
 			}
 		})