From a923810d6507d29aa2317e4c4eacfa65aeecc39d Mon Sep 17 00:00:00 2001 From: David Linus Briemann Date: Tue, 30 Nov 2021 21:57:22 +0100 Subject: [PATCH 1/5] add rate limiter to client --- client.go | 10 ++++++++++ go.mod | 5 ++++- go.sum | 2 ++ request.go | 8 ++++++++ request_test.go | 33 +++++++++++++++++++++++++++++++++ 5 files changed, 57 insertions(+), 1 deletion(-) diff --git a/client.go b/client.go index ae960189..ffff67b6 100644 --- a/client.go +++ b/client.go @@ -23,6 +23,8 @@ import ( "strings" "sync" "time" + + "golang.org/x/time/rate" ) const ( @@ -152,6 +154,7 @@ type Client struct { errorHooks []ErrorHook invalidHooks []ErrorHook panicHooks []ErrorHook + rateLimiter *rate.Limiter } // User type is to hold an username and password information @@ -920,6 +923,13 @@ func (c *Client) SetOutputDirectory(dirPath string) *Client { return c } +// SetRateLimiter sets an optional `*rate.Limiter`. If set the rate limiter will control +// all requests made with this client. +func (c *Client) SetRateLimiter(rl *rate.Limiter) *Client { + c.rateLimiter = rl + return c +} + // SetTransport method sets custom `*http.Transport` or any `http.RoundTripper` // compatible interface implementation in the resty client. // diff --git a/go.mod b/go.mod index 8ec1fd58..09e7e326 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 10bb1cb0..3a3acf42 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/request.go b/request.go index fec09763..67d1830e 100644 --- a/request.go +++ b/request.go @@ -920,6 +920,14 @@ func (r *Request) Execute(method, url string) (*Response, error) { return nil, err } + // If there is a rate limiter set for this client, the Execute call will block + // until the request is allowed to go through. + if r.client.rateLimiter != nil { + if err := r.client.rateLimiter.Wait(r.Context()); err != nil { + return nil, err + } + } + if r.SRV != nil { _, addrs, err = net.LookupSRV(r.SRV.Service, "tcp", r.SRV.Domain) if err != nil { diff --git a/request_test.go b/request_test.go index f48660d4..f63dcd67 100644 --- a/request_test.go +++ b/request_test.go @@ -9,6 +9,7 @@ import ( "crypto/tls" "encoding/xml" "errors" + "fmt" "io" "net" "net/http" @@ -19,6 +20,8 @@ import ( "strings" "testing" "time" + + "golang.org/x/time/rate" ) type AuthSuccess struct { @@ -66,6 +69,36 @@ 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() + + client := dc().SetRateLimiter(rate.NewLimiter(rate.Every(100*time.Millisecond), 10)) + + start := time.Now() + for i := 0; i < 21; i++ { + resp, err := client.R(). + SetQueryParam("request_no", strconv.FormatInt(time.Now().Unix(), 10)). + Get(ts.URL + "/") + assertError(t, err) + assertEqual(t, http.StatusOK, resp.StatusCode()) + assertEqual(t, "HTTP/1.1", resp.Proto()) + assertEqual(t, "200 OK", resp.Status()) + assertNotNil(t, resp.Body()) + assertEqual(t, "TestGet: text response", resp.String()) + + logResponse(t, resp) + } + + assertError(t, func() error { + dur := time.Now().Sub(start) + if dur <= 1*time.Second { + return fmt.Errorf("requests executed too fast (%f). rate limiting not working correctly", dur.Seconds()) + } + return nil + }()) +} + func TestIllegalRetryCount(t *testing.T) { ts := createGetServer(t) defer ts.Close() From 689812e92a12ebeb19d86ed574a4a4362bd250f0 Mon Sep 17 00:00:00 2001 From: David Linus Briemann Date: Wed, 1 Dec 2021 22:48:20 +0100 Subject: [PATCH 2/5] make rate limiter work for retries --- client.go | 8 ++++++++ request.go | 8 -------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/client.go b/client.go index ffff67b6..ee73401f 100644 --- a/client.go +++ b/client.go @@ -1151,6 +1151,14 @@ func (c *Client) execute(req *Request) (*Response, error) { } } + // If there is a rate limiter set for this client, the Execute call will block + // until the request is allowed to go through. + if req.client.rateLimiter != nil { + if err := req.client.rateLimiter.Wait(req.Context()); err != nil { + return nil, err + } + } + // resty middlewares for _, f := range c.beforeRequest { if err = f(c, req); err != nil { diff --git a/request.go b/request.go index 67d1830e..fec09763 100644 --- a/request.go +++ b/request.go @@ -920,14 +920,6 @@ func (r *Request) Execute(method, url string) (*Response, error) { return nil, err } - // If there is a rate limiter set for this client, the Execute call will block - // until the request is allowed to go through. - if r.client.rateLimiter != nil { - if err := r.client.rateLimiter.Wait(r.Context()); err != nil { - return nil, err - } - } - if r.SRV != nil { _, addrs, err = net.LookupSRV(r.SRV.Service, "tcp", r.SRV.Domain) if err != nil { From 5351a660d603640b446f6c160ad0acf8fda02d2e Mon Sep 17 00:00:00 2001 From: David Linus Briemann Date: Fri, 25 Feb 2022 09:03:41 +0100 Subject: [PATCH 3/5] make rate limiter return error instead of blocking --- client.go | 8 ++++---- errors.go | 7 +++++++ 2 files changed, 11 insertions(+), 4 deletions(-) create mode 100644 errors.go diff --git a/client.go b/client.go index ee73401f..23240269 100644 --- a/client.go +++ b/client.go @@ -1151,11 +1151,11 @@ func (c *Client) execute(req *Request) (*Response, error) { } } - // If there is a rate limiter set for this client, the Execute call will block - // until the request is allowed to go through. + // 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 err := req.client.rateLimiter.Wait(req.Context()); err != nil { - return nil, err + if !req.client.rateLimiter.Allow() { + return nil, wrapNoRetryErr(ErrRateLimitExceeded) } } diff --git a/errors.go b/errors.go new file mode 100644 index 00000000..f8efe7ea --- /dev/null +++ b/errors.go @@ -0,0 +1,7 @@ +package resty + +import "errors" + +var ( + ErrRateLimitExceeded = errors.New("rate limit exceeded") +) From 2587aa4b1d43233d8473e6ce87b4228306b96e8f Mon Sep 17 00:00:00 2001 From: David Linus Briemann Date: Fri, 24 Mar 2023 18:49:11 +0100 Subject: [PATCH 4/5] fix test --- request_test.go | 43 +++++++++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/request_test.go b/request_test.go index f63dcd67..15632ab0 100644 --- a/request_test.go +++ b/request_test.go @@ -9,7 +9,6 @@ import ( "crypto/tls" "encoding/xml" "errors" - "fmt" "io" "net" "net/http" @@ -73,30 +72,38 @@ 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)) - start := time.Now() - for i := 0; i < 21; i++ { + // Execute a burst of 10 requests. + for i := 0; i < 10; i++ { resp, err := client.R(). - SetQueryParam("request_no", strconv.FormatInt(time.Now().Unix(), 10)). - Get(ts.URL + "/") + SetQueryParam("request_no", strconv.Itoa(i)).Get(ts.URL + "/") assertError(t, err) assertEqual(t, http.StatusOK, resp.StatusCode()) - assertEqual(t, "HTTP/1.1", resp.Proto()) - assertEqual(t, "200 OK", resp.Status()) - assertNotNil(t, resp.Body()) - assertEqual(t, "TestGet: text response", resp.String()) - - logResponse(t, resp) + } + // 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) } - assertError(t, func() error { - dur := time.Now().Sub(start) - if dur <= 1*time.Second { - return fmt.Errorf("requests executed too fast (%f). rate limiting not working correctly", dur.Seconds()) - } - return nil - }()) + // 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) { From d8914b2e2927be5cdba073b50c1fb41e8efce6bc Mon Sep 17 00:00:00 2001 From: Sergey Vilgelm Date: Fri, 29 Sep 2023 07:21:43 -0700 Subject: [PATCH 5/5] use RateLimiter interface instead of x/time/rate --- client.go | 8 +++----- errors.go | 7 ------- util.go | 11 +++++++++++ 3 files changed, 14 insertions(+), 12 deletions(-) delete mode 100644 errors.go diff --git a/client.go b/client.go index 23240269..9165bf25 100644 --- a/client.go +++ b/client.go @@ -23,8 +23,6 @@ import ( "strings" "sync" "time" - - "golang.org/x/time/rate" ) const ( @@ -154,7 +152,7 @@ type Client struct { errorHooks []ErrorHook invalidHooks []ErrorHook panicHooks []ErrorHook - rateLimiter *rate.Limiter + rateLimiter RateLimiter } // User type is to hold an username and password information @@ -923,9 +921,9 @@ func (c *Client) SetOutputDirectory(dirPath string) *Client { return c } -// SetRateLimiter sets an optional `*rate.Limiter`. If set the rate limiter will control +// SetRateLimiter sets an optional `RateLimiter`. If set the rate limiter will control // all requests made with this client. -func (c *Client) SetRateLimiter(rl *rate.Limiter) *Client { +func (c *Client) SetRateLimiter(rl RateLimiter) *Client { c.rateLimiter = rl return c } diff --git a/errors.go b/errors.go deleted file mode 100644 index f8efe7ea..00000000 --- a/errors.go +++ /dev/null @@ -1,7 +0,0 @@ -package resty - -import "errors" - -var ( - ErrRateLimitExceeded = errors.New("rate limit exceeded") -) diff --git a/util.go b/util.go index 99c203b0..3279116f 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 //_______________________________________________________________________