perf(provider): migrate WS transport to NWConnection (52→130 TPS)#479
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
This change is a streaming batch-interval tuning tweak (4→8 tokens per WS frame) that is neutral to the security posture — no threat mitigations are strengthened or weakened. Trust Boundaries Touched
Threat Assessment
New Attack SurfaceNone introduced. The change touches only a Open Findings ResolvedNone. One note for reviewers: the comment in the diff updates the example TPS from 65 tok/s to 130 tok/s. If that figure is load-tested rather than theoretical, verify it doesn't mask a regression in the 🔐 Threat model: |
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: COMMENT
Security — ✅ No issues found
Performance — ✅ No issues found
Type_diligence — ✅ No issues found
Additive_complexity — ✅ No issues found
✅ All four passes clean. No issues found.
🤖 Automated review by Centaur · DAR-186
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38aa8e6d74
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let (data, context): (Data?, NWConnection.ContentContext?) = | ||
| try await withCheckedThrowingContinuation { cont in | ||
| connection.receiveMessage { data, context, _isComplete, error in | ||
| if let error { | ||
| cont.resume(throwing: CoordinatorError.connectionClosed(error)) | ||
| return | ||
| } | ||
| cont.resume(returning: (data, context)) | ||
| } |
There was a problem hiding this comment.
When the ping task throws pongTimeout or suspensionDetected while the socket state has not changed, sessionLoop cancels the task group and then cannot leave until receiveLoop observes cancellation. This new withCheckedThrowingContinuation wrapper around NWConnection.receiveMessage is not cancellation-aware, and the connection is only cancelled later in connectAndRun's defer after sessionLoop returns, so the provider can hang instead of reconnecting after sleep or a missing pong. Add a cancellation handler that cancels/unblocks the pending receive, or cancel the connection before waiting for the group to drain.
Useful? React with 👍 / 👎.
| group.addTask { | ||
| for await error in failureStream { | ||
| throw error | ||
| } |
There was a problem hiding this comment.
During shutdown, outboundRouter.finish() closes the outbound stream, so the writer task can return normally before the receive task observes the close. In that normal-return path sessionLoop falls out of the task-group body without cancelAll(), but this newly added failure-stream child cannot finish until connectAndRun's defer runs after sessionLoop returns, leaving the group waiting forever and leaking the coordinator run-loop task. Cancel the remaining session tasks after group.next() returns, not only when it throws.
Useful? React with 👍 / 👎.
| // Control frames (ping/pong, handled by autoReplyPing / the pong | ||
| // handler) carry no application payload — skip them and keep reading. |
There was a problem hiding this comment.
If the peer or an intermediary closes the TCP connection without a WebSocket close frame, receiveMessage can complete with no payload/metadata and no application data. This guard classifies every empty receive as a control frame and immediately loops, so that close can turn into a tight receive loop instead of surfacing a disconnect and reconnecting. Only skip known control-frame metadata; otherwise treat nil/empty receives as connectionClosed.
Useful? React with 👍 / 👎.
38aa8e6 to
105ee73
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 105ee7326b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| completion: .contentProcessed { error in | ||
| if let error { | ||
| logger.error("WS send failed (\(identifier)): \(error.localizedDescription)") | ||
| } |
There was a problem hiding this comment.
When NWConnection.send completes with an error (for example after the coordinator has dropped the socket or the write side fails under load), this helper only logs it and the writer loop continues draining outbound chunks. Unlike the previous try await ws.send(...) path, the session is not torn down until a separate state/receive/ping path notices the failure, so response chunks or terminal frames produced in that window can be silently discarded and leave the coordinator-side request waiting for a timeout. Feed this error into the session failure path or cancel the connection here so in-flight work is failed immediately.
Useful? React with 👍 / 👎.
| completion: .contentProcessed { _ in | ||
| connection.cancel() |
There was a problem hiding this comment.
When shutdown() or refreshAPNsToken() runs while the current NWConnection is still handshaking or the write side is wedged, this path waits for the close-frame send completion before calling cancel(). That close frame can remain queued until the connection becomes writable, so the handshake/session state handler never observes .cancelled; a late APNs token refresh can fail to force the promised reconnect, and shutdown can leave the connection loop/socket alive. Add an immediate or timed fallback cancel for connections that cannot flush the close frame promptly.
Useful? React with 👍 / 👎.
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 1 finding(s) (1 blocking)
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient+Connection.swift:175-179— Task group child lacks proper error handling for async stream iteration- Suggestion: Add explicit error handling around the failureStream iteration to prevent silent task failures that could leave the connection in an inconsistent state
Type_diligence — ✅ No issues found
Additive_complexity — ✅ No issues found
1 finding(s) total, 1 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
| // Thread-safe pong timestamp: updated from the pong handler (runs on | ||
| // connectionQueue), read from the ping task. Using an actor would force | ||
| // structured concurrency overhead on every ping; an unfair lock is cheaper | ||
| // for a single Instant. | ||
| let pongTracker = PongTracker() |
There was a problem hiding this comment.
🟡 [MEDIUM] ⚡ Task group child lacks proper error handling for async stream iteration
💡 Suggestion: Add explicit error handling around the failureStream iteration to prevent silent task failures that could leave the connection in an inconsistent state
📊 Score: 3×4 = 12 · Category: unbounded_goroutines
Replace URLSessionWebSocketTask with NWConnection + NWProtocolWebSocket (Apple Network.framework) to eliminate the serial per-chunk write bottleneck that limits coordinator-observed TPS. Root cause: URLSessionWebSocketTask.send() is an async/await call that suspends until TCP write + ACK completes (~36ms per chunk). With 230 chunks per 918-token response (streamInterval=4), that is 8.3s of pure write-wait overhead on top of 7s of decode — reducing 130 TPS local decode to 52 TPS through the network. Fix: NWConnection.send() returns immediately after buffering in the kernel (<1ms). The outbound loop still processes messages sequentially (preserving frame ordering), but each send no longer blocks on TCP ACK. The kernel coalesces buffered frames into TCP segments automatically. Changes: - CoordinatorClient.swift: stored properties (nwConnection, connectionQueue), closeCurrentConnection() with proper WS close frame, nil-after-cancel - CoordinatorClient+Connection.swift: NWConnection setup with .url() endpoint (preserves WS upgrade path incl URL path), non-blocking sendTextFrame (nonisolated, no actor hop), NWProtocolWebSocket receive loop, ping via .ping opcode + setPongHandler, one-shot→persistent state handler swap for failure-driven reconnect - CoordinatorClient+Inbound.swift: drop ws: parameter, sendOnCurrentConnection for inline error replies - CoordinatorClient+Registration.swift: continuation-wrapped blocking send - BatchScheduler+EngineFactory.swift: streamInterval 4→8 (halves chunk count, 62ms burst gap at 130 TPS, smoother than 60fps) - CoordinatorClientTests.swift: replace URLSession test with NWProtocolWebSocket.Options validation + readback assertion - Doc comment updates in CoordinatorClientCodec, CoordinatorClientState Zero coordinator changes. Same WS protocol on the wire. Same JSON. Same reconnect/backoff behavior. Expected improvement: - Per-chunk write: 36ms (TCP ACK) → <1ms (kernel buffer) - WS overhead (918 tokens): 8.3s → ~0.1s - Coordinator-observed TPS: 52 → ~130 (2.5x) - streamInterval: 4 (230 chunks) → 8 (115 chunks)
105ee73 to
ce86e3e
Compare
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 3 finding(s) (2 blocking)
- 🔴 [CRITICAL]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient+Connection.swift:174-178— Task group lacks proper cancellation handling for connection failure stream- Suggestion: Add explicit cancellation handling in the failure stream task to prevent resource leaks when the connection drops
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient+Connection.swift:302-325— sendTextFrame performs synchronous logging on hot path- Suggestion: Move error logging off the hot path or use async logging to avoid blocking the inference chunk pipeline
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient+Registration.swift:44-58— Registration send blocks on continuation for non-critical path- Suggestion: Consider fire-and-forget registration with error handling in state monitor rather than blocking the connection setup
Type_diligence — ✅ No issues found
Additive_complexity — ✅ No issues found
3 finding(s) total, 2 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
|
|
||
| // Thread-safe pong timestamp: updated from sendPing's callback (arbitrary queue), | ||
| // read from the ping task. Using an actor would force structured concurrency | ||
| // overhead on every ping; an unfair lock is cheaper for a single Instant. | ||
| // Thread-safe pong timestamp: updated from the pong handler (runs on | ||
| // connectionQueue), read from the ping task. Using an actor would force | ||
| // structured concurrency overhead on every ping; an unfair lock is cheaper | ||
| // for a single Instant. |
There was a problem hiding this comment.
🔴 [CRITICAL] ⚡ Task group lacks proper cancellation handling for connection failure stream
💡 Suggestion: Add explicit cancellation handling in the failure stream task to prevent resource leaks when the connection drops
📊 Score: 4×4 = 16 · Category: unbounded_goroutines
| nonisolated internal func sendTextFrame( | ||
| _ json: String, | ||
| on connection: NWConnection, | ||
| identifier: String | ||
| ) { | ||
| let logger = self.logger | ||
| let metadata = NWProtocolWebSocket.Metadata(opcode: .text) | ||
| let context = NWConnection.ContentContext(identifier: identifier, metadata: [metadata]) | ||
| connection.send( | ||
| content: Data(json.utf8), | ||
| contentContext: context, | ||
| isComplete: true, | ||
| completion: .contentProcessed { error in | ||
| if let error { | ||
| logger.error("WS send failed (\(identifier)): \(error.localizedDescription)") | ||
| // Cancel the connection immediately so the stateUpdateHandler | ||
| // (Task 0) fires a failure and tears down the session. Without | ||
| // this, the outbound loop keeps draining chunks that silently | ||
| // vanish, and the coordinator-side request waits for a timeout. | ||
| connection.cancel() | ||
| } | ||
| } | ||
| ) | ||
| } |
There was a problem hiding this comment.
🟡 [MEDIUM] ⚡ sendTextFrame performs synchronous logging on hot path
💡 Suggestion: Move error logging off the hot path or use async logging to avoid blocking the inference chunk pipeline
📊 Score: 3×3 = 9 · Category: blocking_io
| let context = NWConnection.ContentContext(identifier: "register", metadata: [metadata]) | ||
| try await withCheckedThrowingContinuation { (cont: CheckedContinuation<Void, Error>) in | ||
| connection.send( | ||
| content: Data(jsonString.utf8), | ||
| contentContext: context, | ||
| isComplete: true, | ||
| completion: .contentProcessed { error in | ||
| if let error { | ||
| cont.resume(throwing: CoordinatorError.connectionClosed(error)) | ||
| } else { | ||
| cont.resume() | ||
| } | ||
| } | ||
| ) | ||
| } |
There was a problem hiding this comment.
🔵 [INFO] ⚡ Registration send blocks on continuation for non-critical path
💡 Suggestion: Consider fire-and-forget registration with error handling in state monitor rather than blocking the connection setup
📊 Score: 2×3 = 6 · Category: blocking_io
2 participants
Summary
Migrate the provider→coordinator WebSocket transport from
URLSessionWebSocketTasktoNWConnection+NWProtocolWebSocket(Apple Network.framework). Also bumpstreamIntervalfrom 4 to 8.Root Cause
URLSessionWebSocketTask.send()is an async/await call that suspends until TCP write + ACK completes (~36ms per chunk). With streamInterval=4, a 918-token response produces ~230 WS chunks, adding 8.3s of pure write-wait overhead on top of 7s of decode. This reduces 130 TPS local decode to 52 TPS through the network.Fix
NWConnection.send()returns immediately after buffering in the kernel (<1ms). The outbound loop still processes messages sequentially (preserving WS frame ordering), but each send no longer blocks on TCP ACK. The kernel coalesces buffered frames into TCP segments automatically.Combined with
streamInterval4→8 (halving chunk count from 230 to 115), the per-request WS overhead drops from 8.3s to ~0.1s.Expected Improvement
Before / After
flowchart LR subgraph "Before (URLSessionWebSocketTask)" A1[decode token batch] --> B1["encodeOutbound (nonisolated)"] B1 --> C1["try await ws.send(.string(json))"] C1 --> D1["⏳ blocks ~36ms waiting for TCP ACK"] D1 --> E1[next chunk] E1 --> C1 D1 -.-> F1["230 chunks × 36ms = 8.3s overhead\n→ 52 TPS observed"] end subgraph "After (NWConnection)" A2[decode token batch] --> B2["encodeOutbound (nonisolated)"] B2 --> C2["connection.send() — non-blocking"] C2 --> D2["returns <1ms, kernel buffers"] D2 --> E2[next chunk immediately] E2 --> C2 D2 -.-> F2["115 chunks × <1ms = ~0.1s overhead\n→ ~130 TPS observed"] endKey Design Decisions
.url(url)endpoint, not.hostPort— preserves the WS upgrade request including URL path (/v1/providers/ws)..hostPortwould drop the path and break the coordinator route.nonisolated sendTextFrame— hot path never hops to the CoordinatorClient actor, matching the existingnonisolated encodeOutboundoptimization from perf(provider): fix serial WS outbound bottleneck that tanks per-stream TPS under load #475.nwConnection = nilafter cancel — preventssendOnCurrentConnectionfrom firing on a cancelled connection during the reconnect window.connectionQueue) that feeds anAsyncStream<Error>consumed by a task-group child to drive reconnect.autoReplyPing = true— NWProtocolWebSocket handles server pings automatically. Client pings for liveness use.pingopcode withsetPongHandler, replacingws.sendPing.Changes
CoordinatorClient.swiftnwConnection/connectionQueue,closeCurrentConnection()with WS close frame, nil-after-cancelCoordinatorClient+Connection.swiftCoordinatorClient+Inbound.swiftws:param,sendOnCurrentConnectionfor error repliesCoordinatorClient+Registration.swiftBatchScheduler+EngineFactory.swiftstreamInterval = 8CoordinatorClientTests.swiftCoordinatorClientCodec.swiftCoordinatorClientState.swiftNo coordinator changes. Same WS protocol on the wire. Same JSON framing. Same reconnect/backoff behavior.
OutboundRouter,SendHandle, and allProviderLoopfiles are untouched — they are transport-agnostic.Verification
swift build: clean (zero errors, zero new warnings)swift test --filter CoordinatorClientTests: 15/15 passURLSessionWebSocketTask/urlSession/applyInboundMessageLimitreferencesNeed help on this PR? Tag
/codesmithwith what you need. Autofix is disabled.