proxy: replace streamBuffer with chunkedResponseWriter for capped response storage#88
Open
elevran wants to merge 6 commits into
Open
proxy: replace streamBuffer with chunkedResponseWriter for capped response storage#88elevran wants to merge 6 commits into
elevran wants to merge 6 commits into
Conversation
…ponse storage
Replace the all-at-once Store call with a staging-based chunked protocol:
- POST /staging → PUT /staging/{sid}/chunks/{k} (×N) → PUT /staging/{sid}/complete
- Abort via PUT /staging/{sid}/abort when no bytes were written
Key changes:
- pkg/charon/client.go: extend Backend interface with AppendChunk, Complete, Abort
- cmd/proxy/chunk.go: new chunkedResponseWriter — byte-level splitter with
configurable MaxChunkBytes cap; splits within a single Add call
- internal/proxyconfig/proxyconfig.go: add --max-chunk-bytes flag (default 1 MiB)
- cmd/proxy/handler.go: wire chunkedResponseWriter into HandleCreate; add
WithMaxChunkBytes option
- cmd/proxy/sse.go: remove streamBuffer; accumulate output items, marshal once,
write via chunkedResponseWriter
- cmd/proxy/ws.go: same pattern as sse.go
- cmd/proxy/main.go: pass MaxChunkBytes from config to handler
Tests: unit tests in chunk_test.go; integration tests in handler_test.go,
backend_routing_test.go, and disruptive_test.go verify chunk counts, cap
splitting, store:false no-op, and mid-stream failure non-persistence.
Apply consensus findings from /simplify and /ponytail-review: - chunk.go: drop errWriterClosed sentinel. No caller or test compares against it; closedErr now stores backend.Abort's error directly. Removes the 'errors' import. - chunk.go (Abort): replace the errWriterClosed default + override with a single closedErr = err assignment. - backend_routing_test.go: rewrite the docstring above TestBufferedProxyStoreFalseNoStagingCalls to match what the test actually asserts (no staging calls on store:false). The prior comment named a non-existent test and rambled about aborts. - chunk_test.go (TestWriterCapFlush): fix 'chunk indices must be 0,1,2' comment — the loop iterates only 0,1 for the 2-element chunks slice.
The chunk tests added in this PR (TestBufferedProxySingleChunk,
TestStreamedProxySingleChunk, TestBufferedProxyMultipleChunks,
TestBufferedProxyStoreFalseNoStagingCalls, TestBufferedCapConfigurable)
all manually constructed 'rec := &routingRecorder{}; s := newTestStack(t,
withCharonMiddleware(rec.middleware()))' — the exact pattern that
newRoutingStack(t) already encapsulates.
Extend newRoutingStack to accept variadic stackOption so the cap test
can pass withMaxChunkBytes(64) alongside the recorder middleware.
Net: -5 lines per test, single source of truth for the recorder setup.
- New internal/bytesize package: ByteSize type with UnmarshalJSON, plus KiB/MiB/GiB/TiB constants. The constants are exported so call sites can write 'bytesize.MiB' instead of '1 << 20'. - charonconfig: drop ByteSize; reference bytesize.ByteSize in CharonOptions and fileStorageConfig. JSON parsing is preserved (the unmarshal method moved with the type). - internal/server/handlers.go: defaultMaxChunkBytes = bytesize.MiB, maxChunkBytes = 4 * bytesize.MiB. - cmd/proxy/handler.go: NewHandler's default cap is bytesize.MiB. - internal/proxyconfig/proxyconfig.go: defaultMaxChunkBytes const is bytesize.MiB.
The code-reviewer agent found a real concurrency bug in chunkedResponseWriter.Close: the second concurrent Close caller read w.closedErr under the lock BEFORE waiting on <-ch, racing against the first Close's deferred write of closedErr (which happens after the unlocked Complete network call). On failure, the second caller could return a stale nil, masking the real error. Reproduction (reviewer's 200-iter trial without the fix): ~3/200 failures under go test -race. Fix: in the early-return path, drop the lock, wait on <-ch, then re-acquire the lock to read closedErr. The channel close provides happens-before for closedErr, and the second lock acquisition prevents torn interface reads. Symmetric with Abort, which already reads closedErr after <-ch. Also add TestWriterCloseIdempotentOnError that exercises the failure path concurrently to lock in the contract. Pin both paths at 200 iter under -race. Drive-by: fix the misleading comment in TestBufferedProxyStoreFalseNoStagingCalls — the prior wording pointed at disruptive_test.go for the empty-response abort path, but disruptive_test.go only covers inference-truncation. The unit tests in chunk_test.go (TestWriterZeroBytesAborts) actually exercise that path.
Apply every Important and Minor finding from the code-reviewer agent: 1. Symmetric Abort + log Abort errors (handler.go/sse.go/ws.go): Extract (h *Handler) commitStoredResponse(ctx, stagingID, id, tenantKey, blob) error. The helper owns Add → Abort-on-error → Close → Abort-on-error, logging all stages. Abort failures are logged at warn but never mask the original error. Each call site collapses from ~12 lines to a single if that maps the error to its surface (HTTP / SSE / WS). This addresses the deferred altitude finding (three-way duplication) and the symmetry finding in one shot. 2. chunk.flush(): make flush failure terminal. When AppendChunk returns an error, the writer now sets closed=true and closes closedCh (gated by !w.closed so a Close-driven flush doesn't double-close the channel). Subsequent Add/Close/Abort calls return the same error without retrying the same chunk index. Removed the redundant closedErr = err in Close since flush now owns that write. 3. internal/server/handlers.go: rename defaultMaxChunkBytes to defaultChunkBodyBytes — disambiguates from proxyconfig's defaultMaxChunkBytes (different meaning, same value). Also remove the unused maxChunkBytes = 4*MiB const. 4. internal/proxyconfig/proxyconfig.go: add maxMaxChunkBytes = MiB cap and Validate() check that rejects --max-chunk-bytes values exceeding it with a clear error. Update --max-chunk-bytes help text to advertise the upper bound. Hardcoded rather than imported from internal/server to avoid a config→server import edge; comment cross-references server.defaultChunkBodyBytes. Verified: presubmit green; go test -race -count=200 -run TestWriter ./cmd/proxy/ passes (covers the new terminal-flush state path).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Storecall with a staging-based chunked protocol:POST /staging→PUT /staging/{sid}/chunks/{k}(×N) →PUT /staging/{sid}/complete, orPUT /staging/{sid}/aborton empty outputchunkedResponseWriterincmd/proxy/chunk.gosplits response bytes at a configurableMaxChunkBytescap (default 1 MiB), including within a single largeAddcallstreamBufferis removedChanges
pkg/charon/client.go: extendBackendinterface withAppendChunk,Complete,Abortcmd/proxy/chunk.go:chunkedResponseWriter— byte-level splitter, mutex-protected, idempotentClose/Abortinternal/proxyconfig/proxyconfig.go:--max-chunk-bytesflag (0 → default 1 MiB)cmd/proxy/handler.go: wire writer intoHandleCreate;WithMaxChunkBytesoptioncmd/proxy/sse.go: removestreamBuffer; marshal full blob at end, write via writercmd/proxy/ws.go: same pattern assse.gocmd/proxy/main.go: passMaxChunkBytesfrom config to handlerTest Plan
chunk_test.go: single chunk, cap flush, abort-on-empty, idempotent close, cancel during flush, flush failure propagationTestBufferedCapConfigurable/TestBufferedProxyMultipleChunks: tiny cap forces ≥2AppendChunkcallsTestBufferedProxySingleChunk/TestStreamedProxySingleChunk: default cap → exactly 1 chunk + 1 completeTestBufferedProxyStoreFalseNoStagingCalls:store:falsemakes no staging callsTestStreamingInferenceFailureEmitsFailedNotCompleted: truncated stream → 0 chunks, 0 completego test ./... -raceandmake presubmitclean