From 227ac4b8fcc4e906d129ab17129961d7d3a4b253 Mon Sep 17 00:00:00 2001 From: Anton Novojilov Date: Wed, 13 Nov 2024 00:30:11 +0300 Subject: [PATCH 1/2] [req] Add request limiter --- CHANGELOG.md | 4 ++++ req/example_test.go | 1 + req/req.go | 25 +++++++++++++++++++++++-- req/req_test.go | 10 ++++++++++ version.go | 2 +- 5 files changed, 39 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f0b27f3c..4ae6c6e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## Changelog +### [13.11.0](https://kaos.sh/ek/13.11.0) + +* `[req]` Added request limiter + ### [13.10.1](https://kaos.sh/ek/13.10.1) * `[mathutil]` Added shorthand helper `B` diff --git a/req/example_test.go b/req/example_test.go index f08809ff..1721ed33 100644 --- a/req/example_test.go +++ b/req/example_test.go @@ -24,6 +24,7 @@ func ExampleRequest_Do() { SetUserAgent("my-supper-app", "1.0") SetDialTimeout(30.0) SetRequestTimeout(30.0) + SetLimit(15.0) resp, err := Request{ Method: GET, diff --git a/req/req.go b/req/req.go index 4670858b..b1168d1a 100644 --- a/req/req.go +++ b/req/req.go @@ -224,8 +224,9 @@ type Engine struct { Transport *http.Transport // Transport is default transport struct Client *http.Client // Client is default client struct - dialTimeout float64 // dialTimeout is dial timeout in seconds - requestTimeout float64 // requestTimeout is request timeout in seconds + limiter *limiter // Request limiter + dialTimeout float64 // dialTimeout is dial timeout in seconds + requestTimeout float64 // requestTimeout is request timeout in seconds initialized bool } @@ -283,6 +284,12 @@ func SetRequestTimeout(timeout float64) { Global.SetRequestTimeout(timeout) } +// SetLimit sets a hard limit on the number of requests per second (useful for +// working with APIs) +func SetLimit(rps float64) { + Global.SetLimit(rps) +} + // ////////////////////////////////////////////////////////////////////////////////// // // Init initializes engine @@ -423,6 +430,16 @@ func (e *Engine) SetRequestTimeout(timeout float64) { } } +// SetLimit sets a hard limit on the number of requests per second (useful for +// working with APIs) +func (e *Engine) SetLimit(rps float64) { + if e == nil { + return + } + + e.limiter = createLimiter(rps) +} + // ////////////////////////////////////////////////////////////////////////////////// // // Do sends request and process response @@ -606,6 +623,10 @@ func (e *Engine) doRequest(r Request, method string) (*Response, error) { return nil, err } + if e.limiter != nil { + e.limiter.Wait() + } + resp, err := e.Client.Do(req) if err != nil { diff --git a/req/req_test.go b/req/req_test.go index b361c104..1ec28a88 100644 --- a/req/req_test.go +++ b/req/req_test.go @@ -81,6 +81,7 @@ func (s *ReqSuite) SetUpSuite(c *C) { SetDialTimeout(60.0) SetRequestTimeout(60.0) + SetLimit(1000.0) SetUserAgent("req-test", "5", "Test/5.1.1", "Magic/4.2.1") go runHTTPServer(s, c) @@ -580,6 +581,14 @@ func (s *ReqSuite) TestQueryEncoding(c *C) { c.Assert(qrs, Equals, "a=1&b=abcd&c&d") } +func (s *ReqSuite) TestLimiter(c *C) { + var l *limiter + + c.Assert(createLimiter(0.0), IsNil) + + l.Wait() +} + func (s *ReqSuite) TestNil(c *C) { var e *Engine @@ -589,6 +598,7 @@ func (s *ReqSuite) TestNil(c *C) { c.Assert(func() { e.SetUserAgent("APP", "1") }, NotPanics) c.Assert(func() { e.SetDialTimeout(1) }, NotPanics) c.Assert(func() { e.SetRequestTimeout(1) }, NotPanics) + c.Assert(func() { e.SetLimit(1.0) }, NotPanics) var r *Response diff --git a/version.go b/version.go index 2c240142..21a80d34 100644 --- a/version.go +++ b/version.go @@ -8,4 +8,4 @@ package ek // ////////////////////////////////////////////////////////////////////////////////// // // VERSION is current ek package version -const VERSION = "13.10.1" +const VERSION = "13.11.0" From 88d78a12b75a1da1242d624b1921355913a3f3a3 Mon Sep 17 00:00:00 2001 From: Anton Novojilov Date: Wed, 13 Nov 2024 00:30:28 +0300 Subject: [PATCH 2/2] [req] Add request limiter --- req/limiter.go | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 req/limiter.go diff --git a/req/limiter.go b/req/limiter.go new file mode 100644 index 00000000..210a8142 --- /dev/null +++ b/req/limiter.go @@ -0,0 +1,53 @@ +package req + +// ////////////////////////////////////////////////////////////////////////////////// // +// // +// Copyright (c) 2024 ESSENTIAL KAOS // +// Apache License, Version 2.0 // +// // +// ////////////////////////////////////////////////////////////////////////////////// // + +import "time" + +// ////////////////////////////////////////////////////////////////////////////////// // + +// limiter is request limiter +type limiter struct { + lastCall time.Time + delay time.Duration +} + +// ////////////////////////////////////////////////////////////////////////////////// // + +// createLimiter creates new limiter +func createLimiter(rps float64) *limiter { + if rps <= 0 { + return nil + } + + return &limiter{ + delay: time.Duration(float64(time.Second) / rps), + } +} + +// ////////////////////////////////////////////////////////////////////////////////// // + +// Wait blocks current goroutine execution until next time slot become available +func (l *limiter) Wait() { + if l == nil { + return + } + + if l.lastCall.IsZero() { + l.lastCall = time.Now() + return + } + + w := time.Since(l.lastCall) + + if w < l.delay { + time.Sleep(l.delay - w) + } + + l.lastCall = time.Now() +}