Fix #126: Release mutex before async write in send_to_connection#145
Fix #126: Release mutex before async write in send_to_connection#145mcurrier2 wants to merge 1 commit into
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughBoth TCP and Unix domain socket IPC transports change ChangesIPC connection locking fix
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant SendToConnection
participant ConnectionsMap
participant Stream
Caller->>SendToConnection: send_to_connection(id, message)
SendToConnection->>ConnectionsMap: lock and remove stream
ConnectionsMap-->>SendToConnection: stream
SendToConnection->>Stream: write_message(message).await
Stream-->>SendToConnection: write result
SendToConnection->>ConnectionsMap: lock and reinsert stream
SendToConnection-->>Caller: result
Related issues: Suggested labels: bug, concurrency, performance Suggested reviewers: (repository maintainers familiar with the IPC transport module) Poem 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
dustinblack
left a comment
dustinblack
left a comment
There was a problem hiding this comment.
Reviewed: correct remove/write/re-insert pattern for mutex-across-await fix. send_to_connection has zero callers, so no runtime impact.
📈 Changed lines coverage: 0.00% (0/10)🚨 Uncovered lines in this PR
📊 Code Coverage Summary
|
|
Definitely not an expert on rust, but claude brought up a point (perhaps already discussed elsewhere) that I thought was interesting. Apparently the tokio::sync::Mutex used in rusty is actually designed for this purpose as a async-aware lock so it doesn't block on .await and can be used to hold a lock across an async operation. In theory, the existing code should be fine as is. Is there more context for the change? |
Both TCP and UDS send_to_connection() held the connections mutex for the entire duration of write_message().await, serializing all multi- client writes and risking deadlock under high contention. Additionally, get_active_connections() silently returned an empty vec when the lock was contended, making it appear as if there were no active connections. Fix send_to_connection in both transports by removing the stream from the map, releasing the lock, performing the async write, then re-inserting the stream. This allows other connection operations to proceed concurrently during writes. For get_active_connections, add a warning log when the lock cannot be acquired so the behavior is observable (the try_lock pattern is retained since this is a sync fn called from async context and must not block the executor). - TCP: remove/write/re-insert pattern in send_to_connection - UDS: same remove/write/re-insert pattern - Both: add warn! log in get_active_connections try_lock failure path - All tests passing including multi-client tests, clippy clean AI-assisted-by: Claude Opus 4.6 Co-authored-by: Cursor <cursoragent@cursor.com>
9ed422d to
1e063e0
Compare
I'll double-check this, but a lot of the problem with tokio for us has been that it adds a significant amount of overhead and leads to negatively impacting our test results. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/ipc/tcp_socket.rs (1)
459-493: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftNew race window: connection_id briefly absent from map during write.
Removing the stream before releasing the lock and reinserting afterward fixes the "lock held across await" issue, but introduces a TOCTOU window where
connection_iddoesn't exist inconnections:
- A concurrent
send_to_connectionfor the sameconnection_idnow fails with a spurious "Connection not found" error instead of serializing as it did previously.- A concurrent
close_connection(connection_id)called while a write is in-flight will also get "not found" (since the entry is removed), signaling failure to the caller — but the stream is silently reinserted once the write completes, effectively resurrecting a connection the caller believed was closed.A more robust fix is to make the per-connection lock granular instead of removing from the map, e.g.
Arc<Mutex<HashMap<ConnectionId, Arc<Mutex<TcpStream>>>>>. Thensend_to_connectiononly holds the outer map lock briefly to clone the per-connectionArc<Mutex<_>>>, and locks that specific stream across the write — no removal, no spurious "not found", and other connections remain unaffected.Sketch of per-connection lock approach
// connections: Arc<Mutex<HashMap<ConnectionId, Arc<Mutex<TcpStream>>>>> async fn send_to_connection(&mut self, connection_id: ConnectionId, message: &Message) -> Result<()> { let stream_lock = { let conns = self.connections.lock().await; conns .get(&connection_id) .cloned() .ok_or_else(|| anyhow!("Connection {} not found", connection_id))? }; let mut stream = stream_lock.lock().await; Self::write_message(&mut stream, message).await.map_err(anyhow::Error::from) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ipc/tcp_socket.rs` around lines 459 - 493, The current send_to_connection flow removes the TcpStream from self.connections, creating a race where the connection is temporarily absent and concurrent send_to_connection or close_connection calls can see spurious "not found" errors. Change the connection handling in TcpSocket so send_to_connection only briefly locks the connections map to clone a per-connection handle, then locks that specific stream across write_message without removing the entry. Update the connections storage to use a per-connection Arc<Mutex<TcpStream>> (or equivalent) so close_connection and other operations can coordinate safely without resurrecting closed connections or blocking unrelated connections.Source: Path instructions
src/ipc/unix_domain_socket.rs (1)
424-458: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftSame remove/reinsert race as tcp_socket.rs.
Identical pattern here creates the same TOCTOU window: concurrent
send_to_connectioncalls for the same connection_id will spuriously fail with "not found", andclose_connectionracing with an in-flight write can be defeated (stream gets reinserted after the "close" already reported failure/removed nothing). See the parallel comment insrc/ipc/tcp_socket.rs(Line 459) for a suggested per-connection-lock fix (Arc<Mutex<HashMap<ConnectionId, Arc<Mutex<UnixStream>>>>>) that avoids removing entries from the map entirely.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ipc/unix_domain_socket.rs` around lines 424 - 458, The send_to_connection flow in UnixDomainSocketManager has the same remove/reinsert race as the TCP path, causing concurrent writes to fail spuriously and close_connection to lose races with in-flight sends. Update send_to_connection and the surrounding connection handling to keep connections in the map and lock per-connection instead of removing the UnixStream, using the same Arc<Mutex<HashMap<ConnectionId, Arc<Mutex<UnixStream>>>>> pattern referenced by the tcp_socket.rs implementation so write_message no longer depends on a temporary removal.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/ipc/tcp_socket.rs`:
- Around line 459-493: The current send_to_connection flow removes the TcpStream
from self.connections, creating a race where the connection is temporarily
absent and concurrent send_to_connection or close_connection calls can see
spurious "not found" errors. Change the connection handling in TcpSocket so
send_to_connection only briefly locks the connections map to clone a
per-connection handle, then locks that specific stream across write_message
without removing the entry. Update the connections storage to use a
per-connection Arc<Mutex<TcpStream>> (or equivalent) so close_connection and
other operations can coordinate safely without resurrecting closed connections
or blocking unrelated connections.
In `@src/ipc/unix_domain_socket.rs`:
- Around line 424-458: The send_to_connection flow in UnixDomainSocketManager
has the same remove/reinsert race as the TCP path, causing concurrent writes to
fail spuriously and close_connection to lose races with in-flight sends. Update
send_to_connection and the surrounding connection handling to keep connections
in the map and lock per-connection instead of removing the UnixStream, using the
same Arc<Mutex<HashMap<ConnectionId, Arc<Mutex<UnixStream>>>>> pattern
referenced by the tcp_socket.rs implementation so write_message no longer
depends on a temporary removal.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Enterprise
Run ID: 96981e28-a141-41c9-b346-d3840dd4f577
📒 Files selected for processing (2)
src/ipc/tcp_socket.rssrc/ipc/unix_domain_socket.rs
|
Good question. You're right that However, safe != optimal here. The issue is about lock duration, not lock safety.
The remove/re-insert pattern reduces lock hold time from "duration of network I/O" (milliseconds to seconds) to "duration of HashMap lookup" (nanoseconds). The cost is two lock acquisitions instead of one, and the connection is briefly invisible in the map during the write. Since In short: the existing code wouldn't deadlock or block the runtime, but it would serialize multi-client throughput unnecessarily. The fix is a net improvement. |
|
I guess the question for me is should this PR consider the send_to_connection case (which currently has no callers), or is that out of scope for this PR? |
3 participants
Summary
send_to_connection()to not hold the connections mutex across thewrite_message().awaitcallwarn!log inget_active_connections()whentry_lock()fails, making the behavior observableTest plan
test_tcp_multi_client)test_unix_domain_socket_multi_client)Fixes #126
Made with Cursor