use crate::network::{NetworkConfig, NetworkError};
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::Instant;
#[derive(Debug, Clone)]
pub struct BatchRequest {
pub id: u64,
pub data: Vec<u8>,
pub created_at: Instant,
pub priority: u8,
}
#[derive(Debug, Clone)]
pub struct BatchProcessorStats {
pub total_requests: u64,
pub batch_operations: usize,
pub average_batch_size: f64,
pub batch_latency_us: u64,
pub dropped_requests: u64,
}
pub struct BatchProcessor {
config: NetworkConfig,
requests: Arc<Mutex<Vec<BatchRequest>>>,
pending_count: Arc<Mutex<usize>>,
stats: Arc<Mutex<BatchProcessorStats>>,
next_request_id: Arc<Mutex<u64>>,
}
impl BatchProcessor {
pub fn new(config: NetworkConfig) -> Result<Self, NetworkError> {
Ok(Self {
config,
requests: Arc::new(Mutex::new(Vec::new())),
pending_count: Arc::new(Mutex::new(0)),
stats: Arc::new(Mutex::new(BatchProcessorStats {
total_requests: 0,
batch_operations: 0,
average_batch_size: 0.0,
batch_latency_us: 0,
dropped_requests: 0,
})),
next_request_id: Arc::new(Mutex::new(0)),
})
}
pub fn add_request(&mut self, data: Vec<u8>, priority: u8) -> Result<u64, NetworkError> {
let mut requests = self.requests.lock().unwrap();
let mut count = self.pending_count.lock().unwrap();
let mut next_id = self.next_request_id.lock().unwrap();
if requests.len() >= self.config.batch_size {
let mut stats = self.stats.lock().unwrap();
stats.dropped_requests += 1;
return Err(NetworkError::ResourceExhausted);
}
let request: _ = BatchRequest {
id: *next_id,
data,
priority,
created_at: Instant::now(),
};
*next_id += 1;
requests.push(request);
*count += 1;
Ok(*next_id - 1)
}
pub fn add_simple_request(&mut self, data: Vec<u8>) -> Result<u64, NetworkError> {
self.add_request(data, 128) }
pub fn process_batch(&mut self) -> Result<Vec<BatchRequest>, NetworkError> {
let mut requests = self.requests.lock().unwrap();
let mut count = self.pending_count.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
if requests.is_empty() {
return Ok(Vec::new());
}
requests.sort_by(|a, b| b.priority.cmp(&a.priority));
let batch_size: _ = std::cmp::min(requests.len(), self.config.batch_size);
let processed_requests: _ = requests.drain(0..batch_size).collect();
*count = requests.len();
stats.total_requests += batch_size as u64;
stats.batch_operations += 1;
stats.average_batch_size = stats.total_requests as f64 / stats.batch_operations as f64;
Ok(processed_requests)
}
pub fn pending_count(&self) -> usize {
*self.pending_count.lock().unwrap()
}
pub fn get_stats(&self) -> BatchProcessorStats {
self.stats.lock().unwrap().clone()
}
pub fn clear(&mut self) {
let mut requests = self.requests.lock().unwrap();
let mut count = self.pending_count.lock().unwrap();
requests.clear();
*count = 0;
}
pub fn should_process(&self) -> bool {
let count: _ = *self.pending_count.lock().unwrap();
count >= self.config.batch_size / 2 || count > 0 && self.is_timeout() }
fn is_timeout(&self) -> bool {
let requests: _ = self.requests.lock().unwrap();
if let Some(oldest) = requests.first() {
oldest.created_at.elapsed() >= self.config.batch_timeout
} else {
false
}
}
}
impl Default for BatchProcessor {
fn default() -> Self {
BatchProcessor::new(NetworkConfig::default()).unwrap()
}
}