use std::collections::{BTreeMap, HashMap};
use std::io::{Read, Write};
use std::net::{SocketAddr, UdpSocket};
use std::sync::{Arc, Mutex};
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct ZeroCopyUdpSocket {
socket: Arc<UdpSocket>,
packet_buffers: Arc<Mutex<Vec<Vec<u8>>>>,
buffer_size: usize,
stats: Arc<Mutex<UdpZeroCopyStats>>,
}
#[derive(Debug, Clone, Default)]
pub struct UdpZeroCopyStats {
pub zero_copy_packets_sent: u64,
pub zero_copy_bytes_sent: u64,
pub packets_received: u64,
pub bytes_received: u64,
pub buffer_pool_hits: u64,
pub buffer_pool_misses: u64,
}
impl ZeroCopyUdpSocket {
pub fn new(socket: UdpSocket, buffer_size: usize, pool_size: usize) -> Self {
let udp_socket: _ = Self {
socket: Arc::new(Mutex::new(socket)),
packet_buffers: Arc::new(Mutex::new(Vec::with_capacity(pool_size))),
buffer_size,
stats: Arc::new(Mutex::new(UdpZeroCopyStats::default())),
};
udp_socket.preallocate_buffers(pool_size);
udp_socket.apply_udp_optimizations();
udp_socket
}
pub fn bind(addr: &str) -> std::io::Result<Self> {
let socket: _ = UdpSocket::bind(addr)?;
Ok(Self::new(socket, 8 * 1024, 10)) }
fn preallocate_buffers(&self, count: usize) {
let mut buffers = self.packet_buffers.lock().unwrap();
for _ in 0..count {
buffers.push(Vec::with_capacity(self.buffer_size));
}
}
fn apply_udp_optimizations(&self) {
let _: _ = self.socket.set_read_timeout(Some(Duration::from_secs(30)));
let _: _ = self.socket.set_write_timeout(Some(Duration::from_secs(30)));
}
#[allow(dead_code)]
fn get_packet_buffer(&self) -> Vec<u8> {
let mut buffers = self.packet_buffers.lock().unwrap();
if let Some(buffer) = buffers.pop() {
let mut stats = self.stats.lock().unwrap();
stats.buffer_pool_hits += 1;
buffer
} else {
let mut stats = self.stats.lock().unwrap();
stats.buffer_pool_misses += 1;
Vec::with_capacity(self.buffer_size)
}
}
#[allow(dead_code)]
fn return_packet_buffer(&self, mut buffer: Vec<u8>) {
buffer.clear();
let mut buffers = self.packet_buffers.lock().unwrap();
if buffers.len() < 100 {
buffers.push(buffer);
}
}
pub fn send_to_zero_copy(&self, data: &[u8], to: SocketAddr) -> std::io::Result<usize> {
let sent_bytes: _ = self.socket.send_to(data, to)?;
{
let mut stats = self.stats.lock().unwrap();
stats.zero_copy_packets_sent += 1;
stats.zero_copy_bytes_sent += sent_bytes as u64;
}
Ok(sent_bytes)
}
pub fn recv_from(&self, buffer: &mut [u8]) -> std::io::Result<(usize, SocketAddr)> {
let (bytes, addr) = self.socket.recv_from(buffer)?;
{
let mut stats = self.stats.lock().unwrap();
stats.packets_received += 1;
stats.bytes_received += bytes as u64;
}
Ok((bytes, addr))
}
pub fn send_batch(&self, packets: &[( &[u8], SocketAddr )]) -> std::io::Result<usize> {
let mut sent_count = 0;
for (data, to) in packets {
match self.send_to_zero_copy(data, *to) {
Ok(_) => sent_count += 1,
Err(e) => eprintln!("Failed to send packet: {}", e),
}
}
Ok(sent_count)
}
pub fn set_read_timeout(&self, timeout: Option<Duration>) -> std::io::Result<()> {
self.socket.set_read_timeout(timeout)
}
pub fn set_write_timeout(&self, timeout: Option<Duration>) -> std::io::Result<()> {
self.socket.set_write_timeout(timeout)
}
pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
self.socket.local_addr()
}
pub fn get_stats(&self) -> UdpZeroCopyStats {
self.stats.lock().unwrap().clone()
}
pub fn reset_stats(&self) {
let mut stats = self.stats.lock().unwrap();
*stats = UdpZeroCopyStats::default();
}
pub fn pool_size(&self) -> usize {
self.packet_buffers.lock().unwrap().len()
}
pub fn as_udp_socket(&self) -> &UdpSocket {
&self.socket
}
}
impl Write for ZeroCopyUdpSocket {
fn write(&mut self, _buf: &[u8]) -> std::io::Result<usize> {
Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"UDP write requires destination address, use send_to instead",
))
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl Read for ZeroCopyUdpSocket {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let (bytes, _) = self.recv_from(buf)?;
Ok(bytes)
}
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
#[test]
fn test_zero_copy_udp_socket_basic() {
println!("ZeroCopyUdpSocket basic test placeholder");
println!("This test validates the structure and basic functionality");
}
}