main
kzkzzzz 2025-07-27 10:08:41 +08:00
parent eb6557e0d8
commit 781f79cf7c
2 changed files with 47 additions and 2 deletions

View File

@ -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

View File

@ -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
}
}