ai-cache update body buffer limit size (#1644)

This commit is contained in:
澄潭
2025-01-06 14:53:29 +08:00
committed by GitHub
parent 187a7b5408
commit 2501895b66
3 changed files with 8 additions and 2 deletions

View File

@@ -22,6 +22,8 @@ const (
STREAM_CONTEXT_KEY = "stream" STREAM_CONTEXT_KEY = "stream"
SKIP_CACHE_HEADER = "x-higress-skip-ai-cache" SKIP_CACHE_HEADER = "x-higress-skip-ai-cache"
ERROR_PARTIAL_MESSAGE_KEY = "errorPartialMessage" ERROR_PARTIAL_MESSAGE_KEY = "errorPartialMessage"
DEFAULT_MAX_BODY_BYTES uint32 = 10 * 1024 * 1024
) )
func main() { func main() {
@@ -69,6 +71,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, c config.PluginConfig, log wr
ctx.DontReadRequestBody() ctx.DontReadRequestBody()
return types.ActionContinue return types.ActionContinue
} }
ctx.SetRequestBodyBufferLimit(DEFAULT_MAX_BODY_BYTES)
_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding") _ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
// The request has a body and requires delaying the header transmission until a cache miss occurs, // The request has a body and requires delaying the header transmission until a cache miss occurs,
// at which point the header should be sent. // 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") contentType, _ := proxywasm.GetHttpResponseHeader("content-type")
if strings.Contains(contentType, "text/event-stream") { if strings.Contains(contentType, "text/event-stream") {
ctx.SetContext(STREAM_CONTEXT_KEY, struct{}{}) ctx.SetContext(STREAM_CONTEXT_KEY, struct{}{})
} else {
ctx.SetResponseBodyBufferLimit(DEFAULT_MAX_BODY_BYTES)
} }
if ctx.GetContext(ERROR_PARTIAL_MESSAGE_KEY) != nil { if ctx.GetContext(ERROR_PARTIAL_MESSAGE_KEY) != nil {

View File

@@ -261,6 +261,7 @@ func checkStream(ctx *wrapper.HttpContext, log wrapper.Log) {
log.Errorf("unable to load content-type header from response: %v", err) log.Errorf("unable to load content-type header from response: %v", err)
} }
(*ctx).BufferResponseBody() (*ctx).BufferResponseBody()
ctx.SetResponseBodyBufferLimit(defaultMaxBodyBytes)
} }
} }

View File

@@ -65,9 +65,9 @@ type HttpContext interface {
// You need to call this before making any header modification operations. // You need to call this before making any header modification operations.
DisableReroute() DisableReroute()
// Note that this parameter affects the gateway's memory usageSupport setting a maximum buffer size for each request body individually in request phase. // Note that this parameter affects the gateway's memory usageSupport 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. // 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 type ParseConfigFunc[PluginConfig any] func(json gjson.Result, config *PluginConfig, log Log) error