From 17cb83f8e3da3c05259a01cf540394368de5b412 Mon Sep 17 00:00:00 2001 From: Lukas Jenicek Date: Mon, 13 Jan 2025 14:49:27 +0100 Subject: [PATCH] rate limited http transport --- _example/client/go.mod | 15 ++++++++++ _example/client/go.sum | 9 ++++++ _example/client/main.go | 64 +++++++++++++++++++++++++++++++++++++++++ _example/go.mod | 6 +++- _example/go.sum | 5 ++++ go.mod | 9 ++++-- go.sum | 7 +++-- transport.go | 28 ++++++++++++++++++ 8 files changed, 137 insertions(+), 6 deletions(-) create mode 100644 _example/client/go.mod create mode 100644 _example/client/go.sum create mode 100644 _example/client/main.go create mode 100644 transport.go diff --git a/_example/client/go.mod b/_example/client/go.mod new file mode 100644 index 0000000..9d7e985 --- /dev/null +++ b/_example/client/go.mod @@ -0,0 +1,15 @@ +module github.com/go-chi/httprate/_example + +go 1.23.2 + +replace github.com/go-chi/httprate => ../../ + +require ( + github.com/go-chi/httprate v0.0.0-00010101000000-000000000000 + github.com/go-chi/transport v0.5.0 +) + +require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect + golang.org/x/time v0.9.0 // indirect +) diff --git a/_example/client/go.sum b/_example/client/go.sum new file mode 100644 index 0000000..f652792 --- /dev/null +++ b/_example/client/go.sum @@ -0,0 +1,9 @@ +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/go-chi/transport v0.5.0 h1:xpnYcIOpBRrduJD68gX9YxkJouRGIE1y+rK5yGYnMXE= +github.com/go-chi/transport v0.5.0/go.mod h1:uoCleTaQiFtoatEiiqcXFZ5OxIp6s1DfGeVsCVbalT4= +golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= +golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= +golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY= +golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= diff --git a/_example/client/main.go b/_example/client/main.go new file mode 100644 index 0000000..28d84b2 --- /dev/null +++ b/_example/client/main.go @@ -0,0 +1,64 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + "github.com/go-chi/httprate" + "github.com/go-chi/transport" +) + +type loginPayload struct { + Username string `json:"username"` + Password string `json:"password"` +} + +func main() { + ctx := context.Background() + cl := &http.Client{ + Transport: transport.Chain(http.DefaultTransport, httprate.RateLimitedRequest(httprate.RPMLimit(10), 1)), + } + + req := &loginPayload{ + Username: "alice", + Password: "password", + } + + payload, err := json.Marshal(req) + if err != nil { + log.Fatal(err) + } + + for { + func() { + ctx, cancel := context.WithTimeout(ctx, 60*time.Second) + defer cancel() + + // login accepts only 5req/mint so it gets rate limited eventually + req, err := http.NewRequestWithContext(ctx, "POST", "http://localhost:3333/login", bytes.NewReader(payload)) + if err != nil { + log.Fatal(err) + } + + fmt.Println("request started", time.Now()) + + resp, err := cl.Do(req) + if err != nil { + log.Fatal(err) + } + + defer resp.Body.Close() + + if resp.StatusCode == http.StatusTooManyRequests { + fmt.Println("rate limited") + } + + fmt.Println() + }() + } +} diff --git a/_example/go.mod b/_example/go.mod index e24b12b..70201dc 100644 --- a/_example/go.mod +++ b/_example/go.mod @@ -9,4 +9,8 @@ require ( github.com/go-chi/httprate v0.0.0-00010101000000-000000000000 ) -require github.com/cespare/xxhash/v2 v2.3.0 // indirect +require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/go-chi/transport v0.5.0 // indirect + golang.org/x/time v0.9.0 // indirect +) diff --git a/_example/go.sum b/_example/go.sum index 29685bd..06f99b7 100644 --- a/_example/go.sum +++ b/_example/go.sum @@ -2,5 +2,10 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UF github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/go-chi/chi/v5 v5.1.0 h1:acVI1TYaD+hhedDJ3r54HyA6sExp3HfXq7QWEEY/xMw= github.com/go-chi/chi/v5 v5.1.0/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8= +github.com/go-chi/transport v0.5.0 h1:xpnYcIOpBRrduJD68gX9YxkJouRGIE1y+rK5yGYnMXE= +github.com/go-chi/transport v0.5.0/go.mod h1:uoCleTaQiFtoatEiiqcXFZ5OxIp6s1DfGeVsCVbalT4= +golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY= +golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= diff --git a/go.mod b/go.mod index 998cbf5..1bd3393 100644 --- a/go.mod +++ b/go.mod @@ -2,6 +2,9 @@ module github.com/go-chi/httprate go 1.17 -require github.com/cespare/xxhash/v2 v2.3.0 - -require golang.org/x/sync v0.7.0 // indirect +require ( + github.com/cespare/xxhash/v2 v2.3.0 + github.com/go-chi/transport v0.5.0 + golang.org/x/sync v0.7.0 + golang.org/x/time v0.9.0 +) diff --git a/go.sum b/go.sum index 09aebbf..f652792 100644 --- a/go.sum +++ b/go.sum @@ -1,6 +1,9 @@ -github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE= -github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/go-chi/transport v0.5.0 h1:xpnYcIOpBRrduJD68gX9YxkJouRGIE1y+rK5yGYnMXE= +github.com/go-chi/transport v0.5.0/go.mod h1:uoCleTaQiFtoatEiiqcXFZ5OxIp6s1DfGeVsCVbalT4= +golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY= +golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= diff --git a/transport.go b/transport.go new file mode 100644 index 0000000..69c3c7f --- /dev/null +++ b/transport.go @@ -0,0 +1,28 @@ +package httprate + +import ( + "fmt" + "net/http" + "time" + + "github.com/go-chi/transport" + "golang.org/x/time/rate" +) + +func RateLimitedRequest(limit rate.Limit, burst int) func(http.RoundTripper) http.RoundTripper { + limiter := rate.NewLimiter(limit, burst) + + return func(next http.RoundTripper) http.RoundTripper { + return transport.RoundTripFunc(func(req *http.Request) (resp *http.Response, err error) { + if err := limiter.Wait(req.Context()); err != nil { + return nil, fmt.Errorf("rate limiter wait: %w", err) + } + + return next.RoundTrip(req) + }) + } +} + +func RPMLimit(rpm uint) rate.Limit { + return rate.Every(time.Minute / time.Duration(rpm)) +}