Skip to content

Instantly share code, notes, and snippets.

@r6m
Forked from zdebra/throttled_transport.go
Created August 28, 2022 09:49
Show Gist options
  • Save r6m/75b902604e51f0f3ad216ff10362b4fa to your computer and use it in GitHub Desktop.
Save r6m/75b902604e51f0f3ad216ff10362b4fa 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
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.Second, 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