use memmap2::{Mmap, MmapOptions};
use std::sync::{Arc, Mutex, RwLock};
use super::{NetworkConfig, NetworkStats};
use std::time::Instant;
use std::net::{SocketAddr, TcpListener};
#[derive(Debug, Clone)]
pub struct ZeroCopyConfig {
pub mmap_size: usize,
pub prefetch_threshold: usize,
pub zero_copy_threshold: usize,
}
impl Default for ZeroCopyConfig {
fn default() -> Self {
Self {
mmap_size: 1024 * 1024, prefetch_threshold: 64 * 1024, zero_copy_threshold: 1024, }
}
}
#[derive(Debug, Clone)]
pub struct NetworkZeroCopyStats {
pub mmap_operations: u64,
pub zero_copy_sends: u64,
pub zero_copy_receives: u64,
pub prefetch_operations: u64,
pub average_send_latency_ns: u64,
pub average_receive_latency_ns: u64,
}
impl Default for NetworkZeroCopyStats {
fn default() -> Self {
Self {
mmap_operations: 0,
zero_copy_sends: 0,
zero_copy_receives: 0,
prefetch_operations: 0,
average_send_latency_ns: 0,
average_receive_latency_ns: 0,
}
}
}
pub struct ZeroCopySocket {
config: NetworkConfig,
zero_copy_config: ZeroCopyConfig,
stats: Arc<RwLock<NetworkZeroCopyStats>>,
mmap_pool: Arc<Mutex<Vec<Mmap>>>,
}
impl ZeroCopySocket {
pub fn new(config: NetworkConfig) -> Self {
Self {
zero_copy_config: ZeroCopyConfig::default(),
stats: Arc::new(Mutex::new(NetworkZeroCopyStats::default())),
mmap_pool: Arc::new(Mutex::new(Vec::new())),
config,
}
}
pub async fn bind(addr: &SocketAddr) -> Result<ZeroCopyListener> {
let listener: _ = TcpListener::bind(addr).await?;
Ok(ZeroCopyListener {
listener,
config: NetworkConfig::default(),
stats: Arc::new(Mutex::new(NetworkZeroCopyStats::default())),
})
}
pub async fn send_zero_copy(&self, stream: &mut TcpStream, data: &[u8]) -> Result<usize> {
if data.len() < self.zero_copy_config.zero_copy_threshold {
return stream.write(data).await;
}
let start: _ = std::time::Instant::now();
let mmap: _ = self.create_mmap(data.len())?;
unsafe {
std::ptr::copy_nonoverlapping(
data.as_ptr(),
mmap.as_ptr() as *mut u8,
data.len()
);
}
let result: _ = stream.write(data).await;
let mut stats = self.stats.write().await;
stats.zero_copy_sends += 1;
stats.average_send_latency_ns = (stats.average_send_latency_ns + start.elapsed().as_nanos() as u64) / 2;
result
}
pub async fn recv_zero_copy(&self, stream: &mut TcpStream, buf: &mut [u8]) -> Result<usize> {
let start: _ = std::time::Instant::now();
let mmap: _ = self.allocate_mmap(buf.len())?;
let result: _ = stream.read(buf).await;
let mut stats = self.stats.write().await;
stats.zero_copy_receives += 1;
stats.average_receive_latency_ns = (stats.average_receive_latency_ns + start.elapsed().as_nanos() as u64) / 2;
result
}
fn create_mmap(&self, size: usize) -> Result<Mmap> {
let temp_file: _ = std::env::temp_dir().join(format!("amberjs_zero_copy_{}", std::process::id()));
let file: _ = std::fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(&temp_file)?;
file.set_len(size as u64)?;
unsafe {
let mmap: _ = MmapOptions::new()
.map(&file)?;
Ok(mmap)
}
}
fn allocate_mmap(&self, size: usize) -> Result<Mmap> {
self.create_mmap(size)
}
pub async fn get_stats(&self) -> NetworkZeroCopyStats {
self.stats.read().await.clone()
}
}
pub struct ZeroCopyListener {
listener: TcpListener,
config: NetworkConfig,
stats: Arc<RwLock<NetworkZeroCopyStats>>,
}
impl ZeroCopyListener {
pub async fn accept(&self) -> Result<(ZeroCopyStream, SocketAddr)> {
let (stream, addr) = self.listener.accept().await?;
Ok((ZeroCopyStream { stream }, addr))
}
pub async fn get_stats(&self) -> NetworkZeroCopyStats {
self.stats.read().await.clone()
}
}
pub struct ZeroCopyStream {
stream: tokio::net::TcpStream,
}
impl ZeroCopyStream {
pub async fn send(&mut self, data: &[u8]) -> Result<usize> {
self.stream.write(data).await
}
pub async fn recv(&mut self, buf: &mut [u8]) -> Result<usize> {
self.stream.read(buf).await
}
}
impl Drop for ZeroCopySocket {
fn drop(&mut self) {
let temp_dir: _ = std::env::temp_dir();
let temp_file: _ = temp_dir.join(format!("amberjs_zero_copy_{}", std::process::id()));
let _: _ = std::fs::remove_file(temp_file);
}
}