kreuzberg 4.0.0.pre.rc.29 → 4.0.0.rc1

data/vendor/kreuzberg/src/core/batch_optimizations.rs

@@ -1,385 +0,0 @@
-//! 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::core::config::ExtractionConfig;
-use crate::types::ExtractionResult;
-use crate::utils::pool::{ByteBufferPool, StringBufferPool, create_byte_buffer_pool, create_string_buffer_pool};
-use crate::utils::pool_sizing::PoolSizeHint;
-use crate::{KreuzbergError, Result};
-use parking_lot::Mutex;
-use std::path::Path;
-use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, Ordering};
-
-/// Configuration for batch processing with pooling optimizations.
-#[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.
-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 kreuzberg::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 kreuzberg::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
-    ///
-    /// ```rust,no_run
-    /// use kreuzberg::core::batch_optimizations::BatchProcessor;
-    /// use kreuzberg::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
-    }
-
-    /// Process multiple files with optimized pooling.
-    ///
-    /// This is a convenience method that combines file extraction with
-    /// automatic pool management.
-    ///
-    /// # Arguments
-    ///
-    /// * `paths` - Paths to the files to extract
-    /// * `extraction_config` - Extraction configuration
-    ///
-    /// # Returns
-    ///
-    /// A vector of `ExtractionResult` in the same order as input paths.
-    ///
-    /// # Errors
-    ///
-    /// Returns `KreuzbergError` if any file operation fails.
-    #[cfg(feature = "tokio-runtime")]
-    pub async fn process_files(
-        &self,
-        paths: Vec<impl AsRef<Path>>,
-        extraction_config: &ExtractionConfig,
-    ) -> Result<Vec<ExtractionResult>> {
-        use crate::core::extractor::batch_extract_file;
-
-        batch_extract_file(paths, extraction_config).await
-    }
-
-    /// Process multiple byte arrays with optimized pooling.
-    ///
-    /// This is a convenience method that combines bytes extraction with
-    /// automatic pool management.
-    ///
-    /// # Arguments
-    ///
-    /// * `contents` - Vector of (bytes, mime_type) tuples
-    /// * `extraction_config` - Extraction configuration
-    ///
-    /// # Returns
-    ///
-    /// A vector of `ExtractionResult` in the same order as input contents.
-    ///
-    /// # Errors
-    ///
-    /// Returns `KreuzbergError` if extraction fails.
-    #[cfg(feature = "tokio-runtime")]
-    pub async fn process_bytes(
-        &self,
-        contents: Vec<(&[u8], &str)>,
-        extraction_config: &ExtractionConfig,
-    ) -> Result<Vec<ExtractionResult>> {
-        use crate::core::extractor::batch_extract_bytes;
-
-        let owned_contents: Vec<(Vec<u8>, String)> = contents
-            .into_iter()
-            .map(|(bytes, mime)| (bytes.to_vec(), mime.to_string()))
-            .collect();
-
-        batch_extract_bytes(owned_contents, extraction_config).await
-    }
-
-    /// 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) -> Result<()> {
-        let pool_opt = self.string_pool.lock();
-        if let Some(pool) = pool_opt.as_ref() {
-            pool.clear()
-                .map_err(|e| KreuzbergError::Other(format!("string pool error: {}", e)))?;
-        }
-
-        let pool_opt = self.byte_pool.lock();
-        if let Some(pool) = pool_opt.as_ref() {
-            pool.clear()
-                .map_err(|e| KreuzbergError::Other(format!("byte pool error: {}", e)))?;
-        }
-
-        Ok(())
-    }
-}
-
-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().unwrap();
-            s.push_str("test");
-        }
-
-        {
-            let s = pool.acquire().unwrap();
-            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().unwrap();
-            buf.extend_from_slice(b"test");
-        }
-
-        {
-            let buf = pool.acquire().unwrap();
-            assert_eq!(buf.len(), 0);
-        }
-    }
-
-    #[test]
-    fn test_batch_processor_clear_pools() {
-        let processor = BatchProcessor::new();
-
-        let s1 = processor.string_pool().acquire().unwrap();
-        let s2 = processor.byte_pool().acquire().unwrap();
-
-        drop(s1);
-        drop(s2);
-
-        assert!(processor.string_pool_size() > 0);
-        assert!(processor.byte_pool_size() > 0);
-
-        processor.clear_pools().unwrap();
-
-        assert_eq!(processor.string_pool_size(), 0);
-        assert_eq!(processor.byte_pool_size(), 0);
-    }
-}