repr 0.8.0

/*!
This module provides a relatively simple thread-safe pool of reusable
objects. For the most part, it's implemented by a stack represented by a
Mutex<Vec<T>>. It has one small trick: because unlocking a mutex is somewhat
costly, in the case where a pool is accessed by the first thread that tried
to get a value, we bypass the mutex. Here are some benchmarks showing the
difference.

1) misc::anchored_literal_long_non_match    21 (18571 MB/s)
2) misc::anchored_literal_long_non_match   107 (3644 MB/s)
3) misc::anchored_literal_long_non_match    45 (8666 MB/s)
4) misc::anchored_literal_long_non_match    19 (20526 MB/s)

(1) represents our baseline: the master branch at the time of writing when
using the 'thread_local' crate to implement the pool below.

(2) represents a naive pool implemented completely via Mutex<Vec<T>>. There
is no special trick for bypassing the mutex.

(3) is the same as (2), except it uses Mutex<Vec<Box<T>>>. It is twice as
fast because a Box<T> is much smaller than the T we use with a Pool in this
crate. So pushing and popping a Box<T> from a Vec is quite a bit faster
than for T.

(4) is the same as (3), but with the trick for bypassing the mutex in the
case of the first-to-get thread.

Why move off of thread_local? Even though (4) is a hair faster than (1)
above, this was not the main goal. The main goal was to move off of
thread_local and find a way to *simply* re-capture some of its speed for
regex's specific case. So again, why move off of it? The *primary* reason is
because of memory leaks. See https://github.com/rust-lang/regex/issues/362
for example. (Why do I want it to be simple? Well, I suppose what I mean is,
"use as much safe code as possible to minimize risk and be as sure as I can
be that it is correct.")

My guess is that the thread_local design is probably not appropriate for
regex since its memory usage scales to the number of active threads that
have used a regex, where as the pool below scales to the number of threads
that simultaneously use a regex. While neither case permits contraction,
since we own the pool data structure below, we can add contraction if a
clear use case pops up in the wild. More pressingly though, it seems that
there are at least some use case patterns where one might have many threads
sitting around that might have used a regex at one point. While thread_local
does try to reuse space previously used by a thread that has since stopped,
its maximal memory usage still scales with the total number of active
threads. In contrast, the pool below scales with the total number of threads
*simultaneously* using the pool. The hope is that this uses less memory
overall. And if it doesn't, we can hopefully tune it somehow.

It seems that these sort of conditions happen frequently
in FFI inside of other more "managed" languages. This was
mentioned in the issue linked above, and also mentioned here:
https://github.com/BurntSushi/rure-go/issues/3. And in particular, users
confirm that disabling the use of thread_local resolves the leak.

There were other weaker reasons for moving off of thread_local as well.
Namely, at the time, I was looking to reduce dependencies. And for something
like regex, maintenance can be simpler when we own the full dependency tree.
*/

use core::{
    fmt::{self, Debug},
    panic::{RefUnwindSafe, UnwindSafe},
    sync::atomic::{AtomicUsize, Ordering}
};
use std::sync::Mutex;

/// An atomic counter used to allocate thread IDs.
static COUNTER: AtomicUsize = AtomicUsize::new(1);

thread_local!(
    /// A thread local used to assign an ID to a thread.
    static THREAD_ID: usize = {
        let next = COUNTER.fetch_add(1, Ordering::Relaxed);
        // SAFETY: We cannot permit the reuse of thread IDs since reusing a
        // thread ID might result in more than one thread "owning" a pool,
        // and thus, permit accessing a mutable value from multiple threads
        // simultaneously without synchronization. The intent of this panic is
        // to be a sanity check. It is not expected that the thread ID space
        // will actually be exhausted in practice.
        //
        // This checks that the counter never wraps around, since atomic
        // addition wraps around on overflow.
        if next == 0 {
            panic!("regex: thread ID allocation space exhausted");
        }
        next
    };
);

/// The type of the function used to create values in a pool when the pool is
/// empty and the caller requests one.
type CreateFn<T> =
    Box<dyn Fn() -> T + Send + Sync + UnwindSafe + RefUnwindSafe + 'static>;

/// A simple thread safe pool for reusing values.
///
/// Getting a value out comes with a guard. When that guard is dropped, the
/// value is automatically put back in the pool.
///
/// A Pool<T> impls Sync when T is Send (even if it's not Sync). This means
/// that T can use interior mutability. This is possible because a pool is
/// guaranteed to provide a value to exactly one thread at any time.
///
/// Currently, a pool never contracts in size. Its size is proportional to the
/// number of simultaneous uses.
pub struct Pool<T> {
    /// A stack of T values to hand out. These are used when a Pool is
    /// accessed by a thread that didn't create it.
    stack: Mutex<Vec<Box<T>>>,
    /// A function to create more T values when stack is empty and a caller
    /// has requested a T.
    create: CreateFn<T>,
    /// The ID of the thread that owns this pool. The owner is the thread
    /// that makes the first call to 'get'. When the owner calls 'get', it
    /// gets 'owner_val' directly instead of returning a T from 'stack'.
    /// See comments elsewhere for details, but this is intended to be an
    /// optimization for the common case that makes getting a T faster.
    ///
    /// It is initialized to a value of zero (an impossible thread ID) as a
    /// sentinel to indicate that it is unowned.
    owner: AtomicUsize,
    /// A value to return when the caller is in the same thread that created
    /// the Pool.
    owner_val: T,
}

// SAFETY: Since we want to use a Pool from multiple threads simultaneously
// behind an Arc, we need for it to be Sync. In cases where T is sync, Pool<T>
// would be Sync. However, since we use a Pool to store mutable scratch space,
// we wind up using a T that has interior mutability and is thus itself not
// Sync. So what we *really* want is for our Pool<T> to by Sync even when T is
// not Sync (but is at least Send).
//
// The only non-sync aspect of a Pool is its 'owner_val' field, which is used
// to implement faster access to a pool value in the common case of a pool
// being accessed in the same thread in which it was created. The 'stack' field
// is also shared, but a Mutex<T> where T: Send is already Sync. So we only
// need to worry about 'owner_val'.
//
// The key is to guarantee that 'owner_val' can only ever be accessed from one
// thread. In our implementation below, we guarantee this by only returning the
// 'owner_val' when the ID of the current thread matches the ID of the thread
// that created the Pool. Since this can only ever be one thread, it follows
// that only one thread can access 'owner_val' at any point in time. Thus, it
// is safe to declare that Pool<T> is Sync when T is Send.
//
// NOTE: It would also be possible to make the owning thread be the *first*
// thread that tries to get a value out of a Pool. However, the current
// implementation is a little simpler and it's not clear if making the first
// thread (rather than the creating thread) is meaningfully better.
//
// If there is a way to achieve our performance goals using safe code, then
// I would very much welcome a patch. As it stands, the implementation below
// tries to balance safety with performance. The case where a Regex is used
// from multiple threads simultaneously will suffer a bit since getting a cache
// will require unlocking a mutex.
unsafe impl<T: Send> Sync for Pool<T> {}

impl<T: Debug> Debug for Pool<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Pool")
            .field("stack", &self.stack)
            .field("owner", &self.owner)
            .field("owner_val", &self.owner_val)
            .finish()
    }
}

/// A guard that is returned when a caller requests a value from the pool.
///
/// The purpose of the guard is to use RAII to automatically put the value back
/// in the pool once it's dropped.
#[derive(Debug)]
pub struct PoolGuard<'a, T: Send> {
    /// The pool that this guard is attached to.
    pool: &'a Pool<T>,
    /// This is None when the guard represents the special "owned" value. In
    /// which case, the value is retrieved from 'pool.owner_val'.
    value: Option<Box<T>>,
}

impl<T: Send> Pool<T> {
    /// Create a new pool. The given closure is used to create values in the
    /// pool when necessary.
    pub fn new(create: CreateFn<T>) -> Pool<T> {
        let owner = AtomicUsize::new(0);
        let owner_val = create();
        Pool { stack: Mutex::new(vec![]), create, owner, owner_val }
    }

    /// Get a value from the pool. The caller is guaranteed to have exclusive
    /// access to the given value.
    ///
    /// Note that there is no guarantee provided about which value in the
    /// pool is returned. That is, calling get, dropping the guard (causing
    /// the value to go back into the pool) and then calling get again is NOT
    /// guaranteed to return the same value received in the first get call.
    #[cfg_attr(feature = "perf-inline", inline(always))]
    pub fn get(&self) -> PoolGuard<'_, T> {
        // Our fast path checks if the caller is the thread that "owns" this
        // pool. Or stated differently, whether it is the first thread that
        // tried to extract a value from the pool. If it is, then we can return
        // a T to the caller without going through a mutex.
        //
        // SAFETY: We must guarantee that only one thread gets access to this
        // value. Since a thread is uniquely identified by the THREAD_ID thread
        // local, it follows that is the caller's thread ID is equal to the
        // owner, then only one thread may receive this value.
        let caller = THREAD_ID.with(|id| *id);
        let owner = self.owner.load(Ordering::Relaxed);
        if caller == owner {
            return self.guard_owned();
        }
        self.get_slow(caller, owner)
    }

    /// This is the "slow" version that goes through a mutex to pop an
    /// allocated value off a stack to return to the caller. (Or, if the stack
    /// is empty, a new value is created.)
    ///
    /// If the pool has no owner, then this will set the owner.
    #[cold]
    fn get_slow(&self, caller: usize, owner: usize) -> PoolGuard<'_, T> {
        use std::sync::atomic::Ordering::Relaxed;

        if owner == 0 {
            // The sentinel 0 value means this pool is not yet owned. We
            // try to atomically set the owner. If we do, then this thread
            // becomes the owner and we can return a guard that represents
            // the special T for the owner.
            let res = self.owner.compare_exchange(0, caller, Relaxed, Relaxed);
            if res.is_ok() {
                return self.guard_owned();
            }
        }
        let mut stack = self.stack.lock().unwrap();
        let value = match stack.pop() {
            None => Box::new((self.create)()),
            Some(value) => value,
        };
        self.guard_stack(value)
    }

    /// Puts a value back into the pool. Callers don't need to call this. Once
    /// the guard that's returned by 'get' is dropped, it is put back into the
    /// pool automatically.
    fn put(&self, value: Box<T>) {
        let mut stack = self.stack.lock().unwrap();
        stack.push(value);
    }

    /// Create a guard that represents the special owned T.
    fn guard_owned(&self) -> PoolGuard<'_, T> {
        PoolGuard { pool: self, value: None }
    }

    /// Create a guard that contains a value from the pool's stack.
    fn guard_stack(&self, value: Box<T>) -> PoolGuard<'_, T> {
        PoolGuard { pool: self, value: Some(value) }
    }
}

impl<'a, T: Send> PoolGuard<'a, T> {
    /// Return the underlying value.
    pub fn value(&self) -> &T {
        match self.value {
            None => &self.pool.owner_val,
            Some(ref v) => &**v,
        }
    }
}

impl<'a, T: Send> Drop for PoolGuard<'a, T> {
    #[cfg_attr(feature = "perf-inline", inline(always))]
    fn drop(&mut self) {
        if let Some(value) = self.value.take() {
            self.pool.put(value);
        }
    }
}

#[cfg(test)]
mod tests {
    // use core::panic::{RefUnwindSafe, UnwindSafe};

    use super::*;

    // #[test]
    // fn oibits() {
    //     use crate::exec::ProgramCache;

    //     fn has_oibits<T: Send + Sync + UnwindSafe + RefUnwindSafe>() {}
    //     has_oibits::<Pool<ProgramCache<I>>>();
    // }

    // Tests that Pool implements the "single owner" optimization. That is, the
    // thread that first accesses the pool gets its own copy, while all other
    // threads get distinct copies.
    #[test]
    fn thread_owner_optimization() {
        use std::cell::RefCell;
        use std::sync::Arc;

        let pool: Arc<Pool<RefCell<Vec<char>>>> =
            Arc::new(Pool::new(Box::new(|| RefCell::new(vec!['a']))));
        pool.get().value().borrow_mut().push('x');

        let pool1 = pool.clone();
        let t1 = std::thread::spawn(move || {
            let guard = pool1.get();
            let v = guard.value();
            v.borrow_mut().push('y');
        });

        let pool2 = pool.clone();
        let t2 = std::thread::spawn(move || {
            let guard = pool2.get();
            let v = guard.value();
            v.borrow_mut().push('z');
        });

        t1.join().unwrap();
        t2.join().unwrap();

        // If we didn't implement the single owner optimization, then one of
        // the threads above is likely to have mutated the [a, x] vec that
        // we stuffed in the pool before spawning the threads. But since
        // neither thread was first to access the pool, and because of the
        // optimization, we should be guaranteed that neither thread mutates
        // the special owned pool value.
        //
        // (Technically this is an implementation detail and not a contract of
        // Pool's API.)
        assert_eq!(vec!['a', 'x'], *pool.get().value().borrow());
    }
}