refactor: redone concurency model - #7
Conversation
WalkthroughThis update removes the Changes
Sequence Diagram(s)sequenceDiagram
participant Main
participant TunListener
participant UdpListener
participant ResultCoordinator
participant TunDevice
participant UdpSocket
Main->>TunListener: spawn tun_listener()
Main->>UdpListener: spawn udp_listener()
Main->>ResultCoordinator: spawn result_coordinator()
TunDevice->>TunListener: Packet received
TunListener->>ResultCoordinator: EncryptedPacket via channel
UdpSocket->>UdpListener: Packet received
UdpListener->>ResultCoordinator: DecryptedPacket via channel
ResultCoordinator->>TunDevice: Write DecryptedPacket
ResultCoordinator->>UdpSocket: Send EncryptedPacket
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @bernoussama, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly refactors the ipou codebase by overhauling its concurrency model and improving its overall structure. The changes aim to enhance modularity, maintainability, and performance, particularly in packet handling, by distributing responsibilities across dedicated asynchronous tasks and optimizing memory usage.
Highlights
- Concurrency Model Redesign: The core change is a shift from a single
tokio::select!loop insrc/main.rsto a more modular, task-based concurrency model. This involves introducing dedicatedtun_listener,udp_listener, andresult_coordinatortasks, each handling specific aspects of packet processing, managed bytokio::spawnandtokio::try_join!. - Code Structure Reorganization: Shared types like
PeerandIpouErrorhave been consolidated intosrc/lib.rs, and a newsrc/tasksmodule has been created to house the new asynchronous task logic, improving overall modularity and separation of concerns. Thesrc/error.rsandsrc/ipou.rsfiles were removed as part of this consolidation. - Packet Handling Optimization: Packet processing functions (
handle_tun_packet,handle_udp_packet) insrc/net/mod.rsnow utilize pre-allocated fixed-size buffers ([u8; MTU],[u8; MTU + 512]) instead of dynamically sized slices, aiming to reduce memory allocations and potentially improve performance. Theresult_coordinatortask now centralizes the sending of processed packets. - Centralized Constants and Types: Key constants (
MTU,CHANNEL_BUFFER_SIZE,ENCRYPTION_OVERHEAD) and new type aliases (DecryptedPacket,EncryptedPacket) along with message enums (TunMessage,UdpMessage) have been defined insrc/lib.rsfor better clarity and centralized configuration.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Pull Request Overview
This pull request refactors the concurrency model of the ipou codebase to improve modularity and performance. The changes replace a monolithic tokio::select! loop with task-based concurrency using dedicated listener functions and channels for communication.
Key changes include:
- Introduction of task-based concurrency with separate TUN/UDP listeners and result coordinator
- Code consolidation by moving shared types and constants to
src/lib.rs - Optimization of packet handling with pre-allocated fixed-size buffers
Reviewed Changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/tasks/mod.rs | New module implementing task-based concurrency with tun_listener, udp_listener, and result_coordinator functions |
| src/net/mod.rs | Updated packet handlers to use fixed-size buffers instead of slices for better performance |
| src/main.rs | Refactored main loop to spawn concurrent tasks using tokio::try_join! instead of tokio::select! |
| src/lib.rs | Consolidated shared types, constants, and error definitions from separate modules |
| src/ipou.rs | Removed file - Peer struct moved to lib.rs |
| src/error.rs | Removed file - error types moved to lib.rs |
| } | ||
| } | ||
|
|
||
| // Receive enccrypted packets from channel and send to UDP |
There was a problem hiding this comment.
The word 'enccrypted' is misspelled. It should be 'encrypted'.
| // Receive enccrypted packets from channel and send to UDP | |
| // Receive encrypted packets from channel and send to UDP |
| DecryptedPacket, | ||
| Shutdown, | ||
| } | ||
|
|
||
| pub type EncryptedPacket = (Vec<u8>, SocketAddr); | ||
| #[derive(Debug, Clone)] | ||
| pub enum UdpMessage { | ||
| EncryptedPacket, |
There was a problem hiding this comment.
The TunMessage and UdpMessage enums are defined but appear unused. The enum variants don't contain data that would correspond to their names (DecryptedPacket, EncryptedPacket), making the design unclear.
| DecryptedPacket, | |
| Shutdown, | |
| } | |
| pub type EncryptedPacket = (Vec<u8>, SocketAddr); | |
| #[derive(Debug, Clone)] | |
| pub enum UdpMessage { | |
| EncryptedPacket, | |
| DecryptedPacket(DecryptedPacket), // Includes the actual packet data | |
| Shutdown, | |
| } | |
| pub type EncryptedPacket = (Vec<u8>, SocketAddr); | |
| #[derive(Debug, Clone)] | |
| pub enum UdpMessage { | |
| EncryptedPacket(EncryptedPacket), // Includes the encrypted packet and its destination |
| packet: &mut Vec<u8>, | ||
| conf_clone: Arc<Config>, | ||
| runtime_conf: Arc<RuntimeConfig>, | ||
| utx_clone: mpsc::Sender<(Vec<u8>, SocketAddr)>, |
There was a problem hiding this comment.
[nitpick] The parameter name 'utx_clone' suggests it's a clone, but it's actually the original sender passed to the function. Consider renaming to 'etx' or 'encrypted_tx' to match the naming convention used in the tasks module.
| utx_clone: mpsc::Sender<(Vec<u8>, SocketAddr)>, | |
| encrypted_tx: mpsc::Sender<(Vec<u8>, SocketAddr)>, |
There was a problem hiding this comment.
Code Review
This pull request refactors the application's concurrency model, improving modularity. The review focuses on improving robustness and correctness, specifically around error handling and resource management. The error handling for the joined tasks in main.rs is flawed, and errors from I/O operations in the result_coordinator task are currently ignored. The tun_listener spawns a new task for every incoming packet, which can be inefficient.
| tokio::try_join!(tun_listener, udp_listener, result_coordinator) | ||
| .map(|_| ()) | ||
| .expect("Error joining tasks"); |
There was a problem hiding this comment.
The .map(|_| ()) call discards the Results returned by the individual tasks. This means if tun_listener or udp_listener returns an IpouError, this error will be silently ignored, and the program will continue to run in a degraded state. The .expect() will only trigger on a JoinError, which typically happens if a task panics, not when it returns an Err.
You should handle both the JoinError from try_join! and the individual Results from each task.
| tokio::try_join!(tun_listener, udp_listener, result_coordinator) | |
| .map(|_| ()) | |
| .expect("Error joining tasks"); | |
| let (tun_res, udp_res, coord_res) = tokio::try_join!(tun_listener, udp_listener, result_coordinator)?; | |
| tun_res?; | |
| udp_res?; | |
| coord_res?; |
| match dev.send(&decrypted_packet).await { | ||
| Ok(_sent) => {}, | ||
| Err(_e) => {}, | ||
| } |
There was a problem hiding this comment.
Errors from dev.send() are silently ignored. This is dangerous as it can hide critical problems with the TUN device, causing silent packet loss. The error should be logged.
if let Err(e) = dev.send(&decrypted_packet).await {
eprintln!("Error sending packet to TUN device: {e}");
}| match sock.send_to(&encrypted_packet, peer_addr).await { | ||
| Ok(sent) => { | ||
| #[cfg(debug_assertions)] | ||
| println!("Sent {sent} bytes to {peer_addr}"); | ||
| }, | ||
| Err(_e) => {}, | ||
| } |
There was a problem hiding this comment.
Errors from sock.send_to() are silently ignored. This can mask network issues or other problems preventing data from being sent to a peer. This error should be logged.
match sock.send_to(&encrypted_packet, peer_addr).await {
Ok(sent) => {
#[cfg(debug_assertions)]
println!("Sent {sent} bytes to {peer_addr}");
},
Err(e) => {
eprintln!("Error sending encrypted packet to peer {peer_addr}: {e}");
}
}| pub type DecryptedPacket = Vec<u8>; | ||
| #[derive(Debug, Clone)] | ||
| pub enum TunMessage { | ||
| DecryptedPacket, | ||
| Shutdown, | ||
| } | ||
|
|
||
| pub type EncryptedPacket = (Vec<u8>, SocketAddr); | ||
| #[derive(Debug, Clone)] | ||
| pub enum UdpMessage { | ||
| EncryptedPacket, | ||
| Shutdown, | ||
| } |
There was a problem hiding this comment.
The TunMessage and UdpMessage enums are defined but the variants DecryptedPacket in TunMessage and EncryptedPacket in UdpMessage are unit-like and don't carry any data. Consider having them carry data, e.g. TunMessage::Decrypted(DecryptedPacket).
| if len >= 20 { | ||
| tokio::spawn(crate::net::handle_tun_packet( | ||
| tun_buf, | ||
| len, | ||
| Arc::clone(&conf_clone), | ||
| Arc::clone(&runtime_conf), | ||
| etx.clone(), | ||
| )); | ||
| } |
There was a problem hiding this comment.
The tun_listener spawns a new task for every incoming packet. This "spawn-per-packet" model can be inefficient and lead to resource exhaustion under high packet load due to scheduler overhead and memory consumption. Consider using a bounded worker pool to limit concurrency.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (7)
src/lib.rs (2)
22-34: The enum variants should carry data instead of being unit-like.The
DecryptedPacketvariant inTunMessageandEncryptedPacketvariant inUdpMessagedon't carry any data, making the design unclear. These should hold the actual packet data they represent.
23-34: Remove unusedTunMessageandUdpMessageenumsBoth
TunMessageandUdpMessageare only defined insrc/lib.rsand never referenced elsewhere in the codebase. To reduce dead code and avoid confusion, either remove these enums or document their intended future use.Locations to address:
- src/lib.rs: lines 23–34
Suggested change (if removing):
- #[derive(Debug, Clone)] - pub enum TunMessage { - DecryptedPacket, - Shutdown, - } - - #[derive(Debug, Clone)] - pub enum UdpMessage { - EncryptedPacket, - Shutdown, - }src/main.rs (1)
101-103: Handle individual task errors properly instead of discarding them.The
.map(|_| ())discards the Results from individual tasks, which means errors fromtun_listenerorudp_listenerwill be silently ignored. Only JoinErrors from panics will be caught.src/tasks/mod.rs (4)
83-83: Fix typo in comment.
77-80: Log errors when sending to TUN device.Silently ignoring errors can hide critical problems causing packet loss.
87-93: Log errors when sending to UDP socket.Network errors should be logged for debugging and monitoring.
24-32: Consider using a bounded worker pool instead of spawn-per-packet.Spawning a new task for every packet can lead to resource exhaustion under high load.
🧹 Nitpick comments (2)
src/main.rs (1)
57-57: Remove unnecessary Arc clone.The
runtime_config_clonevariable is unnecessary since you can reuseruntime_configdirectly in theudp_listenerspawn.- let runtime_config_clone = Arc::clone(&runtime_config); - let mut tun_config = tun::Configuration::default();let udp_listener = tokio::spawn(tasks::udp_listener( Arc::clone(&sock_arc), - runtime_config_clone, + runtime_config, dtx, ));Also applies to: 89-93
src/net/mod.rs (1)
71-71: Remove unnecessary clone when sending packet.The
packet.clone()is unnecessary sincepacketis a locally owned Vec that can be moved directly.- if let Err(e) = etx.send((packet.clone(), peer.sock_addr)).await { + if let Err(e) = etx.send((packet, peer.sock_addr)).await {
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
src/error.rs(0 hunks)src/ipou.rs(0 hunks)src/lib.rs(1 hunks)src/main.rs(3 hunks)src/net/mod.rs(4 hunks)src/tasks/mod.rs(1 hunks)
💤 Files with no reviewable changes (2)
- src/ipou.rs
- src/error.rs
🧰 Additional context used
🧬 Code Graph Analysis (3)
src/main.rs (1)
src/tasks/mod.rs (3)
tun_listener(12-35)udp_listener(37-59)result_coordinator(61-97)
src/tasks/mod.rs (2)
src/main.rs (2)
mpsc(79-79)mpsc(81-81)src/net/mod.rs (2)
handle_tun_packet(45-90)handle_udp_packet(10-43)
src/net/mod.rs (1)
src/main.rs (2)
mpsc(79-79)mpsc(81-81)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (rust)
🔇 Additional comments (2)
src/tasks/mod.rs (2)
18-32: Critical: Buffer reuse causes data races.The
tun_bufis reused in the loop but passed by value to spawned tasks. Since arrays are copied, and the spawned tasks execute asynchronously, the next loop iteration may overwrite the buffer before the spawned task processes it, causing data corruption.loop { // Listen for TUN packets let len = dev.recv(&mut tun_buf).await?; // Spawn handler task for each packet if len >= 20 { + let packet_data = tun_buf[..len].to_vec(); tokio::spawn(crate::net::handle_tun_packet( - tun_buf, + packet_data, len, Arc::clone(&conf_clone), Arc::clone(&runtime_conf), etx.clone(), )); }Note: This will require updating
handle_tun_packetto accept aVec<u8>instead of a fixed array.Likely an incorrect or invalid review comment.
42-56: Critical: Buffer reuse causes data races.Same issue as tun_listener - the
udp_bufis reused while spawned tasks may still be reading from it.loop { // Listen for UDP packets let (len, peer_addr) = sock.recv_from(&mut udp_buf).await?; // Spawn handler task for each packet if len >= 28 { // 12 bytes nonce + 16 bytes auth tag + let packet_data = udp_buf[..len].to_vec(); tokio::spawn(crate::net::handle_udp_packet( - udp_buf, + packet_data, len, peer_addr, Arc::clone(&runtime_conf), dtx.clone(), )); };Note: This will require updating
handle_udp_packetto accept aVec<u8>instead of a fixed array.Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
src/tasks/mod.rs (2)
24-32: Consider using a bounded worker pool instead of spawn-per-packet modelThe current implementation spawns a new task for every incoming packet, which can lead to resource exhaustion and scheduler overhead under high packet load.
88-88: Fix typo in commentThe word 'enccrypted' should be 'encrypted'.
- // Receive enccrypted packets from channel and send to UDP + // Receive encrypted packets from channel and send to UDP
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/tasks/mod.rs(1 hunks)
🔇 Additional comments (1)
src/tasks/mod.rs (1)
77-85: Good error handling implementationThe error handling for both TUN device and UDP socket operations properly logs errors instead of silently ignoring them, which addresses the previous review concerns.
Also applies to: 92-100
| if len >= 28 { | ||
| // 12 bytes nonce + 16 bytes auth tag | ||
| tokio::spawn(crate::net::handle_udp_packet( | ||
| udp_buf, | ||
| len, | ||
| peer_addr, | ||
| Arc::clone(&runtime_conf), | ||
| dtx.clone(), | ||
| )); | ||
| }; |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Apply the same worker pool pattern here
Similar to tun_listener, this function also spawns a task per packet which can cause performance issues under load.
🤖 Prompt for AI Agents
In src/tasks/mod.rs around lines 47 to 56, the code currently spawns a new Tokio
task for each UDP packet, which can lead to performance degradation under high
load. Refactor this section to use a worker pool pattern similar to the one used
in tun_listener. Instead of spawning a new task per packet, send the packet data
to a bounded channel or task queue that a fixed number of worker tasks consume,
thereby limiting concurrency and improving resource management.
This pull request refactors the
ipoucodebase to improve modularity, maintainability, and performance. The most significant changes include reorganizing the code structure, introducing task-based concurrency, and optimizing packet handling. Below is a categorized summary of the key changes:Code Structure and Organization
Peerstruct andIpouErrortype fromsrc/ipou.rsandsrc/error.rsrespectively tosrc/lib.rs, consolidating shared types into a central location. Theerrormodule was removed. [1] [2] [3]tasksmodule insrc/tasks/mod.rsto encapsulate asynchronous task logic for handling TUN and UDP packets, as well as coordinating results. [1] [2]Constants and Type Definitions
MTU,CHANNEL_BUFFER_SIZE, andENCRYPTION_OVERHEADtosrc/lib.rsfor centralized configuration.DecryptedPacketandEncryptedPacketfor better type clarity and consistency.Task-Based Concurrency
tokio::select!loop insrc/main.rswith three concurrent tasks:tun_listener,udp_listener, andresult_coordinator. These tasks are spawned usingtokio::spawnand managed withtokio::try_join!. [1] [2]tun_listenerandudp_listenerfunctions in thetasksmodule.Packet Handling Optimization
handle_tun_packetandhandle_udp_packetinsrc/net/mod.rsto use pre-allocated buffers ([u8; MTU]and[u8; MTU + 512]) instead of slices, reducing memory allocations. [1] [2]result_coordinatortask to streamline the sending of packets to the TUN device and UDP socket.Minor Cleanup
src/main.rsafter moving them tosrc/lib.rs.These changes improve the modularity of the codebase by separating concerns, enhance performance through optimized packet handling, and make the code easier to maintain by introducing clear abstractions for tasks.
Summary by CodeRabbit
New Features
Refactor
Chores