nntp_proxy/network/
optimizers.rs

1//! Network optimization traits and implementations
2//!
3//! This module provides a trait-based approach to network optimizations,
4//! allowing different optimization strategies for TCP and TLS connections.
5
6use crate::stream::ConnectionStream;
7use std::io;
8use tokio::net::TcpStream;
9use tokio_rustls::client::TlsStream;
10use tracing::debug;
11
12/// Trait for network optimization strategies
13pub trait NetworkOptimizer {
14    /// Apply optimizations to improve network performance
15    fn optimize(&self) -> Result<(), io::Error>;
16
17    /// Get a description of the optimization strategy
18    fn description(&self) -> &'static str;
19}
20
21/// TCP-specific optimizations for high-throughput scenarios
22pub struct TcpOptimizer<'a> {
23    stream: &'a TcpStream,
24    recv_buffer_size: usize,
25    send_buffer_size: usize,
26}
27
28impl<'a> TcpOptimizer<'a> {
29    /// Create a new TCP optimizer with default high-throughput settings
30    pub fn new(stream: &'a TcpStream) -> Self {
31        Self {
32            stream,
33            recv_buffer_size: crate::constants::socket::HIGH_THROUGHPUT_RECV_BUFFER,
34            send_buffer_size: crate::constants::socket::HIGH_THROUGHPUT_SEND_BUFFER,
35        }
36    }
37
38    /// Create a TCP optimizer with custom buffer sizes
39    pub fn with_buffer_sizes(stream: &'a TcpStream, recv_size: usize, send_size: usize) -> Self {
40        Self {
41            stream,
42            recv_buffer_size: recv_size,
43            send_buffer_size: send_size,
44        }
45    }
46}
47
48impl<'a> NetworkOptimizer for TcpOptimizer<'a> {
49    fn optimize(&self) -> Result<(), io::Error> {
50        use socket2::SockRef;
51
52        let sock_ref = SockRef::from(self.stream);
53
54        // Set larger buffer sizes for high throughput
55        sock_ref.set_recv_buffer_size(self.recv_buffer_size)?;
56        sock_ref.set_send_buffer_size(self.send_buffer_size)?;
57
58        debug!(
59            "Applied TCP optimizations: recv_buffer={}, send_buffer={}",
60            self.recv_buffer_size, self.send_buffer_size
61        );
62
63        Ok(())
64    }
65
66    fn description(&self) -> &'static str {
67        "TCP high-throughput optimization"
68    }
69}
70
71/// TLS-specific optimizations that work on the underlying TCP stream
72pub struct TlsOptimizer<'a> {
73    stream: &'a TlsStream<TcpStream>,
74    recv_buffer_size: usize,
75    send_buffer_size: usize,
76}
77
78impl<'a> TlsOptimizer<'a> {
79    /// Create a new TLS optimizer with default settings
80    pub fn new(stream: &'a TlsStream<TcpStream>) -> Self {
81        Self {
82            stream,
83            recv_buffer_size: crate::constants::socket::HIGH_THROUGHPUT_RECV_BUFFER,
84            send_buffer_size: crate::constants::socket::HIGH_THROUGHPUT_SEND_BUFFER,
85        }
86    }
87
88    /// Create a TLS optimizer with custom buffer sizes
89    pub fn with_buffer_sizes(
90        stream: &'a TlsStream<TcpStream>,
91        recv_size: usize,
92        send_size: usize,
93    ) -> Self {
94        Self {
95            stream,
96            recv_buffer_size: recv_size,
97            send_buffer_size: send_size,
98        }
99    }
100}
101
102impl<'a> NetworkOptimizer for TlsOptimizer<'a> {
103    fn optimize(&self) -> Result<(), io::Error> {
104        use socket2::SockRef;
105
106        // Get the underlying TCP stream for optimization
107        let tcp_stream = self.stream.get_ref().0;
108        let sock_ref = SockRef::from(tcp_stream);
109
110        // Apply TCP-level optimizations to the underlying stream
111        sock_ref.set_recv_buffer_size(self.recv_buffer_size)?;
112        sock_ref.set_send_buffer_size(self.send_buffer_size)?;
113
114        debug!(
115            "Applied TLS optimizations to underlying TCP stream: recv_buffer={}, send_buffer={}",
116            self.recv_buffer_size, self.send_buffer_size
117        );
118
119        Ok(())
120    }
121
122    fn description(&self) -> &'static str {
123        "TLS optimization via underlying TCP stream"
124    }
125}
126
127/// High-level optimizer that works with ConnectionStream
128pub struct ConnectionOptimizer<'a> {
129    stream: &'a ConnectionStream,
130    recv_buffer_size: Option<usize>,
131    send_buffer_size: Option<usize>,
132}
133
134impl<'a> ConnectionOptimizer<'a> {
135    /// Create a new connection optimizer with default buffer sizes
136    pub fn new(stream: &'a ConnectionStream) -> Self {
137        Self {
138            stream,
139            recv_buffer_size: None,
140            send_buffer_size: None,
141        }
142    }
143
144    /// Create a connection optimizer with custom buffer sizes
145    pub fn with_buffer_sizes(
146        stream: &'a ConnectionStream,
147        recv_size: usize,
148        send_size: usize,
149    ) -> Self {
150        Self {
151            stream,
152            recv_buffer_size: Some(recv_size),
153            send_buffer_size: Some(send_size),
154        }
155    }
156}
157
158impl<'a> NetworkOptimizer for ConnectionOptimizer<'a> {
159    fn optimize(&self) -> Result<(), io::Error> {
160        match (self.recv_buffer_size, self.send_buffer_size) {
161            (Some(recv), Some(send)) => {
162                // Use custom buffer sizes
163                match self.stream {
164                    ConnectionStream::Plain(tcp) => {
165                        let optimizer = TcpOptimizer::with_buffer_sizes(tcp, recv, send);
166                        debug!("Using {} with custom buffers", optimizer.description());
167                        optimizer.optimize()
168                    }
169                    ConnectionStream::Tls(tls) => {
170                        let optimizer = TlsOptimizer::with_buffer_sizes(tls.as_ref(), recv, send);
171                        debug!("Using {} with custom buffers", optimizer.description());
172                        optimizer.optimize()
173                    }
174                }
175            }
176            _ => {
177                // Use default buffer sizes
178                match self.stream {
179                    ConnectionStream::Plain(tcp) => {
180                        let optimizer = TcpOptimizer::new(tcp);
181                        debug!("Using {}", optimizer.description());
182                        optimizer.optimize()
183                    }
184                    ConnectionStream::Tls(tls) => {
185                        let optimizer = TlsOptimizer::new(tls.as_ref());
186                        debug!("Using {}", optimizer.description());
187                        optimizer.optimize()
188                    }
189                }
190            }
191        }
192    }
193
194    fn description(&self) -> &'static str {
195        match self.stream {
196            ConnectionStream::Plain(_) => "Connection-level TCP optimization",
197            ConnectionStream::Tls(_) => "Connection-level TLS optimization",
198        }
199    }
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205    use tokio::net::TcpListener;
206
207    #[tokio::test]
208    async fn test_tcp_optimizer() {
209        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
210        let addr = listener.local_addr().unwrap();
211
212        let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
213        let optimizer = TcpOptimizer::new(&stream);
214
215        assert_eq!(optimizer.description(), "TCP high-throughput optimization");
216
217        // Should not panic - actual socket optimization might fail in test environment
218        let _ = optimizer.optimize();
219    }
220
221    #[tokio::test]
222    async fn test_connection_optimizer_tcp() {
223        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
224        let addr = listener.local_addr().unwrap();
225
226        let tcp_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
227        let connection_stream = ConnectionStream::Plain(tcp_stream);
228        let optimizer = ConnectionOptimizer::new(&connection_stream);
229
230        assert_eq!(optimizer.description(), "Connection-level TCP optimization");
231
232        // Should not panic - actual socket optimization might fail in test environment
233        let _ = optimizer.optimize();
234    }
235
236    #[tokio::test]
237    async fn test_connection_optimizer_trait_usage() {
238        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
239        let addr = listener.local_addr().unwrap();
240
241        let tcp_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
242        let connection_stream = ConnectionStream::Plain(tcp_stream);
243
244        // Test that ConnectionOptimizer implements NetworkOptimizer trait
245        let optimizer: Box<dyn NetworkOptimizer> =
246            Box::new(ConnectionOptimizer::new(&connection_stream));
247
248        assert_eq!(optimizer.description(), "Connection-level TCP optimization");
249        let _ = optimizer.optimize();
250    }
251
252    #[tokio::test]
253    async fn test_connection_optimizer_with_custom_buffers() {
254        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
255        let addr = listener.local_addr().unwrap();
256
257        let tcp_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
258        let connection_stream = ConnectionStream::Plain(tcp_stream);
259        let optimizer = ConnectionOptimizer::with_buffer_sizes(&connection_stream, 4096, 8192);
260
261        assert_eq!(optimizer.description(), "Connection-level TCP optimization");
262        let _ = optimizer.optimize();
263    }
264
265    #[tokio::test]
266    async fn test_optimizer_creation() {
267        // Test that we can create optimizers with tokio streams
268        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
269        let addr = listener.local_addr().unwrap();
270
271        let tokio_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
272
273        let optimizer = TcpOptimizer::new(&tokio_stream);
274        assert_eq!(
275            optimizer.recv_buffer_size,
276            crate::constants::socket::HIGH_THROUGHPUT_RECV_BUFFER
277        );
278        assert_eq!(
279            optimizer.send_buffer_size,
280            crate::constants::socket::HIGH_THROUGHPUT_SEND_BUFFER
281        );
282    }
283
284    #[tokio::test]
285    async fn test_custom_buffer_sizes() {
286        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
287        let addr = listener.local_addr().unwrap();
288
289        let tokio_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
290
291        let optimizer = TcpOptimizer::with_buffer_sizes(&tokio_stream, 1024, 2048);
292        assert_eq!(optimizer.recv_buffer_size, 1024);
293        assert_eq!(optimizer.send_buffer_size, 2048);
294    }
295}
nntp_proxy/network/optimizers.rs

nntp_proxy/network/
optimizers.rs
