Created
June 15, 2021 14:48
-
-
Save zdebra/10f0e284c4672e99f0cb767298f20c11 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 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), | |
} | |
} |
This is a great solution for rate limiting. Thanks for sharing!
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
That is a really helpful and great code!!
Just as a hint: in line 27 it shall be
10*time.Second
instead of10*time.Seconds
without the suffix s.