update
parent
eb6557e0d8
commit
781f79cf7c
|
@ -7,6 +7,7 @@ import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"git.makemake.in/kzkzzzz/mycommon/mylog"
|
"git.makemake.in/kzkzzzz/mycommon/mylog"
|
||||||
jsoniter "github.com/json-iterator/go"
|
jsoniter "github.com/json-iterator/go"
|
||||||
|
"golang.org/x/time/rate"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
|
@ -64,11 +65,13 @@ type Request struct {
|
||||||
urlQuery url.Values
|
urlQuery url.Values
|
||||||
httpClient *HttpClient
|
httpClient *HttpClient
|
||||||
contentType string
|
contentType string
|
||||||
|
noWaitQps bool
|
||||||
}
|
}
|
||||||
|
|
||||||
type HttpClient struct {
|
type HttpClient struct {
|
||||||
config *Config
|
config *Config
|
||||||
client *http.Client
|
client *http.Client
|
||||||
|
qpsLimiter *rate.Limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
func New(opts ...ConfigOpt) *HttpClient {
|
func New(opts ...ConfigOpt) *HttpClient {
|
||||||
|
@ -103,6 +106,12 @@ func New(opts ...ConfigOpt) *HttpClient {
|
||||||
client: config.client,
|
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
|
return hc
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -188,6 +197,11 @@ func (r *Request) SetHeaders(headers map[string]string) *Request {
|
||||||
return r
|
return r
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (r *Request) NoWaitQps() *Request {
|
||||||
|
r.noWaitQps = true
|
||||||
|
return r
|
||||||
|
}
|
||||||
|
|
||||||
func (r *Request) Get(rawUrl string) (*Response, error) {
|
func (r *Request) Get(rawUrl string) (*Response, error) {
|
||||||
return r.Do(http.MethodGet, rawUrl)
|
return r.Do(http.MethodGet, rawUrl)
|
||||||
}
|
}
|
||||||
|
@ -196,7 +210,23 @@ func (r *Request) Post(rawUrl string) (*Response, error) {
|
||||||
return r.Do(http.MethodPost, rawUrl)
|
return r.Do(http.MethodPost, rawUrl)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
var QpsLimitError = fmt.Errorf("qps limit")
|
||||||
|
|
||||||
func (r *Request) Do(method, rawUrl string) (*Response, error) {
|
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)
|
reqUrl, err := url.Parse(rawUrl)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
|
|
|
@ -1,6 +1,7 @@
|
||||||
package httpc
|
package httpc
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"golang.org/x/time/rate"
|
||||||
"net/http"
|
"net/http"
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
@ -12,6 +13,8 @@ type (
|
||||||
transport *http.Transport
|
transport *http.Transport
|
||||||
redirectFn func(req *http.Request, via []*http.Request) error
|
redirectFn func(req *http.Request, via []*http.Request) error
|
||||||
noCheckStatus bool
|
noCheckStatus bool
|
||||||
|
qps int
|
||||||
|
qpsLimiter *rate.Limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
ConfigOpt func(c *Config)
|
ConfigOpt func(c *Config)
|
||||||
|
@ -54,3 +57,15 @@ func WithNoCheckStatus(v bool) ConfigOpt {
|
||||||
c.noCheckStatus = v
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue