kreuzberg 4.0.0.rc2 → 4.0.1

data/vendor/kreuzberg/src/utils/mod.rs

@@ -1,17 +1,31 @@
-//! Text utility functions for quality processing and string manipulation.
-//!
-//! This module provides:
-//! - Quality processing: clean OCR artifacts, calculate quality scores
-//! - String utilities: safe decoding, mojibake fixing, encoding detection
-
-#[cfg(feature = "quality")]
-pub mod quality;
-
-#[cfg(feature = "quality")]
-pub mod string_utils;
-
-#[cfg(feature = "quality")]
-pub use quality::{calculate_quality_score, clean_extracted_text, normalize_spaces};
-
-#[cfg(feature = "quality")]
-pub use string_utils::{calculate_text_confidence, fix_mojibake, safe_decode};
+//! Text utility functions for quality processing and string manipulation.
+//!
+//! This module provides:
+//! - Quality processing: clean OCR artifacts, calculate quality scores
+//! - String utilities: safe decoding, mojibake fixing, encoding detection
+//! - Object pooling: reusable pools for batch processing to reduce allocations
+
+#[cfg(feature = "quality")]
+pub mod quality;
+
+#[cfg(feature = "quality")]
+pub mod string_utils;
+
+pub mod pool;
+pub mod pool_sizing;
+pub mod string_pool;
+
+#[cfg(feature = "quality")]
+pub use quality::{calculate_quality_score, clean_extracted_text, normalize_spaces};
+
+#[cfg(feature = "quality")]
+pub use string_utils::{calculate_text_confidence, fix_mojibake, safe_decode};
+
+pub use pool::{
+    ByteBufferPool, Pool, PoolError, PoolGuard, Recyclable, StringBufferPool, create_byte_buffer_pool,
+    create_string_buffer_pool,
+};
+
+pub use pool_sizing::{PoolSizeHint, estimate_pool_size};
+
+pub use string_pool::{InternedString, intern_language_code, intern_mime_type};

data/vendor/kreuzberg/src/utils/pool.rs

@@ -0,0 +1,503 @@
+//! Object pooling utilities for reducing allocations in batch processing.
+//!
+//! This module provides reusable object pools to reduce garbage collection and allocator
+//! pressure during batch document extraction. Instead of creating/destroying temporary objects
+//! repeatedly, pools maintain a collection of pre-allocated objects that can be quickly reused.
+//!
+//! # Performance Benefits
+//!
+//! For batch extraction of N documents:
+//! - Without pooling: N allocations + N deallocations per run
+//! - With pooling: 1-2 allocations per run (pool creation), zero during document processing
+//!
+//! Expected improvement: 5-10% overall throughput improvement on batch operations.
+//!
+//! # Thread Safety
+//!
+//! All pools are `Send + Sync` and can be safely shared across threads using `Arc<Pool<T>>`.
+//! Internal state is protected by `Mutex` for proper synchronization.
+//!
+//! # Example
+//!
+//! ```rust,no_run
+//! use kreuzberg::utils::pool::StringBufferPool;
+//!
+//! let pool = StringBufferPool::new(|| String::with_capacity(4096), 10); // 10 buffers of 4KB each
+//! let mut buffer = pool.acquire().unwrap();
+//! buffer.push_str("some content");
+//! // buffer is returned to pool when dropped
+//! ```
+
+use parking_lot::Mutex;
+use std::sync::Arc;
+
+#[cfg(feature = "pool-metrics")]
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+/// Metrics tracking for pool allocations and reuse patterns.
+///
+/// These metrics help identify pool efficiency and allocation patterns.
+/// Only available when the `pool-metrics` feature is enabled.
+#[cfg(feature = "pool-metrics")]
+#[derive(Debug)]
+pub struct PoolMetrics {
+    /// Total number of acquire calls on this pool
+    pub total_acquires: AtomicUsize,
+    /// Total number of cache hits (reused objects from pool)
+    pub total_cache_hits: AtomicUsize,
+    /// Peak number of objects stored simultaneously in this pool
+    pub peak_items_stored: AtomicUsize,
+    /// Total number of objects created by the factory function
+    pub total_creations: AtomicUsize,
+}
+
+#[cfg(feature = "pool-metrics")]
+impl PoolMetrics {
+    /// Create a new metrics tracker with all counters at zero.
+    pub fn new() -> Self {
+        PoolMetrics {
+            total_acquires: AtomicUsize::new(0),
+            total_cache_hits: AtomicUsize::new(0),
+            peak_items_stored: AtomicUsize::new(0),
+            total_creations: AtomicUsize::new(0),
+        }
+    }
+
+    /// Calculate the cache hit rate as a percentage (0.0-100.0).
+    pub fn hit_rate(&self) -> f64 {
+        let acquires = self.total_acquires.load(Ordering::Relaxed);
+        if acquires == 0 {
+            return 0.0;
+        }
+        (self.total_cache_hits.load(Ordering::Relaxed) as f64 / acquires as f64) * 100.0
+    }
+
+    /// Get all metrics as a struct for reporting.
+    pub fn snapshot(&self) -> PoolMetricsSnapshot {
+        PoolMetricsSnapshot {
+            total_acquires: self.total_acquires.load(Ordering::Relaxed),
+            total_cache_hits: self.total_cache_hits.load(Ordering::Relaxed),
+            peak_items_stored: self.peak_items_stored.load(Ordering::Relaxed),
+            total_creations: self.total_creations.load(Ordering::Relaxed),
+        }
+    }
+
+    /// Reset all metrics to zero.
+    pub fn reset(&self) {
+        self.total_acquires.store(0, Ordering::Relaxed);
+        self.total_cache_hits.store(0, Ordering::Relaxed);
+        self.peak_items_stored.store(0, Ordering::Relaxed);
+        self.total_creations.store(0, Ordering::Relaxed);
+    }
+}
+
+#[cfg(feature = "pool-metrics")]
+impl Default for PoolMetrics {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(feature = "pool-metrics")]
+#[derive(Debug, Clone, Copy)]
+pub struct PoolMetricsSnapshot {
+    pub total_acquires: usize,
+    pub total_cache_hits: usize,
+    pub peak_items_stored: usize,
+    pub total_creations: usize,
+}
+
+/// A thread-safe object pool that reuses instances to reduce allocations.
+///
+/// Generic over any type that implements `Recyclable`, allowing pooling of
+/// different object types with custom reset logic.
+#[derive(Clone)]
+pub struct Pool<T: Recyclable> {
+    factory: Arc<dyn Fn() -> T + Send + Sync>,
+    objects: Arc<Mutex<Vec<T>>>,
+    max_size: usize,
+    #[cfg(feature = "pool-metrics")]
+    metrics: Arc<PoolMetrics>,
+}
+
+/// Trait for types that can be pooled and reused.
+///
+/// Implementing this trait allows a type to be used with `Pool<T>`.
+/// The `reset()` method should clear the object's state for reuse.
+pub trait Recyclable: Send + 'static {
+    /// Reset the object to a reusable state.
+    ///
+    /// This is called when returning an object to the pool.
+    /// Should clear any internal data while preserving capacity.
+    fn reset(&mut self);
+}
+
+impl<T: Recyclable> Pool<T> {
+    /// Create a new pool with a given factory and maximum size.
+    ///
+    /// # Arguments
+    ///
+    /// * `factory` - Function that creates new instances
+    /// * `max_size` - Maximum number of objects to keep in the pool
+    ///
+    /// # Returns
+    ///
+    /// A new pool ready to use.
+    ///
+    /// # Example
+    ///
+    /// ```rust,no_run
+    /// use kreuzberg::utils::pool::Pool;
+    ///
+    /// let pool = Pool::new(|| String::new(), 10);
+    /// ```
+    pub fn new<F>(factory: F, max_size: usize) -> Self
+    where
+        F: Fn() -> T + Send + Sync + 'static,
+    {
+        Pool {
+            factory: Arc::new(factory),
+            objects: Arc::new(Mutex::new(Vec::with_capacity(max_size))),
+            max_size,
+            #[cfg(feature = "pool-metrics")]
+            metrics: Arc::new(PoolMetrics::new()),
+        }
+    }
+
+    /// Acquire an object from the pool or create a new one if empty.
+    ///
+    /// # Returns
+    ///
+    /// A `PoolGuard<T>` that will return the object to the pool when dropped.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the mutex is already locked by the current thread (deadlock).
+    /// This is a safety mechanism provided by parking_lot to prevent subtle bugs.
+    pub fn acquire(&self) -> Result<PoolGuard<T>, PoolError> {
+        #[cfg(feature = "pool-metrics")]
+        self.metrics.total_acquires.fetch_add(1, Ordering::Relaxed);
+
+        let mut objects = self.objects.lock();
+
+        let object = if let Some(mut obj) = objects.pop() {
+            #[cfg(feature = "pool-metrics")]
+            self.metrics.total_cache_hits.fetch_add(1, Ordering::Relaxed);
+
+            obj.reset();
+            obj
+        } else {
+            #[cfg(feature = "pool-metrics")]
+            self.metrics.total_creations.fetch_add(1, Ordering::Relaxed);
+
+            (self.factory)()
+        };
+
+        Ok(PoolGuard {
+            object: Some(object),
+            pool: Pool {
+                factory: Arc::clone(&self.factory),
+                objects: Arc::clone(&self.objects),
+                max_size: self.max_size,
+                #[cfg(feature = "pool-metrics")]
+                metrics: Arc::clone(&self.metrics),
+            },
+        })
+    }
+
+    /// Get the current number of objects in the pool.
+    pub fn size(&self) -> usize {
+        self.objects.lock().len()
+    }
+
+    /// Clear the pool, discarding all pooled objects.
+    pub fn clear(&self) -> Result<(), PoolError> {
+        self.objects.lock().clear();
+        Ok(())
+    }
+
+    /// Get a reference to the pool metrics (only available with `pool-metrics` feature).
+    #[cfg(feature = "pool-metrics")]
+    pub fn metrics(&self) -> &PoolMetrics {
+        &self.metrics
+    }
+}
+
+/// RAII guard that returns an object to the pool when dropped.
+pub struct PoolGuard<T: Recyclable> {
+    object: Option<T>,
+    pool: Pool<T>,
+}
+
+impl<T: Recyclable> std::ops::Deref for PoolGuard<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        self.object.as_ref().expect("object should never be None")
+    }
+}
+
+impl<T: Recyclable> std::ops::DerefMut for PoolGuard<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.object.as_mut().expect("object should never be None")
+    }
+}
+
+impl<T: Recyclable> Drop for PoolGuard<T> {
+    fn drop(&mut self) {
+        if let Some(object) = self.object.take() {
+            let mut objects = self.pool.objects.lock();
+            if objects.len() < self.pool.max_size {
+                objects.push(object);
+
+                #[cfg(feature = "pool-metrics")]
+                {
+                    let current_size = objects.len();
+                    let peak = self.pool.metrics.peak_items_stored.load(Ordering::Relaxed);
+                    if current_size > peak {
+                        self.pool
+                            .metrics
+                            .peak_items_stored
+                            .store(current_size, Ordering::Relaxed);
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// Pooled string buffer for text accumulation.
+///
+/// Reuses a single String allocation across multiple uses to reduce allocations
+/// when processing multiple documents.
+impl Recyclable for String {
+    fn reset(&mut self) {
+        self.clear();
+    }
+}
+
+/// Pooled byte vector for binary data handling.
+///
+/// Reuses a Vec<u8> allocation across multiple uses.
+impl Recyclable for Vec<u8> {
+    fn reset(&mut self) {
+        self.clear();
+    }
+}
+
+/// Error type for pool operations.
+#[derive(Debug, Clone)]
+pub enum PoolError {
+    /// The pool's internal mutex was poisoned.
+    ///
+    /// This indicates a panic occurred while holding the lock.
+    /// The pool is in a locked state and cannot be recovered.
+    LockPoisoned,
+}
+
+impl std::fmt::Display for PoolError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PoolError::LockPoisoned => write!(f, "pool lock poisoned"),
+        }
+    }
+}
+
+impl std::error::Error for PoolError {}
+
+/// Convenience type alias for a pooled String.
+pub type StringBufferPool = Pool<String>;
+
+/// Convenience type alias for a pooled Vec<u8>.
+pub type ByteBufferPool = Pool<Vec<u8>>;
+
+/// Create a pre-configured string buffer pool for batch processing.
+///
+/// # Arguments
+///
+/// * `pool_size` - Maximum number of buffers to keep in the pool
+/// * `buffer_capacity` - Initial capacity for each buffer in bytes
+///
+/// # Returns
+///
+/// A pool configured for text accumulation with reasonable defaults.
+///
+/// # Example
+///
+/// ```rust,no_run
+/// use kreuzberg::utils::pool::create_string_buffer_pool;
+///
+/// let pool = create_string_buffer_pool(10, 8192);
+/// let mut buffer = pool.acquire().unwrap();
+/// buffer.push_str("content");
+/// ```
+pub fn create_string_buffer_pool(pool_size: usize, buffer_capacity: usize) -> StringBufferPool {
+    Pool::new(move || String::with_capacity(buffer_capacity), pool_size)
+}
+
+/// Create a pre-configured byte buffer pool for batch processing.
+///
+/// # Arguments
+///
+/// * `pool_size` - Maximum number of buffers to keep in the pool
+/// * `buffer_capacity` - Initial capacity for each buffer in bytes
+///
+/// # Returns
+///
+/// A pool configured for binary data handling with reasonable defaults.
+///
+/// # Example
+///
+/// ```rust,no_run
+/// use kreuzberg::utils::pool::create_byte_buffer_pool;
+///
+/// let pool = create_byte_buffer_pool(10, 65536);
+/// let mut buffer = pool.acquire().unwrap();
+/// buffer.extend_from_slice(b"binary data");
+/// ```
+pub fn create_byte_buffer_pool(pool_size: usize, buffer_capacity: usize) -> ByteBufferPool {
+    Pool::new(move || Vec::with_capacity(buffer_capacity), pool_size)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pool_acquire_and_reuse() {
+        let pool = Pool::new(String::new, 5);
+
+        {
+            let mut s1 = pool.acquire().unwrap();
+            s1.push_str("hello");
+            assert_eq!(s1.len(), 5);
+        }
+
+        {
+            let s2 = pool.acquire().unwrap();
+            assert_eq!(s2.len(), 0, "string should be cleared when reused");
+        }
+    }
+
+    #[test]
+    fn test_pool_respects_max_size() {
+        let pool = Pool::new(String::new, 2);
+
+        let guard1 = pool.acquire().unwrap();
+        let guard2 = pool.acquire().unwrap();
+        let guard3 = pool.acquire().unwrap();
+
+        drop(guard1);
+        drop(guard2);
+        drop(guard3);
+
+        assert!(pool.size() <= 2, "pool size should not exceed max_size");
+    }
+
+    #[test]
+    fn test_pool_clear() {
+        let pool = Pool::new(String::new, 5);
+
+        let _g1 = pool.acquire().unwrap();
+        let _g2 = pool.acquire().unwrap();
+
+        drop(_g1);
+        drop(_g2);
+
+        assert!(pool.size() > 0, "pool should have items");
+        pool.clear().unwrap();
+        assert_eq!(pool.size(), 0, "pool should be empty after clear");
+    }
+
+    #[test]
+    fn test_byte_buffer_pool() {
+        let pool = Pool::new(Vec::new, 5);
+
+        {
+            let mut buf = pool.acquire().unwrap();
+            buf.extend_from_slice(b"hello");
+            assert_eq!(buf.len(), 5);
+        }
+
+        {
+            let buf = pool.acquire().unwrap();
+            assert_eq!(buf.len(), 0, "buffer should be cleared");
+        }
+    }
+
+    #[test]
+    fn test_pool_deref() {
+        let pool = Pool::new(String::new, 5);
+        let mut guard = pool.acquire().unwrap();
+
+        let _len = guard.len();
+
+        guard.push_str("test");
+        assert_eq!(&*guard, "test");
+    }
+
+    #[test]
+    fn test_concurrent_access() {
+        use std::sync::Arc;
+        use std::thread;
+
+        let pool = Arc::new(Pool::new(String::new, 10));
+        let mut handles = vec![];
+
+        for i in 0..5 {
+            let pool_clone = Arc::clone(&pool);
+            let handle = thread::spawn(move || {
+                let mut buf = pool_clone.acquire().unwrap();
+                buf.push_str(&format!("thread_{}", i));
+                std::thread::sleep(std::time::Duration::from_millis(10));
+                drop(buf);
+            });
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.join().unwrap();
+        }
+
+        assert!(pool.size() <= 10);
+    }
+
+    #[test]
+    #[cfg(feature = "pool-metrics")]
+    fn test_pool_metrics_tracking() {
+        let pool = Pool::new(String::new, 5);
+
+        {
+            let _s1 = pool.acquire().unwrap();
+        }
+
+        {
+            let _s2 = pool.acquire().unwrap();
+        }
+
+        let metrics = pool.metrics();
+        assert_eq!(metrics.total_acquires.load(Ordering::Relaxed), 2);
+        assert_eq!(metrics.total_cache_hits.load(Ordering::Relaxed), 1);
+        assert_eq!(metrics.total_creations.load(Ordering::Relaxed), 1);
+
+        let hit_rate = metrics.hit_rate();
+        assert!(hit_rate > 49.0 && hit_rate < 51.0);
+    }
+
+    #[test]
+    #[cfg(feature = "pool-metrics")]
+    fn test_pool_metrics_peak_tracking() {
+        let pool = Pool::new(String::new, 5);
+
+        let g1 = pool.acquire().unwrap();
+        let g2 = pool.acquire().unwrap();
+        let g3 = pool.acquire().unwrap();
+
+        drop(g1);
+        drop(g2);
+        drop(g3);
+
+        let metrics = pool.metrics();
+        assert!(metrics.peak_items_stored.load(Ordering::Relaxed) >= 3);
+    }
+}