feat: Apply the init-concurrency limiter to polling and streaming - #782
Draft
kinyoklion wants to merge 14 commits into
Draft
feat: Apply the init-concurrency limiter to polling and streaming#782kinyoklion wants to merge 14 commits into
kinyoklion wants to merge 14 commits into
Conversation
kinyoklion
force-pushed
the
rlamb/relay-init-concurrency-wiring
branch
from
July 29, 2026 20:57
141e39f to
6349165
Compare
keelerm84
reviewed
Jul 29, 2026
| // Retry-After and does not invoke the wrapped handler. A disabled or nil limiter is a | ||
| // pass-through with zero overhead. | ||
| func LimitConcurrency(limiter *concurrency.Limiter) func(http.Handler) http.Handler { | ||
| return func(next http.Handler) http.Handler { |
Member
There was a problem hiding this comment.
We should add some metrics and spans for this.
Member
Author
There was a problem hiding this comment.
Yes, I am currently deferring the otel instrumentation until we have the ServerTrace implementation in the eventsource. Then I will make a followup PR.
Member
Author
There was a problem hiding this comment.
Technically we only need it for half of the changes, but I think it is nice to put it in one batch.
| // the egress of sending it. The slot is held across the send and released when the | ||
| // send finishes, the client disconnects, or the stall backstop fires. | ||
| if r.initLimiter.Enabled() { | ||
| release, ok := r.initLimiter.Acquire(ctx, r.envKey) |
Member
There was a problem hiding this comment.
Here too if we could manage it.
kinyoklion
marked this pull request as ready for review
July 29, 2026 21:31
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Builds one shared init limiter from the [Concurrency] config and applies it to both initialization-delivery surfaces: - Polling: a middleware acquires a slot before the poll/eval handler runs, bounding the concurrent full-response serialization; it responds 503 + Retry-After when the budget is full. This covers the server FDv2 /sdk/poll, the client-side FDv2 /sdk/poll/eval, and the FDv1 PHP all-flags /sdk/flags. The single-flag and single-segment PHP endpoints return one item and are requested per evaluation, so they are left ungated. - Streaming: builds on the disconnect-aware replay, reading the store once (deduplicated) and sending a small up-to-date reply without a slot when the client's basis already matches. For a full basis it acquires a slot BEFORE serializing, so the budget bounds the payload memory and the egress, holds the slot across the send, and releases it when the send finishes, the client disconnects, or the stall backstop fires. The store read is single-flight-deduplicated with a fixed key so a herd of reconnects at any basis shares one read; the per-basis serialization stays deduplicated. The limiter is disabled unless configured, so behavior is unchanged by default.
kinyoklion
force-pushed
the
rlamb/relay-init-concurrency-wiring
branch
from
July 30, 2026 16:42
6349165 to
fc42e18
Compare
kinyoklion
commented
Jul 30, 2026
| // of queueing them. | ||
| MaxQueued ct.OptInt `conf:"INIT_MAX_QUEUED"` | ||
|
|
||
| // PerEnvMaxPercent limits the share of the budget that any single environment may |
Member
Author
There was a problem hiding this comment.
I've decided not to support this for now, so removing it from the config as well.
| } else { | ||
| events = r.serializePutV1(snapshot) | ||
| } | ||
| r.sendEvents(ctx, out, events) |
There was a problem hiding this comment.
Stale basis after queue wait
Medium Severity
The store is read in peek before admission, then a full-basis replay may block in Acquire for a long time when MaxQueued is non-zero. Serialization still uses that earlier snapshot, so a client that waited through store updates can be initialized from outdated data and miss changes published meanwhile.
Reviewed by Cursor Bugbot for commit fc42e18. Configure here.
Resolves correctness findings from a review of the streaming and polling integration, centered on the fact that closing a replay's batch channel is a benign end-of-batch to eventsource, not a connection close. Streaming: - Remove the idle timeout from sendEvents. Abandoning a partly-sent basis while the connection stays open let a later delta complete it into a corrupt data set, and the timeout was armed even when the limiter was disabled (a default deployment). The send loop is now base behavior: it ends only on completion or a real disconnect (context). - Bound a client that stalls without disconnecting by arming a write deadline on each write (withInitDeadline). A write that blocks past sendTimeout fails, so eventsource closes the connection, which frees the held slot and prompts a clean reconnect. This is installed only when the limiter is enabled, so a disabled deployment is exactly base behavior. - On shed, close the connection (via a context-installed hook) so the SDK reconnects with backoff instead of sitting connected but uninitialized; log at Warn. - Materialize the FDv1 put payload inside the single flight, while the slot is held, so its memory is accounted under the budget and shared rather than re-serialized per connection after the slot is released. Polling: - Acquire read and write deadlines around a held slot so a slow request body or a non-reading client cannot park it. - Acquire lazily in the FDv2 poll handlers, only on the full-basis branch, so a cheap up-to-date reply is never charged (matching the streaming path). The FDv1 all-flags poll, which is always a full delivery, keeps acquire-on-entry. - Shed with a jittered Retry-After and a JSON body. Config: an unset or zero INIT_SEND_TIMEOUT now falls back to the default for both paths; document the [Concurrency] section. Tests for the limiter middleware and the shed-closes-connection behavior.
The init-delivery limiter sets read/write deadlines on the connection via http.NewResponseController to bound how long a slow client can hold a budget slot. That controller can only reach the underlying connection if every ResponseWriter wrapper in the chain implements Unwrap. statusRecorder (duration metrics) and loggingHTTPResponseWriter (request logging) did not, so SetReadDeadline/SetWriteDeadline silently returned ErrNotSupported and the deadline protection was a no-op on every gated poll and stream. Add Unwrap to both wrappers, with tests that assert a deadline set through each reaches the base writer. Found by re-running the stalled-client repro against the built binary: with a short INIT_SEND_TIMEOUT the relay now closes a stalled stream connection at the deadline (freeing the slot); before the fix it delivered the full payload and kept the connection open regardless of the timeout.
Replaces the flat per-write deadline, which had two contradictory and both-wrong semantics (per round-2 review R2-3): on the FDv1 put and poll responses it was an absolute cap on one large write, so a healthy client slower than payload/timeout was cut mid-delivery and reconnect-looped; on FDv2 streams it re-armed per event, imposing no real bound and letting a slow-drip client hold a slot indefinitely. New internal/initwrite uses a throughput floor plus an absolute cap: a throughput floor (64 KB/s) drives a per-chunk write deadline so a slow-but-steady client is never cut for being slow, and a generous absolute cap (INIT_SEND_TIMEOUT, now defaulting to 2m) backstops a client stuck right at the floor on a very large payload. Large writes are sliced only to re-arm the deadline; nothing is buffered. The floor bites a large single write (FDv1 put, poll body); a many-small-event message (FDv2 basis) is bounded by the absolute cap. - Streams: withInitDeadline wraps the response in initwrite.Writer instead of the flat deadlineWriter. - Polls: LimitConcurrency and ProvideInitLimiter wrap the response the same way; the read deadline is dropped (R2-4: it guarded nothing and could cancel a healthy request), and ProvideInitLimiter short-circuits when disabled (R2-11). - Retry-After jitter now uses math/rand/v2 (R2-11). - Docs updated for the new semantics and to stop overclaiming coverage (R2-12). Verified: unit tests for the arming/capping/Unwrap logic; a 100 KB/s reader (above the floor) now completes a full 3 MB delivery instead of being cut at 30s; the distinct-basis memory-slope repro stays bounded.
| return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| ctx := context.WithValue(r.Context(), initLimiterCtxKey{}, initLimiterHolder{limiter: limiter, maxHold: maxHold}) | ||
| next.ServeHTTP(initwrite.Wrap(w, maxHold), r.WithContext(ctx)) | ||
| }) |
There was a problem hiding this comment.
Write deadline not cleared
Medium Severity
initwrite arms a connection SetWriteDeadline during gated responses but never clears it when the handler returns. On HTTP keep-alive, a later ungated request on the same connection can hit that expired deadline and fail the write, even though the safety model says deadlines are cleared so keep-alive is unaffected.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit 8de370d. Configure here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 3 total unresolved issues (including 2 from previous reviews).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit a66fd04. Configure here.
| // panic if it's not an eventsource.Event - as this should be impossible | ||
| return data.([]eventsource.Event), nil | ||
| // The value is always a []eventsource.Event. | ||
| return data.([]eventsource.Event) |
There was a problem hiding this comment.
Stale basis via shared serialize
High Severity
serializePutV1 and serializeBasisV2 single-flight on keys that ignore the current store version, while peek runs in a separate flight. Concurrent replays that peeked different snapshots can join one serialize and receive another caller’s older basis, leaving SDKs on stale data until a later update.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit a66fd04. Configure here.
kinyoklion
marked this pull request as draft
July 31, 2026 15:48
…gression) The acquire-before-serialize rework read the store in a peek() single flight keyed by a fixed string, so all concurrent replays shared one snapshot regardless of basis. A client reconnecting at the current basis could join an older in-flight read taken before its basis existed, fail the up-to-date check against that stale selector, and be sent a full basis at the old state -- losing any delta published between that read and its own subscription (silent data loss, a regression vs base, which keys its replay flight by basis). Key the peek flight by basis, matching base semantics: only same-basis callers share a read (so they also share an identical up-to-date decision), and a client at the current basis reads fresh and is correctly told it is up-to-date. A herd of same-basis reconnects still shares one read. Adds a deterministic regression test (verified to fail with the fixed key restored) where a current-basis client subscribes while an older-basis read is in flight and must still get an up-to-date reply.
…ient copy Addresses the round-3 review's Critical regressions in the progress-aware write deadline, both of which came from arming the deadline for the whole connection lifetime and never clearing it. - C2 (busy streams cut at maxHold): the deadline is now armed only between Begin and the end-of-delivery flush and is cleared there. Live delta/heartbeat traffic after the basis carries no deadline, so a healthy stream that keeps receiving updates past maxHold is not cut. The producer signals the delivery boundary via the connection's writer (End is called before the batch channel closes, so the single end-of-batch flush observes it and clears the deadline). - C3 (idle HTTP/2 streams reset ~21s after the last write): clearing the deadline at end-of-delivery removes the self-firing h2 timer, so idle streams survive. - C1 payload copy: initwrite.Writer now implements io.StringWriter, so eventsource's io.WriteString no longer allocates a []byte copy of the entire basis per connection (the ~105 MB/client OOM driver). The string path copies at most one 1 MiB chunk at a time. - H2/M2: the per-write budget now scales with the bytes written (n/floor + slack) rather than a flat per-chunk value, and the docs and package comment state accurately that the absolute cap governs a delivery larger than ~sendTimeout × floor, so a floor-rate client on a very large data set is cut at the cap. Real-socket tests (httptest + a small SO_SNDBUF/SO_RCVBUF to force a blocking write) drive the actual HandlerV2 -> withInitDeadline -> eventsource -> initwrite chain: a stalled client is cut at the deadline, and a busy stream keeps receiving deltas past maxHold. The busy test is verified to fail if the deadline is scoped to the connection lifetime (the C2 regression). Deferred (documented): C1's slot-held-across-the-FDv1-send (the slot is released on the channel handoff before eventsource writes the single put event). The io.StringWriter change removes its OOM teeth -- the payload is shared, not copied per client -- so the remaining gap is egress-concurrency, not resident memory.
Round-4 review follow-ups (non-blocking): - Fix the inaccurate comment claiming the stream slot is "held across the send" and "bounds egress": the slot is released when the replay goroutine returns (at the channel handoff, before the eventsource handler writes), so it bounds concurrent serialization and resident distinct payloads, not per-connection egress; a stalled send is bounded by the write deadline instead. - Clear the poll connection's write deadline when the handler returns (this server sets no http.Server.WriteTimeout, so net/http does not reset it), so a poll's armed deadline cannot linger on a kept-alive connection and fire during a later request. Correct the initwrite package comment accordingly.
…el handoff Closes M1 (the deferred half of C1/R2-1): MaxConcurrent now bounds concurrent stream sends and resident payloads, including the single-event FDv1 /all put. The budget slot was released when the replay producer goroutine returned, which is at the channel handoff -- the eventsource handler receives the events on another goroutine and only then writes them to the socket. For an FDv1 /all basis (one put event) the producer returned, released the slot, and only afterward did the handler write the multi-MB payload, so the slot bounded nothing for that path (staggered stalled clients measured at 12x payload resident). Now the producer hands the slot's release to the connection's writer via the Begin/Done handshake and holds it across the actual send: initwrite.Writer closes Done at the end-of-basis flush (the last byte written); the producer, after End + closing the batch channel, waits on Done -- or on the request context, which covers a client disconnecting or a send the write deadline cuts -- before releasing. The producer's own goroutine does the wait, so there is no extra goroutine and it remains the single releaser (no double release). Terminal paths (normal completion, disconnect, deadline cut, drain) all release exactly once. Real-socket test TestSocketSlotHeldAcrossSend: with MaxConcurrent=1 and a stalled FDv1 client holding the slot, a second full-basis client is shed. Verified to fail if the slot is released at the handoff (the pre-fix behavior). -race clean.
Comments should describe what the code does now, not what it is modeled on. Remove the two comment references to an internal service (the initwrite package comment and the INIT_SEND_TIMEOUT default constant); the behavior is unchanged.
kinyoklion
force-pushed
the
rlamb/relay-init-concurrency-wiring
branch
2 times, most recently
from
July 31, 2026 20:32
764fc15 to
5f27f66
Compare
…-leak race) The M1 slot-hold release read iw.Done() inside the deferred release, after closeOut(). But closeOut() triggers the eventsource handler's end-of-basis Flush, which both closes the Done channel and nils the writer's reference to it. The producer's Done() read and the handler's nil-ing are driven by the same close(out), with no happens-before, so the producer could observe a nil channel; a receive on nil never fires, dropping the select to ctx.Done() -- which for a healthy long-lived stream only fires at client disconnect. The slot would then be pinned for the whole connection, and under reconnect churn leaked slots accumulate until the budget is exhausted and every init is shed. Capture the channel once, before closeOut(). Done is created in Begin and only nil'd in Flush, so the capture is always the live channel Flush will close. Test TestSocketSlotReleasedAfterHealthyBasis: a healthy client that reads its full basis and stays connected must see the slot return to 0 (released via Done, not pinned until disconnect) -- the existing stall test can't catch this because ctx/the deadline fires there anyway. Verified to fail if Done is read after closeOut. streams -race clean.
…(T1) The round-5 guard passed the whole natural loop even on the buggy (read-Done- after-closeOut) code, because the producer almost always wins the race without scheduler pressure -- so a re-introduction would sail through CI. Add a test-only seam (testHookSlowBasisClose, an atomic, no-op in production) that inserts a small delay between closing the batch channel and the release select, forcing the end-of-basis flush to win. Under that interleaving a stale done-channel capture observes a nil channel and leaks the slot, while capturing before closeOut survives. TestSocketSlotReleasedAfterHealthyBasis now enables the seam and fails deterministically on the buggy form (verified: leaks at iteration 0) instead of ~0.3% of the time.
The eventsource server's Logger was unset, so a connection-write failure -- in particular one the init limiter's write deadline cut to reclaim a slot -- was dropped, making the limiter's cuts invisible in relay logs. Add a WithLogger option that installs a slog adapter as the server-side stream provider's SSE server logger. The adapter distinguishes a relay-initiated cut (a deadline-exceeded error, logged at warn) from an ordinary client disconnect (logged at debug, so normal connection churn does not spam the logs). Wired from relay.go for the server-side stream provider only; other stream kinds are unaffected. Sheds were already logged at warn in the replay path. Tests: the adapter routes a deadline error to warn and a generic write error to debug; WithLogger sets the SSE servers' logger and its absence leaves it unset.
The test seam that forces the end-of-basis flush to win the slot-release race is a package-level atomic. Mark it //nolint:gochecknoglobals with a note that production reads observe the zero value, matching the repo's existing trailing-directive convention.
This was referenced Jul 31, 2026
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
2 participants
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.
Applies the init-concurrency limiter (added to this feature branch by #780) to the two places the Relay Proxy performs an initialization delivery — the full data payload it materializes and sends when an SDK first receives its data — across polling and streaming. This bounds the memory and egress a burst of connecting SDKs can impose, using one shared budget.
The limiter is disabled unless
[Concurrency]/INIT_*is configured, so behavior is unchanged by default.Safety model
Giving up on a delivery must mean closing the connection, never leaving a partly-sent payload on an open one — a dropped connection is safe (the SDK discards the partial and reconnects), but a half-basis left on a live connection can be completed by a later delta into a corrupt data set. Everything below follows from that.
Polling
/sdk/flagsis always a full delivery, so a middleware acquires a slot on entry and sheds with 503 +Retry-After(jittered) when the budget is full./sdk/polland/sdk/poll/evalacquire a slot lazily, only on the full-basis branch — a cheap up-to-date reply (matching basis) builds no payload and is never charged, the same exemption the streaming path uses./sdk/flags/{key}and/sdk/segments/{key}are left ungated (one item, requested per evaluation — not initialization deliveries).Streaming
Builds on the disconnect-aware replay: the producer reads the store once, and if the client's basis already matches it sends a small up-to-date reply without drawing from the budget. For a full basis it acquires a slot before serializing (bounding payload memory), holds it across the send, and releases it when the send finishes or the client disconnects.
A client that stalls without disconnecting is bounded by a write deadline on the connection: a write that blocks past
INIT_SEND_TIMEOUTfails, so the eventsource handler closes the connection — freeing the slot and prompting a clean reconnect — rather than the producer abandoning a partial basis on a still-open connection. When the budget is full, a stream replay likewise closes the connection (the SSE response has already started, so a 503 isn't possible) so the SDK reconnects with backoff instead of stranding uninitialized. Both are armed only when the limiter is enabled, so a disabled deployment is exactly the base behavior.The store read is single-flight-deduplicated with a fixed key (a reconnect herd at any basis shares one read); the per-basis serialization stays deduplicated; and the FDv1 put is materialized under the slot so its memory is accounted there and shared rather than re-serialized per connection.
Behavior
One budget shared by polls and full-basis stream replays:
maxConcurrentin flight,maxQueuedwaiting then shed. Up-to-date replies, deltas, heartbeats, pings, and single-item lookups are never gated. Per-environment fairness was removed for now (#780's cap was never wired); it's safe as a single global pool because a stalled or leaked client can no longer park a slot indefinitely (the deadlines above reclaim it).Testing
Retry-After+ JSON body; disabled → pass-through; slot released on return; up-to-date exemption via provide/acquire) and for streaming (a shed replay closes the connection and delivers nothing). Fullrelayend-to-end suite passes.maxConcurrent=4,maxQueued=0): server FDv2 poll, PHP/sdk/flags, and client-side/sdk/poll/evaleach admit 4 and shed the rest; a cold FDv2/FDv1 stream delivers the full basis; an up-to-date stream/poll is exempt; single-item lookups are served even when the budget is full.maxConcurrenttotal (not per surface); stalled streams occupying the budget cause polls to shed until released.Deferred (follow-ups)
Stats/Rejectedexport, spans) is the planned stacked follow-up PR./flagsstreaming for old server SDKs) is not yet gated and still uses the pre-fix: stop server-side stream replay producer when the SDK client disconnects #774 context-less replay; gating it would replicate this mechanism for a low-traffic legacy endpoint.Note
Medium Risk
Touches hot paths for SDK connect (poll + stream) and connection teardown semantics; misconfiguration or deadline bugs could cause extra 503s or reconnect storms, though default-off behavior limits blast radius.
Overview
Initialization deliveries (full dataset payloads on connect/reconnect) now share one optional global budget across FDv1
/sdk/flags, FDv2 polls, and server-side/sdk/streamreplays. The feature stays off unless[Concurrency]/INIT_MAX_CONCURRENTis set.The limiter is simplified:
PerEnvMaxPercentand per-environment gates are removed;Acquireno longer takes an env key.SendTimeoutis documented as a 2m absolute cap on slot hold time, paired with a newinitwritewrapper that enforces a 64 KB/s throughput floor so slow-but-steady clients are not cut off while stalled clients lose the connection and free the slot.Polling: always-full paths use
LimitConcurrency(slot on entry, 503 + jitteredRetry-Afterwhen full). FDv2 polls useProvideInitLimiterso slots are taken only on the full-basis branch; up-to-date replies stay ungated.Streaming: server-side stream provider gets
WithInitLimiter; full-basis replays acquire before serialize, shed by closing the SSE connection when the budget is full, and wrap responses withinitwrite. Replay logic is refactored to peek the store once, skip budget for matching-basis FDv2 clients, and materialize FDv1 put payloads under the held slot.HTTP middleware (
logging, metricsstatusRecorder) gainsUnwrapso write deadlines reach the real connection. Docs add a[Concurrency]section.Reviewed by Cursor Bugbot for commit a66fd04. Bugbot is set up for automated code reviews on this repo. Configure here.