diff --git a/myhttp/httpc/httpclient.go b/myhttp/httpc/httpclient.go index 275c29a..119a82f 100644 --- a/myhttp/httpc/httpclient.go +++ b/myhttp/httpc/httpclient.go @@ -7,6 +7,7 @@ import ( "fmt" "git.makemake.in/kzkzzzz/mycommon/mylog" jsoniter "github.com/json-iterator/go" + "golang.org/x/time/rate" "io" "net/http" "net/url" @@ -64,11 +65,13 @@ type Request struct { urlQuery url.Values httpClient *HttpClient contentType string + noWaitQps bool } type HttpClient struct { - config *Config - client *http.Client + config *Config + client *http.Client + qpsLimiter *rate.Limiter } func New(opts ...ConfigOpt) *HttpClient { @@ -103,6 +106,12 @@ func New(opts ...ConfigOpt) *HttpClient { client: config.client, } + if config.qpsLimiter == nil { + hc.qpsLimiter = config.qpsLimiter + } else if config.qps > 0 { + hc.qpsLimiter = rate.NewLimiter(rate.Every(time.Second/time.Duration(config.qps)), config.qps) + } + return hc } @@ -188,6 +197,11 @@ func (r *Request) SetHeaders(headers map[string]string) *Request { return r } +func (r *Request) NoWaitQps() *Request { + r.noWaitQps = true + return r +} + func (r *Request) Get(rawUrl string) (*Response, error) { return r.Do(http.MethodGet, rawUrl) } @@ -196,7 +210,23 @@ func (r *Request) Post(rawUrl string) (*Response, error) { return r.Do(http.MethodPost, rawUrl) } +var QpsLimitError = fmt.Errorf("qps limit") + func (r *Request) Do(method, rawUrl string) (*Response, error) { + if r.httpClient.qpsLimiter != nil { + if r.noWaitQps { + allow := r.httpClient.qpsLimiter.Allow() + if !allow { + return nil, QpsLimitError + } + } else { + err := r.httpClient.qpsLimiter.Wait(r.ctx) + if err != nil { + return nil, err + } + } + } + reqUrl, err := url.Parse(rawUrl) if err != nil { return nil, err diff --git a/myhttp/httpc/option.go b/myhttp/httpc/option.go index 61b8e8c..f09b4f2 100644 --- a/myhttp/httpc/option.go +++ b/myhttp/httpc/option.go @@ -1,6 +1,7 @@ package httpc import ( + "golang.org/x/time/rate" "net/http" "time" ) @@ -12,6 +13,8 @@ type ( transport *http.Transport redirectFn func(req *http.Request, via []*http.Request) error noCheckStatus bool + qps int + qpsLimiter *rate.Limiter } ConfigOpt func(c *Config) @@ -54,3 +57,15 @@ func WithNoCheckStatus(v bool) ConfigOpt { c.noCheckStatus = v } } + +func WithQps(v int) ConfigOpt { + return func(c *Config) { + c.qps = v + } +} + +func WithQpsLimiter(v *rate.Limiter) ConfigOpt { + return func(c *Config) { + c.qpsLimiter = v + } +}