Repository: go-resty/resty
Pull request #715 (2023-09-29 14:32:45)

Add rate limiter to client
Add a simple per client rate limiting based on `Allow() bool` interface

Continue work on https://github.com/go-resty/resty/pull/494

SVilgelm: I'm unable to push directly to @dbriemann's PR, so had to create this one.
Changed to use an interface instead of `x/time/rate`, also rebased
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/go-resty/resty/pull/715?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`41199c3`)](https://app.codecov.io/gh/go-resty/resty/commit/41199c3df1af740f587bf704674a2afcf31c68cb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) 96.33% compared to head [(`8c264e1`)](https://app.codecov.io/gh/go-resty/resty/pull/715?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) 96.34%.

> :exclamation: Current head 8c264e1 differs from pull request most recent head d8914b2. Consider uploading reports for the commit d8914b2 to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #715      +/-   ##
==========================================
+ Coverage   96.33%   96.34%   +0.01%     
==========================================
  Files          12       12              
  Lines        1610     1616       +6     
==========================================
+ Hits         1551     1557       +6     
  Misses         37       37              
  Partials       22       22              
```


| [Files](https://app.codecov.io/gh/go-resty/resty/pull/715?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) | Coverage Δ | |
|---|---|---|
| [client.go](https://app.codecov.io/gh/go-resty/resty/pull/715?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty#diff-Y2xpZW50Lmdv) | `98.04% <100.00%> (+0.02%)` | :arrow_up: |
| [util.go](https://app.codecov.io/gh/go-resty/resty/pull/715?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty#diff-dXRpbC5nbw==) | `94.44% <ø> (ø)` | |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/go-resty/resty/pull/715?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty).


Patch:
diff --git a/client.go b/client.go
index ae96018..9165bf2 100644
--- a/client.go
+++ b/client.go
@@ -152,6 +152,7 @@ type Client struct {
 	errorHooks          []ErrorHook
 	invalidHooks        []ErrorHook
 	panicHooks          []ErrorHook
+	rateLimiter         RateLimiter
 }
 
 // User type is to hold an username and password information
@@ -920,6 +921,13 @@ func (c *Client) SetOutputDirectory(dirPath string) *Client {
 	return c
 }
 
+// SetRateLimiter sets an optional `RateLimiter`. If set the rate limiter will control
+// all requests made with this client.
+func (c *Client) SetRateLimiter(rl RateLimiter) *Client {
+	c.rateLimiter = rl
+	return c
+}
+
 // SetTransport method sets custom `*http.Transport` or any `http.RoundTripper`
 // compatible interface implementation in the resty client.
 //
@@ -1141,6 +1149,14 @@ func (c *Client) execute(req *Request) (*Response, error) {
 		}
 	}
 
+	// If there is a rate limiter set for this client, the Execute call
+	// will return an error if the rate limit is exceeded.
+	if req.client.rateLimiter != nil {
+		if !req.client.rateLimiter.Allow() {
+			return nil, wrapNoRetryErr(ErrRateLimitExceeded)
+		}
+	}
+
 	// resty middlewares
 	for _, f := range c.beforeRequest {
 		if err = f(c, req); err != nil {
diff --git a/go.mod b/go.mod
index 8ec1fd5..09e7e32 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,7 @@ module github.com/go-resty/resty/v2
 
 go 1.16
 
-require golang.org/x/net v0.15.0
+require (
+	golang.org/x/net v0.15.0
+	golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11
+)
diff --git a/go.sum b/go.sum
index 10bb1cb..3a3acf4 100644
--- a/go.sum
+++ b/go.sum
@@ -33,6 +33,8 @@ golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
+golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11 h1:GZokNIeuVkl3aZHJchRrr13WCsols02MLUcz1U9is6M=
+golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
diff --git a/util.go b/util.go
index 99c203b..3279116 100644
--- a/util.go
+++ b/util.go
@@ -6,6 +6,7 @@ package resty
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -64,6 +65,16 @@ func (l *logger) output(format string, v ...interface{}) {
 	l.l.Printf(format, v...)
 }
 
+//‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
+// Rate Limiter interface
+//_______________________________________________________________________
+
+type RateLimiter interface {
+	Allow() bool
+}
+
+var ErrRateLimitExceeded = errors.New("rate limit exceeded")
+
 //‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
 // Package Helper methods
 //_______________________________________________________________________


Test patch:
diff --git a/request_test.go b/request_test.go
index f48660d..15632ab 100644
--- a/request_test.go
+++ b/request_test.go
@@ -19,6 +19,8 @@ import (
 	"strings"
 	"testing"
 	"time"
+
+	"golang.org/x/time/rate"
 )
 
 type AuthSuccess struct {
@@ -66,6 +68,44 @@ func TestGetGH524(t *testing.T) {
 	assertEqual(t, resp.Request.Header.Get("Content-Type"), "") //  unable to reproduce reported issue
 }
 
+func TestRateLimiter(t *testing.T) {
+	ts := createGetServer(t)
+	defer ts.Close()
+
+	// Test a burst with a valid capacity and then a consecutive request that must fail.
+
+	// Allow a rate of 1 every 100 ms but also allow bursts of 10 requests.
+	client := dc().SetRateLimiter(rate.NewLimiter(rate.Every(100*time.Millisecond), 10))
+
+	// Execute a burst of 10 requests.
+	for i := 0; i < 10; i++ {
+		resp, err := client.R().
+			SetQueryParam("request_no", strconv.Itoa(i)).Get(ts.URL + "/")
+		assertError(t, err)
+		assertEqual(t, http.StatusOK, resp.StatusCode())
+	}
+	// Next request issued directly should fail because burst of 10 has been consumed.
+	{
+		_, err := client.R().
+			SetQueryParam("request_no", strconv.Itoa(11)).Get(ts.URL + "/")
+		assertErrorIs(t, ErrRateLimitExceeded, err)
+	}
+
+	// Test continues request at a valid rate
+
+	// Allow a rate of 1 every ms with no burst.
+	client = dc().SetRateLimiter(rate.NewLimiter(rate.Every(1*time.Millisecond), 1))
+
+	// Sending requests every ms+tiny delta must succeed.
+	for i := 0; i < 100; i++ {
+		resp, err := client.R().
+			SetQueryParam("request_no", strconv.Itoa(i)).Get(ts.URL + "/")
+		assertError(t, err)
+		assertEqual(t, http.StatusOK, resp.StatusCode())
+		time.Sleep(1*time.Millisecond + 100*time.Microsecond)
+	}
+}
+
 func TestIllegalRetryCount(t *testing.T) {
 	ts := createGetServer(t)
 	defer ts.Close()