feat: support retry on http status code

alibaba · Feb 25, 2025 · a3cd560 · a3cd560
1 parent e27d3d0
commit a3cd560
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 9 deletions.
diff --git a/plugins/wasm-go/extensions/ai-proxy/README.md b/plugins/wasm-go/extensions/ai-proxy/README.md
@@ -92,11 +92,12 @@ custom-setting会遵循如下表格，根据`name`和协议来替换对应的字
 
 `retryOnFailure` 的配置字段说明如下：
 
-| 名称               | 数据类型   | 填写要求            | 默认值   | 描述          |
-|------------------|--------|-----------------|-------|-------------|
-| enabled | bool   | 非必填             | false | 是否启用失败请求重试  |
-| maxRetries | int    | 非必填             | 1     | 最大重试次数      |
-| retryTimeout | int    | 非必填             | 30000 | 重试超时时间，单位毫秒 |
+| 名称               | 数据类型   | 填写要求   | 默认值   | 描述          |
+|------------------|--------|--------|-------|-------------|
+| enabled | bool   | 非必填    | false | 是否启用失败请求重试  |
+| maxRetries | int    | 非必填    | 1     | 最大重试次数      |
+| retryTimeout | int    | 非必填    | 30000 | 重试超时时间，单位毫秒 |
+| retryOnStatus | array of string | 非必填    |      |     需要进行重试的原始请求的状态码，支持正则表达式匹配     |
 
 ### 提供商特有配置
 

diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go
@@ -187,7 +187,7 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, pluginConfig config.PluginCo
 			log.Errorf("unable to load :status header from response: %v", err)
 		}
 		ctx.DontReadResponseBody()
-		return providerConfig.OnRequestFailed(activeProvider, ctx, apiTokenInUse, log)
+		return providerConfig.OnRequestFailed(activeProvider, ctx, apiTokenInUse, status, log)
 	}
 
 	// Reset ctxApiTokenRequestFailureCount if the request is successful,

diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/failover.go b/plugins/wasm-go/extensions/ai-proxy/provider/failover.go
@@ -539,17 +539,21 @@ func (c *ProviderConfig) resetSharedData() {
 	_ = proxywasm.SetSharedData(c.failover.ctxApiTokenRequestFailureCount, nil, 0)
 }
 
-func (c *ProviderConfig) OnRequestFailed(activeProvider Provider, ctx wrapper.HttpContext, apiTokenInUse string, log wrapper.Log) types.Action {
+func (c *ProviderConfig) OnRequestFailed(activeProvider Provider, ctx wrapper.HttpContext, apiTokenInUse string, status string, log wrapper.Log) types.Action {
 	if c.isFailoverEnabled() {
 		c.handleUnavailableApiToken(ctx, apiTokenInUse, log)
 	}
-	if c.isRetryOnFailureEnabled() && ctx.GetContext(ctxKeyIsStreaming) != nil && !ctx.GetContext(ctxKeyIsStreaming).(bool) {
+	if c.isRetryOnFailureEnabled() && c.matchRetryStatus(status) && isNotStreamingResponse(ctx) {
 		c.retryFailedRequest(activeProvider, ctx, log)
 		return types.HeaderStopAllIterationAndWatermark
 	}
 	return types.ActionContinue
 }
 
+func isNotStreamingResponse(ctx wrapper.HttpContext) bool {
+	return ctx.GetContext(ctxKeyIsStreaming) != nil && !ctx.GetContext(ctxKeyIsStreaming).(bool)
+}
+
 func (c *ProviderConfig) GetApiTokenInUse(ctx wrapper.HttpContext) string {
 	token, _ := ctx.GetContext(c.failover.ctxApiTokenInUse).(string)
 	return token

diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/retry.go b/plugins/wasm-go/extensions/ai-proxy/provider/retry.go
@@ -1,11 +1,13 @@
 package provider
 
 import (
+	"net/http"
+	"regexp"
+
 	"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/util"
 	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
 	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
 	"github.com/tidwall/gjson"
-	"net/http"
 )
 
 const (
@@ -20,6 +22,8 @@ type retryOnFailure struct {
 	maxRetries int64 `required:"false" yaml:"maxRetries" json:"maxRetries"`
 	// @Title zh-CN 重试超时时间
 	retryTimeout int64 `required:"false" yaml:"retryTimeout" json:"retryTimeout"`
+	// @Title zh-CN 需要进行重试的原始请求的状态码，支持正则表达式匹配
+	retryOnStatus []string `required:"false" yaml:"retryOnStatus" json:"retryOnStatus"`
 }
 
 func (r *retryOnFailure) FromJson(json gjson.Result) {
@@ -32,12 +36,29 @@ func (r *retryOnFailure) FromJson(json gjson.Result) {
 	if r.retryTimeout == 0 {
 		r.retryTimeout = 30 * 1000
 	}
+	for _, status := range json.Get("retryOnStatus").Array() {
+		r.retryOnStatus = append(r.retryOnStatus, status.String())
+	}
+	// If retryOnStatus is empty, default to retry on 4xx and 5xx
+	if len(r.retryOnStatus) == 0 {
+		r.retryOnStatus = []string{"4.*", "5.*"}
+	}
 }
 
 func (c *ProviderConfig) isRetryOnFailureEnabled() bool {
 	return c.retryOnFailure.enabled
 }
 
+func (c *ProviderConfig) matchRetryStatus(status string) bool {
+	for _, pattern := range c.retryOnFailure.retryOnStatus {
+		matched, _ := regexp.MatchString(pattern, status)
+		if matched {
+			return true
+		}
+	}
+	return false
+}
+
 func (c *ProviderConfig) retryFailedRequest(activeProvider Provider, ctx wrapper.HttpContext, log wrapper.Log) {
 	log.Debugf("Retry failed request: provider=%s", activeProvider.GetProviderType())
 	retryClient := createRetryClient(ctx)