automata: make MAX_POOL_STACKS user configurable

This exposes a knob inside of `regex-automata` (but not `regex`) that
permits end users to control how many cache lines are used inside of the
cache pool. Previously, this was always fixed to `8`, but using the
benchmark in #934 shows that this can cause performance to suffer when
the number of threads is greater than `8`.

In this commit, we not only expose it as user configurable, but we also
change the default to be `std::thread::available_parallelism()`. This
does increase memory usage for environments with higher core counts, but
it's still probably the better default. If this proves burdensome, we
can always roll the default back to `8`.

Fixes #1241

Author: BurntSushi

regex-automata/src/meta/regex.rs

@@ -1919,7 +1919,10 @@ impl Clone for Regex {
         let pool = {
             let strat = Arc::clone(&imp.strat);
             let create: CachePoolFn = Box::new(move || strat.create_cache());
-            Pool::new(create)
+            Pool::with_capacity(
+                self.imp.info.config().get_pool_capacity(),
+                create,
+            )
         };
         Regex { imp, pool }
     }
@@ -2483,6 +2486,7 @@ pub struct Config {
     backtrack: Option<bool>,
     byte_classes: Option<bool>,
     line_terminator: Option<u8>,
+    pool_capacity: Option<usize>,
 }
 
 impl Config {
@@ -3039,6 +3043,17 @@ impl Config {
         Config { line_terminator: Some(byte), ..self }
     }
 
+    /// Sets the capacity used to manage a pool of [`Cache`] values in the
+    /// higher level convenience APIs.
+    ///
+    /// When not configured explicitly, a reasonable default is selected. It
+    /// is rarely expected that a number large than the number of logical CPUs
+    /// makes sense as a value. A smaller number could result in slowdowns if
+    /// many regex queries are run under contention.
+    pub fn pool_capacity(self, capacity: usize) -> Config {
+        Config { pool_capacity: Some(capacity), ..self }
+    }
+
     /// Toggle whether the hybrid NFA/DFA (also known as the "lazy DFA") should
     /// be available for use by the meta regex engine.
     ///
@@ -3210,6 +3225,30 @@ impl Config {
         self.line_terminator.unwrap_or(b'\n')
     }
 
+    /// Returns the configured pool capacity, as set by
+    /// [`Config::pool_capacity`].
+    ///
+    /// If it was not explicitly set, then a default value is returned.
+    pub fn get_pool_capacity(&self) -> usize {
+        // The default is an empirically chosen value that balances memory
+        // usage with runtime performance. In practice, with `std` enabled,
+        // we choose a value that matches the total number of CPUs.
+        const DEFAULT_POOL_CAPACITY: usize = 8;
+
+        self.pool_capacity.unwrap_or_else(|| {
+            #[cfg(feature = "std")]
+            {
+                std::thread::available_parallelism()
+                    .map(|n| n.get())
+                    .unwrap_or(DEFAULT_POOL_CAPACITY)
+            }
+            #[cfg(not(feature = "std"))]
+            {
+                DEFAULT_POOL_CAPACITY
+            }
+        })
+    }
+
     /// Returns whether the hybrid NFA/DFA regex engine may be used, as set by
     /// [`Config::hybrid`].
     ///
@@ -3317,6 +3356,7 @@ impl Config {
             backtrack: o.backtrack.or(self.backtrack),
             byte_classes: o.byte_classes.or(self.byte_classes),
             line_terminator: o.line_terminator.or(self.line_terminator),
+            pool_capacity: o.pool_capacity.or(self.pool_capacity),
         }
     }
 }
@@ -3641,7 +3681,7 @@ impl Builder {
         let pool = {
             let strat = Arc::clone(&strat);
             let create: CachePoolFn = Box::new(move || strat.create_cache());
-            Pool::new(create)
+            Pool::with_capacity(self.config.get_pool_capacity(), create)
         };
         Ok(Regex { imp: Arc::new(RegexI { strat, info }), pool })
     }

regex-automata/src/util/pool.rs

@@ -159,6 +159,50 @@ impl<T, F> Pool<T, F> {
     pub fn new(create: F) -> Pool<T, F> {
         Pool(alloc::boxed::Box::new(inner::Pool::new(create)))
     }
+
+    /// Create a new pool. The given closure is used to create values in
+    /// the pool when necessary.
+    ///
+    /// When the `std` feature is enabled, a `Pool` is thread-aware and spreads
+    /// its memory out across multiple cache lines. The number of cache lines
+    /// is determined by the `capacity` parameter passed here. By default, a
+    /// fixed reasonable number is used. A smaller number means less memory is
+    /// used, but a higher number means there may be less contention on this
+    /// pool in highly threaded environments doing a lot of searches using the
+    /// same `Regex` value.
+    ///
+    /// When `std` is not enabled, then the capacity parameter is ignored
+    /// because the underlying pool implementation is not thread-aware.
+    ///
+    /// The capacity must be at least 1. If it's less than 1, then it is
+    /// forced to be 1.
+    pub fn with_capacity(capacity: usize, create: F) -> Pool<T, F> {
+        Pool(alloc::boxed::Box::new(inner::Pool::with_capacity(
+            capacity, create,
+        )))
+    }
+
+    /// Create a new pool. The given closure is used to create values in
+    /// the pool when necessary.
+    ///
+    /// This is a convenience routine for calling `Pool::with_capacity` with
+    /// a number equivalent to the available parallelism for this environment.
+    ///
+    /// If `std` is not enabled or if the query for available parallelism
+    /// failed, then this is equivalent to calling `Pool::new`.
+    pub fn with_available_parallelism_capacity(create: F) -> Pool<T, F> {
+        #[cfg(feature = "std")]
+        {
+            let Ok(n) = std::thread::available_parallelism() else {
+                return Pool::new(create);
+            };
+            Pool::with_capacity(n.get(), create)
+        }
+        #[cfg(not(feature = "std"))]
+        {
+            Pool::new(create)
+        }
+    }
 }
 
 impl<T: Send, F: Fn() -> T> Pool<T, F> {
@@ -455,6 +499,18 @@ mod inner {
         /// Create a new pool. The given closure is used to create values in
         /// the pool when necessary.
         pub(super) fn new(create: F) -> Pool<T, F> {
+            Pool::with_capacity(MAX_POOL_STACKS, create)
+        }
+
+        /// Create a new pool. The given closure is used to create values in
+        /// the pool when necessary.
+        ///
+        /// The given capacity is used to determine how many cache lines to
+        /// maintain. Each cache line contains a stack of cached entries.
+        ///
+        /// The capacity must be at least 1. If it's less than 1, then it is
+        /// forced to be 1.
+        pub(super) fn with_capacity(capacity: usize, create: F) -> Pool<T, F> {
             // FIXME: Now that we require 1.65+, Mutex::new is available as
             // const... So we can almost mark this function as const. But of
             // course, we're creating a Vec of stacks below (we didn't when I
@@ -493,7 +549,7 @@ mod inner {
             // Back to square one. I maybe we just don't make a pool's
             // constructor const and live with it. It's probably not a huge
             // deal.
-            let mut stacks = Vec::with_capacity(MAX_POOL_STACKS);
+            let mut stacks = Vec::with_capacity(capacity.max(1));
             for _ in 0..stacks.capacity() {
                 stacks.push(CacheLine(Mutex::new(vec![])));
             }
@@ -847,6 +903,14 @@ mod inner {
         pub(super) const fn new(create: F) -> Pool<T, F> {
             Pool { stack: Mutex::new(vec![]), create }
         }
+
+        /// This is a no-op since this pool implementation isn't thread-aware.
+        pub(super) const fn with_capacity(
+            _capacity: usize,
+            create: F,
+        ) -> Pool<T, F> {
+            Pool::new(create)
+        }
     }
 
     impl<T: Send, F: Fn() -> T> Pool<T, F> {