From 4767aeb50419729c93cadf416d70f7c8284f1b5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=99=BD=E8=8C=B6=E6=B8=85=E6=AC=A2?=
 <go_developer@163.com>
Date: Thu, 8 May 2025 11:37:34 +0800
Subject: [PATCH] =?UTF-8?q?=E9=99=90=E6=B5=81=E5=99=A8=E4=BD=BF=E7=94=A8?=
 =?UTF-8?q?=E4=B8=AD=E9=97=B4=E4=BB=B6=E5=AE=8C=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 httpclient/client.go       | 22 ++++++++--------------
 httpclient/define/error.go |  5 +++++
 httpclient/resty.go        | 15 ++++++++++++++-
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/httpclient/client.go b/httpclient/client.go
index d415aa0..273e969 100644
--- a/httpclient/client.go
+++ b/httpclient/client.go
@@ -97,9 +97,11 @@ func NewHttpClient(reqConfig *define.Request, reqOption *RequestOption) (*HttpCl
 		reqConfig.Header = make(map[string]any)
 	}
 	if ua, exist := reqConfig.Header[consts.HeaderKeyUserAgent.String()]; !exist || nil == ua || fmt.Sprintf("%v", ua) == "" {
-		reqConfig.Header[consts.HeaderKeyUserAgent.String()] = "resty-v2@network/httpclient"
+		reqConfig.Header[consts.HeaderKeyUserAgent.String()] = "resty-v3@network/httpclient"
 	}
-	restyClient, restyRequest := NewRestyClient(reqConfig)
+	restyClient, restyRequest := NewRestyClient(reqConfig, reqOption)
+	defer restyClient.Close()
+
 	hc := &HttpClient{
 		Client:               restyClient,
 		request:              restyRequest,
@@ -296,20 +298,12 @@ func (hc *HttpClient) requestBackendApi() *define.Response {
 		response = hc.newResponse()
 		response.Seq = i
 		response.RequestCount = i + 1
-		if nil != hc.reqOption.RateLimiter {
-			if !hc.reqOption.RateLimiter.Allow() {
-				response.FailInfo = &define.ResponseFailInfo{
-					Type:    define.RequestFailTypeRateLimit,
-					Message: "rate limit exceeded",
-				}
-				// 命中限流, 忽略重试
-				log.RecordDebug("请求命中限流, 忽略重试策略, 不进行重试", nil, hc.reqCfg)
-				break
-			}
-		}
 		if response.RestyResponse, err = hc.request.Send(); nil != err {
 			errType := define.RequestFailTypeSend
-			if netErr, ok := err.(net.Error); ok {
+			if err.Error() == define.ErrRateLimitExceeded.Error() {
+				// 命中限流
+				errType = define.RequestFailTypeRateLimit
+			} else if netErr, ok := err.(net.Error); ok {
 				if netErr.Timeout() {
 					// 请求超时
 					errType = define.RequestFailTypeTimeoutError
diff --git a/httpclient/define/error.go b/httpclient/define/error.go
index bee6f77..f9b327b 100644
--- a/httpclient/define/error.go
+++ b/httpclient/define/error.go
@@ -16,3 +16,8 @@ var (
 	ErrMethodIsEmpty      = errors.New("METHOD_IS_EMPTY")       // 请求 方法为空
 	ErrMethodIsNotSupport = errors.New("METHOD_IS_NOT_SUPPORT") // 请求 method不支持
 )
+
+var (
+	// ErrRateLimitExceeded v2版本中命中流控的错误信息
+	ErrRateLimitExceeded = errors.New("rate limit exceeded")
+)
diff --git a/httpclient/resty.go b/httpclient/resty.go
index 987883c..1cd9757 100644
--- a/httpclient/resty.go
+++ b/httpclient/resty.go
@@ -75,12 +75,25 @@ func initRequestConfig(reqConfig *define.Request) {
 // Author : go_developer@163.com<白茶清欢>
 //
 // Date : 15:00 2024/5/31
-func NewRestyClient(reqConfig *define.Request) (*resty.Client, *resty.Request) {
+func NewRestyClient(reqConfig *define.Request, reqOption *RequestOption) (*resty.Client, *resty.Request) {
 	client := resty.New()
 	request := client.R()
 	if nil == reqConfig {
 		return client, request
 	}
+	// 限流处理, 增加限流中间件
+	client.AddRequestMiddleware(func(client *resty.Client, request *resty.Request) error {
+		if nil != reqOption && nil != reqOption.RateLimiter {
+			// 未配置流控
+			return nil
+		}
+		if !reqOption.RateLimiter.Allow() {
+			// 命中流控
+			return define.ErrRateLimitExceeded
+		}
+		return nil
+	})
+	
 	initRequestConfig(reqConfig)             // 初始化 + 格式化配置
 	client.SetAllowMethodGetPayload(true)    // 配置 GET 请求允许带 Body
 	client.SetAllowMethodDeletePayload(true) // 配置 DELETE 请求允许带 Body