Repository navigation
perf(epp): skip the RawBody copy for bodies above the pool cap - #3883
Conversation
| if s.maxPoolBufferSize > 0 && buf.Cap() > s.maxPoolBufferSize { | ||
| // putBuffer drops a buffer above the pool cap, so RawBody takes its | ||
| // memory. Parsers and the Director replace RawBody on a rewrite and | ||
| // never write into it, so inboundBody stays the comparison baseline. | ||
| reqCtx.Request.RawBody = inboundBody | ||
| buf = nil |
There was a problem hiding this comment.
Merge this after #3884, and say in the description that it depends on it.
Why: RawBody now holds the buffer's whole underlying array, not just the body. bytes.Buffer roughly doubles as it grows, so that array is 1.6 to 1.9 times the body size. A 20 MB body written in 64 KiB chunks ends up in a 33.5 MiB array. On main, RawBody is held until the stream ends, so without #3884 each long streamed response keeps 33.5 MiB alive instead of a 20 MB copy. Peak memory while the request is processed goes down, but memory held for the rest of the generation goes up.
| inboundBody := buf.Bytes() | ||
| reqCtx.Request.RawBody = make([]byte, buf.Len()) | ||
| copy(reqCtx.Request.RawBody, inboundBody) | ||
| if s.maxPoolBufferSize > 0 && buf.Cap() > s.maxPoolBufferSize { |
There was a problem hiding this comment.
Move the pool-cap rule into one helper and use it in both putBuffer and this branch.
func (s *StreamingServer) poolKeeps(buf *bytes.Buffer) bool {
return s.maxPoolBufferSize == 0 || buf.Cap() <= s.maxPoolBufferSize
}Why: this branch is only safe because putBuffer never returns these buffers to the pool, and the rule is written out twice. If putBuffer changes, for example to keep larger buffers, the pool hands out a buffer that RawBody still points into. The next request overwrites it, and the earlier request forwards a corrupted body.
There was a problem hiding this comment.
Done. poolKeeps holds the rule, and putBuffer and the RawBody branch both call it.
| return reqCtx, nil | ||
| } | ||
|
|
||
| func TestProcessRawBodyUsesBufferAbovePoolCap(t *testing.T) { |
There was a problem hiding this comment.
Add a case with tool-calling telemetry enabled and a buffer above the cap.
Why: that path now passes buf == nil to the new if buf != nil guard, while the snapshot and compareEPPToolCallingSnapshotToBody read the same memory RawBody uses. No test covers this combination, so a nil dereference or a broken comparison there would not be caught.
There was a problem hiding this comment.
Done. TestProcessToolCallingSnapshotWithBufferAbovePoolCap enables telemetry with a buffer above the cap. It checks that RawBody shares the buffer memory and that the comparison reports tools as preserved.
|
@thomas-fahrner-parasail pls rebase, I think your other PR that got merged cause a conflict on this one |
|
@thomas-fahrner-parasail This PR has merge conflicts with gh pr checkout 3883
git fetch origin main
git rebase origin/main
# resolve conflicts
git push --force-with-lease |
A request buffer larger than the pool cap (the gRPC max receive size, 4 MiB by default) is dropped after the request phase, so copying it into RawBody only doubles the memory a large body takes while it is parsed and scheduled. RawBody takes the buffer's memory instead. Smaller bodies keep the copy so their buffers return to the pool. Signed-off-by: Thomas Fahrner <thomas.fahrner@parasail.io>
RawBody can take a buffer's memory only because putBuffer never returns that buffer to the pool. poolKeeps now holds that rule for both, so a change to which buffers the pool keeps cannot leave RawBody pointing into a reused buffer. A new test covers tool-calling telemetry with a buffer above the pool cap: the snapshot and RawBody share memory, and the comparison reports the tools field as preserved. Signed-off-by: Thomas Fahrner <thomas.fahrner@parasail.io>
3051db1 to
dfd33cc
Compare
|
Rebased on main. The only conflict was in |
|
@ahg-g @thomas-fahrner-parasail has rebased this PR and it is now conflict-free. Please take another look and re-approve when ready. |
What type of PR is this?
/kind performance
What this PR does / why we need it:
putBufferdrops a request buffer larger than the pool cap (--grpc-max-recv-msg-size, 4 MiB by default). The handler still copies that buffer intoRawBody, so a large body takes 2 times its size while it is parsed and scheduled.With this PR,
RawBodyuses the memory of a buffer above the pool cap, and the handler does not return that buffer to the pool. A buffer at or below the cap keeps the copy and returns to the pool. Parsers and the Director replaceRawBodywhen they rewrite the body and do not write into it, so the inbound bytes stay valid for the tool-calling snapshot comparison.For a 20 MB body, this removes 1 copy of 20 MB for the duration of the request phase.
Test plan:
RawBodyshares its memoryRawBodyis a copytoolsas preservedmake presubmitWhich issue(s) this PR fixes:
Part of #3881
Depends on #3884: merge this PR after #3884.
RawBodyholds the buffer's whole backing array, which is up to 2 times the body size (32 MiB for a 20 MB body). Without #3884,RawBodystays live until the stream ends, so a long streamed response would hold more memory than the copy it replaces.Release note (write
NONEif no user-facing change):