Skip to content

phsl / new-api

Go to a project
Commit fddf54cc by CaIon
Options

perf: reduce heap residency for large base64 relay requests 

Three layered optimizations targeting Gemini-style 5MB base64 payloads where
RSS could balloon to tens of GB under concurrent load:

1. Byte-based param override (relay/common/override.go)
   - Switch legacy/operations hot paths from common.Marshal round-trips and
     map[string]any conversions to gjson/sjson on []byte directly.
   - Avoids cloning 5MB strings during each Set/Delete operation.

2. strings.Builder for Gemini response markdown (relay/channel/gemini/relay-gemini.go)
   - Replace string concatenation + strings.Join when assembling
     "![image](data:...;base64,DATA)" content for inline image responses.
   - Pre-allocates capacity from inline_data byte sizes.

3. Outbound BodyStorage + streaming Decoder (this commit's core)
   - New relay/common/outbound_body.go helper wraps marshaled upstream bodies
     in common.BodyStorage, allowing disk-cache mode to offload jsonData to
     a temp file while waiting for upstream TTFB. The original []byte can
     then be GC'd, removing ~5MB/req of heap residency during the longest
     window of a request.
   - All 7 relay handlers (gemini/claude/responses/embedding/image/compatible/
     rerank) plus chat_completions_via_responses adopt the helper with
     defer closer.Close() and explicit jsonData = nil.
   - relay/common/relay_info.go: new UpstreamRequestBodySize so
     relay/channel/api_request.go can populate req.ContentLength (lost when
     body becomes a type-erased io.Reader).
   - common/gin.go UnmarshalBodyReusable: when storage is disk-backed and
     content-type is JSON, decode via DecodeJson(storage) instead of
     storage.Bytes()+Unmarshal, removing one transient 5MB copy per request.
     memory mode and form/multipart paths unchanged.
parent b9bc6f0e
Showing with 239 additions and 36 deletions
common/custom-event.go
......@@ -37,7 +37,7 @@ func checkWriter(writer io.Writer) stringWriter {
// W3C Working Draft 29 October 2009
// http://www.w3.org/TR/2009/WD-eventsource-20091029/
var contentType = []string{"text/event-stream"}
var writeContentType = []string{"text/event-stream"}
var noCache = []string{"no-cache"}
var fieldReplacer = strings.NewReplacer(
......@@ -79,7 +79,7 @@ func (r CustomEvent) WriteContentType(w http.ResponseWriter) {
r.Mutex.Lock()
defer r.Mutex.Unlock()
header := w.Header()
header["Content-Type"] = contentType
header["Content-Type"] = writeContentType
if _, exist := header["Cache-Control"]; !exist {
header["Cache-Control"] = noCache
......
common/gin.go
......@@ -110,11 +110,29 @@ func UnmarshalBodyReusable(c *gin.Context, v any) error {
if err != nil {
return err
}
contentType := c.Request.Header.Get("Content-Type")
// disk-backed JSON: stream-decode directly from the file to avoid
// materializing the entire payload back into a transient []byte
// (diskStorage.Bytes() would ReadFull the whole file into the heap).
if storage.IsDisk() && strings.HasPrefix(contentType, "application/json") {
if _, seekErr := storage.Seek(0, io.SeekStart); seekErr != nil {
return seekErr
}
if err := DecodeJson(storage, v); err != nil {
return err
}
if _, seekErr := storage.Seek(0, io.SeekStart); seekErr != nil {
return seekErr
}
c.Request.Body = io.NopCloser(storage)
return nil
}
requestBody, err := storage.Bytes()
if err != nil {
return err
}
contentType := c.Request.Header.Get("Content-Type")
if strings.HasPrefix(contentType, "application/json") {
err = Unmarshal(requestBody, v)
} else if strings.Contains(contentType, gin.MIMEPOSTForm) {
......
relay/channel/api_request.go
......@@ -25,6 +25,23 @@ import (
"github.com/gorilla/websocket"
)
// applyUpstreamContentLength populates req.ContentLength when the upstream
// body is wrapped in a BodyStorage (see relay/common/outbound_body.go).
//
// net/http.NewRequest only auto-detects ContentLength for *bytes.Reader,
// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader
// (which is the case for ReaderOnly(BodyStorage)), the Content-Length header
// would otherwise be omitted, forcing chunked transfer encoding and breaking
// some upstreams that require an explicit Content-Length.
func applyUpstreamContentLength(req *http.Request, info *common.RelayInfo) {
if info == nil {
return
}
if info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 {
req.ContentLength = info.UpstreamRequestBodySize
}
}
func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) {
if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
// multipart/form-data
......@@ -297,6 +314,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
applyUpstreamContentLength(req, info)
headers := req.Header
err = a.SetupRequestHeader(c, &headers, info)
if err != nil {
......@@ -326,6 +344,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
applyUpstreamContentLength(req, info)
// set form data
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
headers := req.Header
......@@ -522,6 +541,7 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
applyUpstreamContentLength(req, info)
req.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(requestBody), nil
}
......
relay/channel/gemini/relay-gemini.go
......@@ -1079,17 +1079,47 @@ func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse)
FinishReason: constant.FinishReasonStop,
}
if len(candidate.Content.Parts) > 0 {
var texts []string
// 使用 strings.Builder 直接累积最终 content，避免:
// 1) 每张 inline image 生成一次中间 "![image](...)" 字符串
// 2) 末尾 strings.Join 再分配一份等大缓冲
// Gemini 图片返回时 InlineData.Data 可能是数 MB 的 base64，
// 上述两份临时分配在高并发下会显著放大堆驻留。
var content strings.Builder
var inlineGrow int
for _, part := range candidate.Content.Parts {
if part.InlineData != nil {
inlineGrow += len(part.InlineData.MimeType) + len(part.InlineData.Data) + 32
}
}
if inlineGrow > 0 {
content.Grow(inlineGrow)
}
appended := 0
writeSep := func() {
if appended > 0 {
content.WriteByte('\n')
}
appended++
}
var toolCalls []dto.ToolCallResponse
for _, part := range candidate.Content.Parts {
if part.InlineData != nil {
// 媒体内容
if strings.HasPrefix(part.InlineData.MimeType, "image") {
imgText := "![image](data:" + part.InlineData.MimeType + ";base64," + part.InlineData.Data + ")"
texts = append(texts, imgText)
writeSep()
content.WriteString("![image](data:")
content.WriteString(part.InlineData.MimeType)
content.WriteString(";base64,")
content.WriteString(part.InlineData.Data)
content.WriteByte(')')
} else {
// 其他媒体类型，直接显示链接
texts = append(texts, fmt.Sprintf("[media](data:%s;base64,%s)", part.InlineData.MimeType, part.InlineData.Data))
writeSep()
content.WriteString("[media](data:")
content.WriteString(part.InlineData.MimeType)
content.WriteString(";base64,")
content.WriteString(part.InlineData.Data)
content.WriteByte(')')
}
} else if part.FunctionCall != nil {
choice.FinishReason = constant.FinishReasonToolCalls
......@@ -1100,13 +1130,22 @@ func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse)
choice.Message.ReasoningContent = &part.Text
} else {
if part.ExecutableCode != nil {
texts = append(texts, "```"+part.ExecutableCode.Language+"\n"+part.ExecutableCode.Code+"\n```")
writeSep()
content.WriteString("```")
content.WriteString(part.ExecutableCode.Language)
content.WriteByte('\n')
content.WriteString(part.ExecutableCode.Code)
content.WriteString("\n```")
} else if part.CodeExecutionResult != nil {
texts = append(texts, "```output\n"+part.CodeExecutionResult.Output+"\n```")
writeSep()
content.WriteString("```output\n")
content.WriteString(part.CodeExecutionResult.Output)
content.WriteString("\n```")
} else {
// 过滤掉空行
if part.Text != "\n" {
texts = append(texts, part.Text)
writeSep()
content.WriteString(part.Text)
}
}
}
......@@ -1115,7 +1154,7 @@ func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse)
choice.Message.SetToolCalls(toolCalls)
isToolCall = true
}
choice.Message.SetStringContent(strings.Join(texts, "\n"))
choice.Message.SetStringContent(content.String())
}
if candidate.FinishReason != nil {
......@@ -1169,7 +1208,25 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d
//Role: "assistant",
},
}
var texts []string
// 使用 strings.Builder 直接累积 delta content，避免每张 image / 每个
// 文本片段都先 `+` 拼出一份临时 string，再 strings.Join 再拷贝一遍。
var content strings.Builder
var inlineGrow int
for _, part := range candidate.Content.Parts {
if part.InlineData != nil {
inlineGrow += len(part.InlineData.MimeType) + len(part.InlineData.Data) + 32
}
}
if inlineGrow > 0 {
content.Grow(inlineGrow)
}
appended := 0
writeSep := func() {
if appended > 0 {
content.WriteByte('\n')
}
appended++
}
isTools := false
isThought := false
if candidate.FinishReason != nil {
......@@ -1207,8 +1264,12 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d
for _, part := range candidate.Content.Parts {
if part.InlineData != nil {
if strings.HasPrefix(part.InlineData.MimeType, "image") {
imgText := "![image](data:" + part.InlineData.MimeType + ";base64," + part.InlineData.Data + ")"
texts = append(texts, imgText)
writeSep()
content.WriteString("![image](data:")
content.WriteString(part.InlineData.MimeType)
content.WriteString(";base64,")
content.WriteString(part.InlineData.Data)
content.WriteByte(')')
}
} else if part.FunctionCall != nil {
isTools = true
......@@ -1219,23 +1280,33 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d
} else if part.Thought {
isThought = true
texts = append(texts, part.Text)
writeSep()
content.WriteString(part.Text)
} else {
if part.ExecutableCode != nil {
texts = append(texts, "```"+part.ExecutableCode.Language+"\n"+part.ExecutableCode.Code+"\n```\n")
writeSep()
content.WriteString("```")
content.WriteString(part.ExecutableCode.Language)
content.WriteByte('\n')
content.WriteString(part.ExecutableCode.Code)
content.WriteString("\n```\n")
} else if part.CodeExecutionResult != nil {
texts = append(texts, "```output\n"+part.CodeExecutionResult.Output+"\n```\n")
writeSep()
content.WriteString("```output\n")
content.WriteString(part.CodeExecutionResult.Output)
content.WriteString("\n```\n")
} else {
if part.Text != "\n" {
texts = append(texts, part.Text)
writeSep()
content.WriteString(part.Text)
}
}
}
}
if isThought {
choice.Delta.SetReasoningContent(strings.Join(texts, "\n"))
choice.Delta.SetReasoningContent(content.String())
} else {
choice.Delta.SetContentString(strings.Join(texts, "\n"))
choice.Delta.SetContentString(content.String())
}
if isTools {
choice.FinishReason = &constant.FinishReasonToolCalls
......
relay/chat_completions_via_responses.go
package relay
import (
"bytes"
"io"
"net/http"
"strings"
......@@ -125,7 +124,14 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
var requestBody io.Reader = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
var requestBody io.Reader = body
var httpResp *http.Response
resp, err := adaptor.DoRequest(c, info, requestBody)
......
relay/claude_handler.go
package relay
import (
"bytes"
"encoding/json"
"fmt"
"io"
......@@ -179,7 +178,14 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
logger.LogDebug(c, "requestBody: %s", jsonData)
requestBody = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
}
statusCodeMappingStr := c.GetString("status_code_mapping")
......
relay/common/outbound_body.go 0 → 100644
package common
import (
"io"
"github.com/QuantumNous/new-api/common"
)
// NewOutboundJSONBody wraps the already-marshaled upstream request body into a
// BodyStorage. When disk cache is enabled and the payload exceeds the configured
// threshold, the data is written to a temp file and the original []byte can be
// GC'd, significantly reducing the heap residency while waiting for the
// upstream provider to respond (the dominant cost for large base64 payloads).
//
// In memory mode the underlying memoryStorage reuses the same backing array,
// so this is equivalent to bytes.NewReader(data) in terms of memory usage.
//
// The caller MUST invoke closer.Close() once the upstream call has finished
// (typically via defer) to release the disk file / memory accounting.
//
// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP
// transport from prematurely closing the underlying BodyStorage. The returned
// size is meant to be propagated to http.Request.ContentLength because the
// type-erased io.Reader prevents net/http from auto-detecting it.
func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, closer io.Closer, err error) {
storage, err := common.CreateBodyStorage(data)
if err != nil {
return nil, 0, nil, err
}
return common.ReaderOnly(storage), storage.Size(), storage, nil
}
relay/common/relay_info.go
......@@ -154,6 +154,13 @@ type RelayInfo struct {
UseRuntimeHeadersOverride bool
ParamOverrideAudit []string
// UpstreamRequestBodySize is the byte size of the marshaled upstream request
// body. It is set when the body is wrapped in a BodyStorage (see
// relay/common/outbound_body.go), so that DoApiRequest can populate
// http.Request.ContentLength manually (net/http only auto-detects it for
// *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide".
UpstreamRequestBodySize int64
PriceData types.PriceData
// TieredBillingSnapshot is a frozen snapshot of tiered billing rules
......
relay/compatible_handler.go
package relay
import (
"bytes"
"fmt"
"io"
"net/http"
......@@ -176,7 +175,14 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
logger.LogDebug(c, "text request body: %s", jsonData)
requestBody = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
}
var httpResp *http.Response
......
relay/embedding_handler.go
package relay
import (
"bytes"
"fmt"
"io"
"net/http"
......@@ -59,7 +58,14 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}
logger.LogDebug(c, "converted embedding request body: %s", jsonData)
var requestBody io.Reader = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
var requestBody io.Reader = body
statusCodeMappingStr := c.GetString("status_code_mapping")
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
......
relay/gemini_handler.go
package relay
import (
"bytes"
"fmt"
"io"
"net/http"
......@@ -165,7 +164,14 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
logger.LogDebug(c, "Gemini request body: %s", jsonData)
requestBody = bytes.NewReader(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
}
resp, err := adaptor.DoRequest(c, info, requestBody)
......@@ -263,7 +269,14 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
}
}
logger.LogDebug(c, "Gemini embedding request body: %s", jsonData)
requestBody = bytes.NewReader(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
......
relay/image_handler.go
......@@ -77,7 +77,14 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}
logger.LogDebug(c, "image request body: %s", jsonData)
requestBody = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
}
}
......
relay/rerank_handler.go
package relay
import (
"bytes"
"fmt"
"io"
"net/http"
......@@ -69,7 +68,14 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
logger.LogDebug(c, "Rerank request body: %s", jsonData)
requestBody = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
}
resp, err := adaptor.DoRequest(c, info, requestBody)
......
relay/responses_handler.go
package relay
import (
"bytes"
"fmt"
"io"
"net/http"
......@@ -104,7 +103,14 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}
logger.LogDebug(c, "requestBody: %s", jsonData)
requestBody = bytes.NewBuffer(jsonData)
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
requestBody = body
}
var httpResp *http.Response
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment