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

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

@@ -1,503 +0,0 @@
-//! 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);
-    }
-}