diff --git a/httpclient/define/request.go b/httpclient/define/request.go index 70ea06b..0d8db94 100644 --- a/httpclient/define/request.go +++ b/httpclient/define/request.go @@ -9,6 +9,7 @@ package define import ( "context" + "github.com/go-resty/resty/v2" "go.uber.org/zap" ) @@ -37,6 +38,7 @@ type Request struct { ReadTimeout int64 `json:"read_timeout"` // 读取超时时间 RetryRule *RequestRetryRule `json:"retry_rule"` // 重试规则 Logger *zap.Logger `json:"-"` // 日志记录器 + RateLimiter resty.RateLimiter `json:"-"` // 流控实例 } // RequestRetryRule 重试规则 diff --git a/httpclient/resty.go b/httpclient/resty.go index faaf57c..8e2f9fc 100644 --- a/httpclient/resty.go +++ b/httpclient/resty.go @@ -31,6 +31,10 @@ func NewRestyClient(reqConfig *define.Request) (*resty.Client, *resty.Request) { if nil == reqConfig { return client, request } + if nil != reqConfig.RateLimiter { + // 设置流控实例 + client.SetRateLimiter(reqConfig.RateLimiter) + } formatHeader(reqConfig) client.SetAllowGetMethodPayload(true) // 配置 GET 请求允许带 Body client.SetJSONMarshaler(json.Marshal) // 序列化方法