From 2501895b6618a2e88bc8cb426a3d8d65b1280fba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BE=84=E6=BD=AD?= Date: Mon, 6 Jan 2025 14:53:29 +0800 Subject: [PATCH] ai-cache update body buffer limit size (#1644) --- plugins/wasm-go/extensions/ai-cache/main.go | 5 +++++ plugins/wasm-go/extensions/ai-proxy/main.go | 1 + plugins/wasm-go/pkg/wrapper/plugin_wrapper.go | 4 ++-- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-cache/main.go b/plugins/wasm-go/extensions/ai-cache/main.go index 650d3805ac..0ab323de13 100644 --- a/plugins/wasm-go/extensions/ai-cache/main.go +++ b/plugins/wasm-go/extensions/ai-cache/main.go @@ -22,6 +22,8 @@ const ( STREAM_CONTEXT_KEY = "stream" SKIP_CACHE_HEADER = "x-higress-skip-ai-cache" ERROR_PARTIAL_MESSAGE_KEY = "errorPartialMessage" + + DEFAULT_MAX_BODY_BYTES uint32 = 10 * 1024 * 1024 ) func main() { @@ -69,6 +71,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, c config.PluginConfig, log wr ctx.DontReadRequestBody() return types.ActionContinue } + ctx.SetRequestBodyBufferLimit(DEFAULT_MAX_BODY_BYTES) _ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding") // The request has a body and requires delaying the header transmission until a cache miss occurs, // at which point the header should be sent. @@ -140,6 +143,8 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, c config.PluginConfig, log w contentType, _ := proxywasm.GetHttpResponseHeader("content-type") if strings.Contains(contentType, "text/event-stream") { ctx.SetContext(STREAM_CONTEXT_KEY, struct{}{}) + } else { + ctx.SetResponseBodyBufferLimit(DEFAULT_MAX_BODY_BYTES) } if ctx.GetContext(ERROR_PARTIAL_MESSAGE_KEY) != nil { diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go index 6c7756e457..c08dcdb16f 100644 --- a/plugins/wasm-go/extensions/ai-proxy/main.go +++ b/plugins/wasm-go/extensions/ai-proxy/main.go @@ -261,6 +261,7 @@ func checkStream(ctx *wrapper.HttpContext, log wrapper.Log) { log.Errorf("unable to load content-type header from response: %v", err) } (*ctx).BufferResponseBody() + ctx.SetResponseBodyBufferLimit(defaultMaxBodyBytes) } } diff --git a/plugins/wasm-go/pkg/wrapper/plugin_wrapper.go b/plugins/wasm-go/pkg/wrapper/plugin_wrapper.go index 8b342d57b5..0e3fc14782 100644 --- a/plugins/wasm-go/pkg/wrapper/plugin_wrapper.go +++ b/plugins/wasm-go/pkg/wrapper/plugin_wrapper.go @@ -65,9 +65,9 @@ type HttpContext interface { // You need to call this before making any header modification operations. DisableReroute() // Note that this parameter affects the gateway's memory usageļ¼Support setting a maximum buffer size for each request body individually in request phase. - SetRequestBodyBufferLimit(size uint32) + SetRequestBodyBufferLimit(byteSize uint32) // Note that this parameter affects the gateway's memory usage! Support setting a maximum buffer size for each response body individually in response phase. - SetResponseBodyBufferLimit(size uint32) + SetResponseBodyBufferLimit(byteSize uint32) } type ParseConfigFunc[PluginConfig any] func(json gjson.Result, config *PluginConfig, log Log) error