fix(openai-ws): share immutable replay bodies to stop O(T²) copying

OpenAI Responses WS/HTTP-bridge 的 replay 历史此前每轮对全部正文做多次
深拷贝(构建、双份序列、保存点各一遍),长会话累计分配 O(T²),实测
128 turn x 10KiB 达 160MB+,是 #6382 内存持续增长直至 OOM 的主因。

改为共享不可变正文所有权:
- extract 经零拷贝 string 视图 + gjson Index 直接切 payload 子切片,
  并以 json.Valid 保持对断裂数组的既有错误语义
- 序列构建/sanitize/collector/保存点全部只传引用;跨持有者保存经
  combineOpenAIWSReplayItems 新建头数组,正文不复制
- bridge 每轮一次解析,正常 replay 与 account-failover 序列共享正文
- 前缀比较增加字节级快路径,避免逐轮全历史 decode+encode
- strict 状态可用时不再整包保留上一轮 payload

新增 128 turn 分配回归测试锚定上界(实测 2.5MB,深拷贝回归会击穿),
以及正文共享的所有权断言。
This commit is contained in:
feeeei
2026-08-31 01:01:20 +08:00
parent b5827cfd54
commit ca9b4d73f2
5 changed files with 270 additions and 71 deletions
@@ -586,24 +586,27 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
toolOutputCoverage := AnalyzeToolCallOutputContextCoverageBytes(currentBridgePayload.payloadRaw)
needsBridgeReplay := currentBridgePayload.previousResponseID != "" ||
(toolOutputCoverage.HasFunctionCallOutput && !toolOutputCoverage.ContextCoversAllCallIDs)
turnReplayInput, turnReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence(
// 一次解析当前 input,正常 replay 与 account-failover 两份序列共享同一批正文。
bridgeCurrentItems, bridgeCurrentItemsExist, extractErr := openAIWSExtractNormalizedInputSequence(
currentBridgePayload.payloadRaw,
)
if extractErr != nil {
return fmt.Errorf("build websocket http bridge replay input: %w", extractErr)
}
turnReplayInput, turnReplayInputExists := buildOpenAIWSReplayInputSequenceFromItems(
bridgeReplayInput,
bridgeReplayInputExists,
currentBridgePayload.payloadRaw,
bridgeCurrentItems,
bridgeCurrentItemsExist,
needsBridgeReplay,
)
if replayInputErr != nil {
return fmt.Errorf("build websocket http bridge replay input: %w", replayInputErr)
}
turnAccountFailoverInput, turnAccountFailoverInputExists, failoverInputErr := buildOpenAIWSReplayInputSequence(
turnAccountFailoverInput, turnAccountFailoverInputExists := buildOpenAIWSReplayInputSequenceFromItems(
bridgeAccountFailoverInput,
bridgeAccountFailoverInputExists,
currentBridgePayload.payloadRaw,
bridgeCurrentItems,
bridgeCurrentItemsExist,
needsBridgeReplay,
)
if failoverInputErr != nil {
return fmt.Errorf("build websocket account failover input: %w", failoverInputErr)
}
if needsBridgeReplay && turnReplayInputExists {
updatedPayload, setInputErr := setOpenAIWSPayloadInputSequence(
currentBridgePayload.payloadRaw,
@@ -680,18 +683,20 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
if result == nil {
return errors.New("websocket http bridge turn result is nil")
}
bridgeReplayInput = cloneOpenAIWSRawMessages(turnReplayInput)
// turnReplayInput/turnAccountFailoverInput 可能共享同一头数组(转移自
// bridgeCurrentItems),保存历史必须经 combine 新建头,禁止就地 append。
bridgeReplayInput = turnReplayInput
bridgeReplayInputExists = turnReplayInputExists
if result.wsReplayInputExists {
bridgeReplayInput = append(bridgeReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...)
bridgeReplayInput = combineOpenAIWSReplayItems(bridgeReplayInput, result.wsReplayInput)
bridgeReplayInputExists = true
}
bridgeAccountFailoverInput = cloneOpenAIWSRawMessages(turnAccountFailoverInput)
bridgeAccountFailoverInput = turnAccountFailoverInput
bridgeAccountFailoverInputExists = turnAccountFailoverInputExists
if len(result.wsAccountFailoverReplayInput) > 0 {
bridgeAccountFailoverInput = append(
bridgeAccountFailoverInput = combineOpenAIWSReplayItems(
bridgeAccountFailoverInput,
cloneOpenAIWSRawMessages(result.wsAccountFailoverReplayInput)...,
result.wsAccountFailoverReplayInput,
)
bridgeAccountFailoverInputExists = true
}
@@ -1744,16 +1749,19 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
}
responseID := strings.TrimSpace(result.RequestID)
lastTurnResponseID = responseID
lastTurnPayload = cloneOpenAIWSPayloadBytes(currentPayload)
lastTurnReplayInput = cloneOpenAIWSRawMessages(currentTurnReplayInput)
// 正文共享:currentPayload/currentTurnReplayInput 均不可变,历史直接引用;
// collector 增量经 combine 合并(新头数组)。
lastTurnReplayInput = currentTurnReplayInput
lastTurnReplayInputExists = currentTurnReplayInputExists
if result.wsReplayInputExists {
lastTurnReplayInput = append(lastTurnReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...)
lastTurnReplayInput = combineOpenAIWSReplayItems(lastTurnReplayInput, result.wsReplayInput)
lastTurnReplayInputExists = true
}
nextStrictState, strictStateErr := buildOpenAIWSIngressPreviousTurnStrictState(currentPayload)
if strictStateErr != nil {
lastTurnStrictState = nil
// strict 状态不可用时保留整份上一轮 payload 供慢路径比较。
lastTurnPayload = currentPayload
logOpenAIWSModeInfo(
"ingress_ws_prev_response_strict_state_skip account_id=%d turn=%d conn_id=%s reason=build_error cause=%s",
account.ID,
@@ -1763,6 +1771,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
)
} else {
lastTurnStrictState = nextStrictState
lastTurnPayload = nil
}
if responseID != "" && stateStore != nil {
@@ -1,6 +1,7 @@
package service
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -1007,18 +1008,79 @@ func TestSetOpenAIWSPayloadInputSequence(t *testing.T) {
})
}
func TestCloneOpenAIWSRawMessages(t *testing.T) {
func TestCombineOpenAIWSReplayItems(t *testing.T) {
t.Parallel()
t.Run("nil_slice", func(t *testing.T) {
cloned := cloneOpenAIWSRawMessages(nil)
require.Nil(t, cloned)
t.Run("empty_delta_returns_history", func(t *testing.T) {
history := []json.RawMessage{json.RawMessage(`{"a":1}`)}
require.Nil(t, combineOpenAIWSReplayItems(nil, nil))
combined := combineOpenAIWSReplayItems(history, nil)
require.Len(t, combined, 1)
})
t.Run("empty_slice", func(t *testing.T) {
items := make([]json.RawMessage, 0)
cloned := cloneOpenAIWSRawMessages(items)
require.NotNil(t, cloned)
require.Len(t, cloned, 0)
t.Run("new_header_shares_bodies", func(t *testing.T) {
history := []json.RawMessage{json.RawMessage(`{"a":1}`)}
delta := []json.RawMessage{json.RawMessage(`{"b":2}`)}
combined := combineOpenAIWSReplayItems(history, delta)
require.Len(t, combined, 2)
// 头数组必须是新建的:对 combined 追加不影响 history。
require.NotSame(t, &history[0], &combined[0])
// 正文共享:不发生字节级深拷贝。
require.Same(t, &history[0][0], &combined[0][0])
require.Same(t, &delta[0][0], &combined[1][0])
})
}
func TestOpenAIWSReplaySequenceSharesBodies(t *testing.T) {
t.Parallel()
t.Run("extract_shares_payload_backing_array", func(t *testing.T) {
payload := []byte(`{"input":[{"type":"input_text","text":"hello"},{"type":"input_text","text":"world"}]}`)
items, exists, err := openAIWSExtractNormalizedInputSequence(payload)
require.NoError(t, err)
require.True(t, exists)
require.Len(t, items, 2)
for _, item := range items {
start := bytes.Index(payload, []byte(item))
require.GreaterOrEqual(t, start, 0)
require.Same(t, &payload[start], &item[0], "extract 应零拷贝共享 payload 底层数组")
}
})
t.Run("build_transfers_current_items_ownership", func(t *testing.T) {
payload := []byte(`{"input":[{"type":"input_text","text":"hello"}]}`)
items, exists, err := buildOpenAIWSReplayInputSequence(nil, false, payload, false)
require.NoError(t, err)
require.True(t, exists)
require.Len(t, items, 1)
start := bytes.Index(payload, []byte(items[0]))
require.GreaterOrEqual(t, start, 0)
require.Same(t, &payload[start], &items[0][0])
})
t.Run("build_merge_shares_history_bodies", func(t *testing.T) {
history := []json.RawMessage{json.RawMessage(`{"type":"input_text","text":"hello"}`)}
items, exists, err := buildOpenAIWSReplayInputSequence(
history,
true,
[]byte(`{"previous_response_id":"resp_1","input":[{"type":"input_text","text":"world"}]}`),
true,
)
require.NoError(t, err)
require.True(t, exists)
require.Len(t, items, 2)
require.Same(t, &history[0][0], &items[0][0], "历史正文应共享而非深拷贝")
})
t.Run("build_prefix_hit_transfers_current_items", func(t *testing.T) {
history := []json.RawMessage{json.RawMessage(`{"type":"input_text","text":"hello"}`)}
payload := []byte(`{"previous_response_id":"resp_1","input":[{"type":"input_text","text":"hello"},{"type":"input_text","text":"world"}]}`)
items, exists, err := buildOpenAIWSReplayInputSequence(history, true, payload, true)
require.NoError(t, err)
require.True(t, exists)
require.Len(t, items, 2)
start := bytes.Index(payload, []byte(items[1]))
require.GreaterOrEqual(t, start, 0)
require.Same(t, &payload[start], &items[1][0], "prefix 命中应转移当前 items 所有权并共享 payload 底层数组")
})
}
@@ -9,6 +9,7 @@ import (
"net/http"
"net/url"
"strings"
"unsafe"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
@@ -431,24 +432,37 @@ func alignStoreDisabledPreviousResponseID(
return updated, true, nil
}
func cloneOpenAIWSPayloadBytes(payload []byte) []byte {
if len(payload) == 0 {
return nil
// Replay 状态所有权不变式:replay 序列中的 json.RawMessage 正文一经放入即视为
// 不可变,所有持有者共享同一份字节,任何修改都必须整体替换元素或重建 payload。
// 序列头数组在跨持有者保存时必须新建(combineOpenAIWSReplayItems),禁止通过
// 共享头 append,否则会写入其他持有者可见的底层数组。
// combineOpenAIWSReplayItems 合并历史与增量为新头数组,正文共享不复制。
func combineOpenAIWSReplayItems(history, delta []json.RawMessage) []json.RawMessage {
if len(delta) == 0 {
return history
}
cloned := make([]byte, len(payload))
copy(cloned, payload)
return cloned
combined := make([]json.RawMessage, 0, len(history)+len(delta))
combined = append(combined, history...)
return append(combined, delta...)
}
func cloneOpenAIWSRawMessages(items []json.RawMessage) []json.RawMessage {
if items == nil {
return nil
// openAIWSPayloadStringView 返回与 payload 共享底层数组的零拷贝 string 视图,
// 供 gjson.Get 使用(gjson.GetBytes 会整段复制结果 Raw,对 input 这类占
// payload 主体的字段是每次 O(payload) 分配)。调用方必须保证 payload 在结果
// 存活期间不可变(replay 所有权不变式)。
func openAIWSPayloadStringView(payload []byte) string {
return unsafe.String(unsafe.SliceData(payload), len(payload))
}
// openAIWSRawMessageFromResult 优先返回 parent 的子切片(gjson 值零拷贝共享),
// Index 不可用时回退为复制。共享要求 parent 遵守上面的不可变约定。
func openAIWSRawMessageFromResult(parent []byte, value gjson.Result) json.RawMessage {
idx := value.Index
if idx > 0 && idx+len(value.Raw) <= len(parent) && string(parent[idx:idx+len(value.Raw)]) == value.Raw {
return json.RawMessage(parent[idx : idx+len(value.Raw)])
}
cloned := make([]json.RawMessage, 0, len(items))
for idx := range items {
cloned = append(cloned, json.RawMessage(cloneOpenAIWSPayloadBytes(items[idx])))
}
return cloned
return json.RawMessage(value.Raw)
}
func normalizeOpenAIWSJSONForCompare(raw []byte) ([]byte, error) {
@@ -495,30 +509,38 @@ func normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(payload []byte) (
return json.Marshal(decoded)
}
// openAIWSExtractNormalizedInputSequence 拆出 input 序列。返回的正文尽可能与
// payload 共享底层数组(零拷贝),受 replay 所有权不变式保护。
func openAIWSExtractNormalizedInputSequence(payload []byte) ([]json.RawMessage, bool, error) {
if len(payload) == 0 {
return nil, false, nil
}
inputValue := gjson.GetBytes(payload, "input")
inputValue := gjson.Get(openAIWSPayloadStringView(payload), "input")
if !inputValue.Exists() {
return nil, false, nil
}
if inputValue.Type == gjson.JSON {
raw := strings.TrimSpace(inputValue.Raw)
if strings.HasPrefix(raw, "[") {
var items []json.RawMessage
if err := json.Unmarshal([]byte(raw), &items); err != nil {
return nil, true, err
if inputValue.IsArray() {
// gjson 宽容解析;数组整体先做零分配合法性校验,避免把断裂
// JSON 塞进 replay 历史。
arrayRaw := openAIWSRawMessageFromResult(payload, inputValue)
if !json.Valid(arrayRaw) {
return nil, true, errors.New("input array json is invalid")
}
elems := inputValue.Array()
items := make([]json.RawMessage, 0, len(elems))
for _, elem := range elems {
items = append(items, openAIWSRawMessageFromResult(payload, elem))
}
return items, true, nil
}
return []json.RawMessage{json.RawMessage(raw)}, true, nil
return []json.RawMessage{openAIWSRawMessageFromResult(payload, inputValue)}, true, nil
}
if inputValue.Type == gjson.String {
encoded, _ := json.Marshal(inputValue.String())
return []json.RawMessage{encoded}, true, nil
}
return []json.RawMessage{json.RawMessage(inputValue.Raw)}, true, nil
return []json.RawMessage{openAIWSRawMessageFromResult(payload, inputValue)}, true, nil
}
func openAIWSInputIsPrefixExtended(previousPayload, currentPayload []byte) (bool, error) {
@@ -561,6 +583,10 @@ func openAIWSRawItemsHasPrefix(items []json.RawMessage, prefix []json.RawMessage
return false
}
for idx := range prefix {
// 快路径:客户端逐字节重发历史时直接比较,避免整轮历史的解码/再编码。
if bytes.Equal(bytes.TrimSpace(prefix[idx]), bytes.TrimSpace(items[idx])) {
continue
}
previousNormalized := normalizeOpenAIWSJSONForCompareOrRaw(prefix[idx])
currentNormalized := normalizeOpenAIWSJSONForCompareOrRaw(items[idx])
if !bytes.Equal(previousNormalized, currentNormalized) {
@@ -611,12 +637,13 @@ func openAIWSRawItemsHaveToolCallContextForOutputs(items []json.RawMessage) bool
return true
}
// sanitizeOpenAIWSHistoricalReplayToolCalls 返回的新头数组与 previousItems 共享正文。
func sanitizeOpenAIWSHistoricalReplayToolCalls(
previousItems []json.RawMessage,
currentItems []json.RawMessage,
) []json.RawMessage {
if len(previousItems) == 0 {
return cloneOpenAIWSRawMessages(previousItems)
return previousItems
}
outputCallIDs := make(map[string]struct{})
collectOutputCallIDs := func(items []json.RawMessage) {
@@ -640,7 +667,7 @@ func sanitizeOpenAIWSHistoricalReplayToolCalls(
continue
}
}
sanitized = append(sanitized, append(json.RawMessage(nil), item...))
sanitized = append(sanitized, item)
}
return sanitized
}
@@ -649,7 +676,7 @@ func openAIWSRawPayloadHasToolCallOutput(payload []byte) bool {
if len(payload) == 0 {
return false
}
input := gjson.GetBytes(payload, "input")
input := gjson.Get(openAIWSPayloadStringView(payload), "input")
if !input.Exists() {
return false
}
@@ -667,6 +694,33 @@ func openAIWSRawPayloadHasToolCallOutput(payload []byte) bool {
return false
}
// buildOpenAIWSReplayInputSequenceFromItems 基于已解析的当前 turn input 构建
// replay 序列。返回序列的正文与 previousFullInput/currentItems 共享所有权
// (见 combineOpenAIWSReplayItems 上方的所有权不变式),头数组可能直接转移自
// currentItems。
func buildOpenAIWSReplayInputSequenceFromItems(
previousFullInput []json.RawMessage,
previousFullInputExists bool,
currentItems []json.RawMessage,
currentExists bool,
hasPreviousResponseID bool,
) ([]json.RawMessage, bool) {
if !hasPreviousResponseID || !previousFullInputExists {
return currentItems, currentExists
}
previousFullInput = sanitizeOpenAIWSHistoricalReplayToolCalls(previousFullInput, currentItems)
if !currentExists || len(currentItems) == 0 {
return previousFullInput, true
}
if openAIWSRawItemsHasPrefix(currentItems, previousFullInput) {
return currentItems, true
}
merged := make([]json.RawMessage, 0, len(previousFullInput)+len(currentItems))
merged = append(merged, previousFullInput...)
merged = append(merged, currentItems...)
return merged, true
}
func buildOpenAIWSReplayInputSequence(
previousFullInput []json.RawMessage,
previousFullInputExists bool,
@@ -677,23 +731,14 @@ func buildOpenAIWSReplayInputSequence(
if currentErr != nil {
return nil, false, currentErr
}
if !hasPreviousResponseID {
return cloneOpenAIWSRawMessages(currentItems), currentExists, nil
}
if !previousFullInputExists {
return cloneOpenAIWSRawMessages(currentItems), currentExists, nil
}
previousFullInput = sanitizeOpenAIWSHistoricalReplayToolCalls(previousFullInput, currentItems)
if !currentExists || len(currentItems) == 0 {
return cloneOpenAIWSRawMessages(previousFullInput), true, nil
}
if openAIWSRawItemsHasPrefix(currentItems, previousFullInput) {
return cloneOpenAIWSRawMessages(currentItems), true, nil
}
merged := make([]json.RawMessage, 0, len(previousFullInput)+len(currentItems))
merged = append(merged, cloneOpenAIWSRawMessages(previousFullInput)...)
merged = append(merged, cloneOpenAIWSRawMessages(currentItems)...)
return merged, true, nil
items, exists := buildOpenAIWSReplayInputSequenceFromItems(
previousFullInput,
previousFullInputExists,
currentItems,
currentExists,
hasPreviousResponseID,
)
return items, exists, nil
}
func setOpenAIWSPayloadInputSequence(
@@ -8,6 +8,7 @@ import (
"fmt"
"io"
"net/http"
"slices"
"strings"
"time"
@@ -156,12 +157,14 @@ func (c *openAIWSToolCallReplayCollector) AddEvent(eventType string, message []b
}
}
// Items/AllItems 返回浅拷贝头数组;正文由 collector 独立分配且此后不可变,
// 调用方按 replay 所有权不变式共享持有。
func (c *openAIWSToolCallReplayCollector) Items() []json.RawMessage {
return cloneOpenAIWSRawMessages(c.items)
return slices.Clone(c.items)
}
func (c *openAIWSToolCallReplayCollector) AllItems() []json.RawMessage {
return cloneOpenAIWSRawMessages(c.allItems)
return slices.Clone(c.allItems)
}
func (c *openAIWSToolCallReplayCollector) addAllItem(item gjson.Result) {
@@ -0,0 +1,80 @@
package service
import (
"encoding/json"
"fmt"
"runtime"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
// buildReplayTurnPayload 构造第 turn 轮的全量重发 payload:客户端历史逐轮追加
// itemBytes 大小的项(Codex store=false 模式的真实形态)。
func buildReplayTurnPayload(turn, itemBytes int) []byte {
var b strings.Builder
_, _ = b.WriteString(`{"type":"response.create","model":"gpt-5.5","stream":true,"input":[`)
filler := strings.Repeat("x", itemBytes)
for i := 1; i <= turn; i++ {
if i > 1 {
_, _ = b.WriteString(",")
}
_, _ = fmt.Fprintf(&b, `{"type":"input_text","text":"turn-%d-%s"}`, i, filler)
}
_, _ = b.WriteString(`]}`)
return []byte(b.String())
}
// TestOpenAIWSReplayStateBuildAllocationBounded 长会话分配回归:128 turn、每轮
// ~10KiB 增量,全量历史逐轮重发。replay 状态构建(build + 历史保存)必须共享
// 正文而非逐轮深拷贝,否则累计分配为 O(T²)(本场景 >160MB);共享实现的累计
// 分配与客户端重发总量同阶(~85MB payload 本身之外仅头数组与解析开销)。
func TestOpenAIWSReplayStateBuildAllocationBounded(t *testing.T) {
const (
turns = 128
itemBytes = 10 * 1024
)
payloads := make([][]byte, 0, turns)
for turn := 1; turn <= turns; turn++ {
payloads = append(payloads, buildReplayTurnPayload(turn, itemBytes))
}
var history []json.RawMessage
historyExists := false
delta := []json.RawMessage{json.RawMessage(`{"type":"function_call","id":"item_1","call_id":"call_1","name":"exec","arguments":"{}"}`)}
runtime.GC()
var before runtime.MemStats
runtime.ReadMemStats(&before)
for turn := 1; turn <= turns; turn++ {
items, exists, err := buildOpenAIWSReplayInputSequence(history, historyExists, payloads[turn-1], turn > 1)
require.NoError(t, err)
require.True(t, exists)
require.Len(t, items, turn)
// 保存历史 + collector 增量合并(与 ingress/bridge 保存点同构)。
history = combineOpenAIWSReplayItems(items, delta)
historyExists = true
// 下一轮 payload 含全部用户项但不含 collector 项,触发 sanitize+merge 路径中
// 最常见的 prefix 分支比较。
history = history[:len(history)-1]
}
var after runtime.MemStats
runtime.ReadMemStats(&after)
allocated := after.TotalAlloc - before.TotalAlloc
// 正文共享实现全程约 2.5MB(头数组与 gjson 解析开销)。任何逐轮 O(H) 的
// 字节级复制(extract 拷贝、正文 clone、保存点深拷贝)都会叠加 ≥84MB
// (sum(t×10KiB))的量级,直接击穿此上界。
const maxAllocatedBytes = 32 * 1024 * 1024
require.Lessf(
t,
allocated,
uint64(maxAllocatedBytes),
"replay 状态构建累计分配 %d bytes 超出上界,疑似回归为逐轮深拷贝",
allocated,
)
}