-
-
Save lucasepe/2f49a61c56f8130165008f537234faaf to your computer and use it in GitHub Desktop.
NewThrottledTransport wraps transportWrap with a rate limitter, improvement of https://gist.github.com/MelchiSalins/27c11566184116ec1629a0726e0f9af5 since it allows use of *http.Client
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| package main | |
| import ( | |
| "net/http" | |
| "time" | |
| "golang.org/x/time/rate" | |
| ) | |
| // ThrottledTransport Rate Limited HTTP Client | |
| type ThrottledTransport struct { | |
| roundTripperWrap http.RoundTripper | |
| ratelimiter *rate.Limiter | |
| } | |
| func (c *ThrottledTransport) RoundTrip(r *http.Request) (*http.Response, error) { | |
| err := c.ratelimiter.Wait(r.Context()) // This is a blocking call. Honors the rate limit | |
| if err != nil { | |
| return nil, err | |
| } | |
| return c.roundTripperWrap.RoundTrip(r) | |
| } | |
| // NewThrottledTransport wraps transportWrap with a rate limitter | |
| // examle usage: | |
| // client := http.DefaultClient | |
| // client.Transport = NewThrottledTransport(10*time.Seconds, 60, http.DefaultTransport) allows 60 requests every 10 seconds | |
| func NewThrottledTransport(limitPeriod time.Duration, requestCount int, transportWrap http.RoundTripper) http.RoundTripper { | |
| return &ThrottledTransport{ | |
| roundTripperWrap: transportWrap, | |
| ratelimiter: rate.NewLimiter(rate.Every(limitPeriod), requestCount), | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment