memlink_runtime/

batch.rs

//! Request batching for improved throughput.
//!
//! Combines multiple requests to the same module into batches
//! for more efficient processing.

use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use tokio::sync::{mpsc, Mutex};

/// Batch configuration.
#[derive(Debug, Clone)]
pub struct BatchConfig {
    /// Maximum batch size.
    pub max_size: usize,
    /// Maximum time to wait before flushing.
    pub max_wait: Duration,
    /// Minimum batch size to trigger flush.
    pub min_size: usize,
}

impl Default for BatchConfig {
    fn default() -> Self {
        Self {
            max_size: 100,
            max_wait: Duration::from_millis(50),
            min_size: 10,
        }
    }
}

/// A single batched request.
#[derive(Debug, Clone)]
pub struct BatchedRequest {
    /// Request arguments.
    pub args: Vec<u8>,
    /// Method name.
    pub method: String,
    /// Response sender.
    pub response_tx: mpsc::Sender<Result<Vec<u8>, String>>,
    /// Timestamp when queued.
    pub queued_at: Instant,
}

/// A batch of requests ready for execution.
#[derive(Debug)]
pub struct RequestBatch {
    /// Requests in the batch.
    pub requests: Vec<BatchedRequest>,
    /// Module name.
    pub module: String,
    /// Batch creation time.
    pub created_at: Instant,
}

impl RequestBatch {
    /// Creates a new batch.
    pub fn new(module: String) -> Self {
        Self {
            requests: Vec::new(),
            module,
            created_at: Instant::now(),
        }
    }

    /// Adds a request to the batch.
    pub fn push(&mut self, request: BatchedRequest) {
        self.requests.push(request);
    }

    /// Returns the batch size.
    pub fn len(&self) -> usize {
        self.requests.len()
    }

    /// Returns whether the batch is empty.
    pub fn is_empty(&self) -> bool {
        self.requests.is_empty()
    }
}

/// Batch statistics.
#[derive(Debug, Clone)]
pub struct BatchStats {
    pub total_batches: u64,
    pub total_requests: u64,
    pub avg_batch_size: f64,
    pub avg_latency_us: f64,
    pub current_queue_size: usize,
}

/// Request batcher for a module.
#[derive(Debug)]
pub struct RequestBatcher {
    /// Module name.
    module_name: String,
    /// Configuration.
    config: BatchConfig,
    /// Current pending requests.
    pending: Mutex<Vec<BatchedRequest>>,
    /// Batch sender.
    batch_tx: mpsc::Sender<RequestBatch>,
    /// Batch receiver.
    batch_rx: Mutex<Option<mpsc::Receiver<RequestBatch>>>,
    /// Total batches sent.
    total_batches: AtomicU64,
    /// Total requests batched.
    total_requests: AtomicU64,
}

impl RequestBatcher {
    /// Creates a new request batcher.
    pub fn new(module_name: String, config: BatchConfig) -> Self {
        let (batch_tx, batch_rx) = mpsc::channel(100);

        Self {
            module_name,
            config,
            pending: Mutex::new(Vec::new()),
            batch_tx,
            batch_rx: Mutex::new(Some(batch_rx)),
            total_batches: AtomicU64::new(0),
            total_requests: AtomicU64::new(0),
        }
    }

    /// Queues a request for batching.
    pub async fn queue(&self, request: BatchedRequest) -> Result<(), String> {
        let mut pending = self.pending.lock().await;
        pending.push(request);

        let should_flush = pending.len() >= self.config.max_size;
        drop(pending);

        if should_flush {
            self.flush().await;
        }

        Ok(())
    }

    /// Flushes pending requests into a batch.
    pub async fn flush(&self) {
        let mut pending = self.pending.lock().await;

        if pending.is_empty() {
            return;
        }

        let mut batch = RequestBatch::new(self.module_name.clone());
        let requests: Vec<BatchedRequest> = std::mem::take(&mut *pending);

        for req in requests {
            batch.push(req);
        }

        let batch_size = batch.len();
        let _ = self.batch_tx.send(batch).await;

        self.total_batches.fetch_add(1, Ordering::Relaxed);
        self.total_requests.fetch_add(batch_size as u64, Ordering::Relaxed);
    }

    /// Returns the batch receiver for consumption.
    pub fn take_receiver(&self) -> Option<mpsc::Receiver<RequestBatch>> {
        self.batch_rx.try_lock().ok().and_then(|mut rx| rx.take())
    }

    /// Returns statistics.
    pub fn stats(&self) -> BatchStats {
        let total_batches = self.total_batches.load(Ordering::Relaxed);
        let total_requests = self.total_requests.load(Ordering::Relaxed);

        BatchStats {
            total_batches,
            total_requests,
            avg_batch_size: if total_batches > 0 {
                total_requests as f64 / total_batches as f64
            } else {
                0.0
            },
            avg_latency_us: 0.0, // Would need tracking
            current_queue_size: 0, // Would need to read pending
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_batch_config() {
        let config = BatchConfig::default();
        assert_eq!(config.max_size, 100);
        assert_eq!(config.max_wait.as_millis(), 50);
        assert_eq!(config.min_size, 10);
    }

    #[test]
    fn test_batch_creation() {
        let batch = RequestBatch::new("test".to_string());
        assert_eq!(batch.len(), 0);
        assert!(batch.is_empty());
    }
}