xberg 1.0.5

High-performance document intelligence library for Rust. Extract text, metadata, and structured data from PDFs, Office documents, images, and 98 formats and 306 programming languages via tree-sitter code intelligence with async/sync APIs.
Documentation
//! Batch extraction optimizations using object pooling.
//!
//! This module provides optimized batch processing utilities that leverage
//! object pooling to reduce allocations during concurrent extraction of
//! multiple documents.
//!
//! # Performance Impact
//!
//! - Reuses temporary string/buffer allocations across documents
//! - Reduces garbage collection pressure by ~5-10%
//! - Overall throughput improvement of 5-10% for batch operations
//!
//! # Usage
//!
//! The batch extraction functions automatically use pooling internally.
//! For manual control, use `BatchProcessor` to create pools and manage
//! extraction with custom pool sizes.

use crate::utils::pool::{ByteBufferPool, StringBufferPool, create_byte_buffer_pool, create_string_buffer_pool};
use crate::utils::pool_sizing::PoolSizeHint;
use parking_lot::Mutex;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

/// Configuration for batch processing with pooling optimizations.
#[cfg_attr(alef, alef(skip))]
#[derive(Debug, Clone)]
pub struct BatchProcessorConfig {
    /// Maximum number of string buffers to maintain in the pool
    pub string_pool_size: usize,

    /// Initial capacity for pooled string buffers in bytes
    pub string_buffer_capacity: usize,

    /// Maximum number of byte buffers to maintain in the pool
    pub byte_pool_size: usize,

    /// Initial capacity for pooled byte buffers in bytes
    pub byte_buffer_capacity: usize,

    /// Maximum concurrent extractions (for concurrency control)
    pub max_concurrent: Option<usize>,
}

impl Default for BatchProcessorConfig {
    fn default() -> Self {
        BatchProcessorConfig {
            string_pool_size: 10,
            string_buffer_capacity: 8192,
            byte_pool_size: 10,
            byte_buffer_capacity: 65536,
            max_concurrent: None,
        }
    }
}

/// Batch processor that manages object pools for optimized extraction.
///
/// This struct manages the lifecycle of reusable object pools used during
/// batch extraction. Pools are created lazily on first use and reused across
/// all documents processed by this batch processor.
///
/// # Lazy Initialization
///
/// Pools are initialized on demand to reduce memory usage for applications
/// that may not use batch processing immediately or at all.
#[cfg_attr(alef, alef(skip))]
pub struct BatchProcessor {
    string_pool: Mutex<Option<Arc<StringBufferPool>>>,
    byte_pool: Mutex<Option<Arc<ByteBufferPool>>>,
    config: BatchProcessorConfig,
    string_pool_initialized: AtomicBool,
    byte_pool_initialized: AtomicBool,
}

impl BatchProcessor {
    /// Create a new batch processor with default pool configuration.
    ///
    /// # Returns
    ///
    /// A new `BatchProcessor` ready to process documents.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use xberg::core::batch_optimizations::BatchProcessor;
    ///
    /// let processor = BatchProcessor::new();
    /// ```
    pub fn new() -> Self {
        Self::with_config(BatchProcessorConfig::default())
    }

    /// Create a new batch processor with custom pool configuration.
    ///
    /// Pools are not created immediately but lazily on first access.
    ///
    /// # Arguments
    ///
    /// * `config` - Custom batch processor configuration
    ///
    /// # Returns
    ///
    /// A new `BatchProcessor` configured with the provided settings.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use xberg::core::batch_optimizations::{BatchProcessor, BatchProcessorConfig};
    ///
    /// let mut config = BatchProcessorConfig::default();
    /// config.string_pool_size = 20;
    /// config.string_buffer_capacity = 16384;
    /// let processor = BatchProcessor::with_config(config);
    /// ```
    pub fn with_config(config: BatchProcessorConfig) -> Self {
        BatchProcessor {
            string_pool: Mutex::new(None),
            byte_pool: Mutex::new(None),
            config,
            string_pool_initialized: AtomicBool::new(false),
            byte_pool_initialized: AtomicBool::new(false),
        }
    }

    /// Create a batch processor with pool sizes optimized for a specific document.
    ///
    /// This method uses a `PoolSizeHint` (derived from file size and MIME type)
    /// to create a batch processor with appropriately sized pools. This reduces
    /// memory waste by tailoring pool allocation to actual document complexity.
    ///
    /// # Arguments
    ///
    /// * `hint` - Pool sizing hint containing recommended buffer counts and capacities
    ///
    /// # Returns
    ///
    /// A new `BatchProcessor` configured with the hint-based pool sizes
    ///
    /// # Example
    ///
    /// ```ignore
    /// use xberg::core::batch_optimizations::BatchProcessor;
    /// use xberg::utils::pool_sizing::estimate_pool_size;
    ///
    /// let hint = estimate_pool_size(5_000_000, "application/pdf");
    /// let processor = BatchProcessor::with_pool_hint(&hint);
    /// ```
    pub fn with_pool_hint(hint: &PoolSizeHint) -> Self {
        let config = BatchProcessorConfig {
            string_pool_size: hint.string_buffer_count,
            string_buffer_capacity: hint.string_buffer_capacity,
            byte_pool_size: hint.byte_buffer_count,
            byte_buffer_capacity: hint.byte_buffer_capacity,
            max_concurrent: None,
        };
        Self::with_config(config)
    }

    /// Get a reference to the string buffer pool.
    ///
    /// Creates the pool lazily on first access.
    /// Useful for custom pooling implementations that need direct pool access.
    pub fn string_pool(&self) -> Arc<StringBufferPool> {
        if self.string_pool_initialized.load(Ordering::Acquire) {
            return Arc::clone(self.string_pool.lock().as_ref().unwrap());
        }

        let mut pool_opt = self.string_pool.lock();
        if pool_opt.is_none() {
            let pool = Arc::new(create_string_buffer_pool(
                self.config.string_pool_size,
                self.config.string_buffer_capacity,
            ));
            *pool_opt = Some(pool);
            self.string_pool_initialized.store(true, Ordering::Release);
        }

        Arc::clone(pool_opt.as_ref().unwrap())
    }

    /// Get a reference to the byte buffer pool.
    ///
    /// Creates the pool lazily on first access.
    /// Useful for custom pooling implementations that need direct pool access.
    pub fn byte_pool(&self) -> Arc<ByteBufferPool> {
        if self.byte_pool_initialized.load(Ordering::Acquire) {
            return Arc::clone(self.byte_pool.lock().as_ref().unwrap());
        }

        let mut pool_opt = self.byte_pool.lock();
        if pool_opt.is_none() {
            let pool = Arc::new(create_byte_buffer_pool(
                self.config.byte_pool_size,
                self.config.byte_buffer_capacity,
            ));
            *pool_opt = Some(pool);
            self.byte_pool_initialized.store(true, Ordering::Release);
        }

        Arc::clone(pool_opt.as_ref().unwrap())
    }

    /// Get the current configuration.
    pub fn config(&self) -> &BatchProcessorConfig {
        &self.config
    }

    /// Get the number of pooled string buffers currently available.
    pub fn string_pool_size(&self) -> usize {
        self.string_pool.lock().as_ref().map(|p| p.size()).unwrap_or(0)
    }

    /// Get the number of pooled byte buffers currently available.
    pub fn byte_pool_size(&self) -> usize {
        self.byte_pool.lock().as_ref().map(|p| p.size()).unwrap_or(0)
    }

    /// Clear all pooled objects, forcing new allocations on next acquire.
    ///
    /// Useful for memory-constrained environments or to reclaim memory
    /// after processing large batches.
    pub fn clear_pools(&self) {
        if let Some(pool) = self.string_pool.lock().as_ref() {
            pool.clear();
        }
        if let Some(pool) = self.byte_pool.lock().as_ref() {
            pool.clear();
        }
    }
}

impl Default for BatchProcessor {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_batch_processor_creation() {
        let processor = BatchProcessor::new();
        assert_eq!(processor.string_pool_size(), 0);
        assert_eq!(processor.byte_pool_size(), 0);
    }

    #[test]
    fn test_batch_processor_with_config() {
        let config = BatchProcessorConfig {
            string_pool_size: 5,
            string_buffer_capacity: 1024,
            byte_pool_size: 3,
            byte_buffer_capacity: 4096,
            max_concurrent: None,
        };

        let processor = BatchProcessor::with_config(config);
        assert_eq!(processor.config().string_pool_size, 5);
        assert_eq!(processor.config().byte_pool_size, 3);
    }

    #[test]
    fn test_batch_processor_string_pool_usage() {
        let processor = BatchProcessor::new();
        let pool = processor.string_pool();

        {
            let mut s = pool.acquire();
            s.push_str("test");
        }

        {
            let s = pool.acquire();
            assert_eq!(s.len(), 0);
        }
    }

    #[test]
    fn test_batch_processor_byte_pool_usage() {
        let processor = BatchProcessor::new();
        let pool = processor.byte_pool();

        {
            let mut buf = pool.acquire();
            buf.extend_from_slice(b"test");
        }

        {
            let buf = pool.acquire();
            assert_eq!(buf.len(), 0);
        }
    }

    #[test]
    fn test_batch_processor_clear_pools() {
        let processor = BatchProcessor::new();

        let s1 = processor.string_pool().acquire();
        let s2 = processor.byte_pool().acquire();

        drop(s1);
        drop(s2);

        assert!(processor.string_pool_size() > 0);
        assert!(processor.byte_pool_size() > 0);

        processor.clear_pools();

        assert_eq!(processor.string_pool_size(), 0);
        assert_eq!(processor.byte_pool_size(), 0);
    }
}