1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

//! Efficient per-object thread-local storage implementation.
//!
//! ```rust
//! use std::thread;
//! use std::sync::Arc;
//! use std::cell::RefCell;
//! use per_thread_object::ThreadLocal;
//!
//! fn default() -> RefCell<u32> {
//!     RefCell::new(0x0)
//! }
//!
//! let tl: Arc<ThreadLocal<RefCell<u32>>> = Arc::new(ThreadLocal::new());
//! let tl2 = tl.clone();
//!
//! thread::spawn(move || {
//!     *tl2.get_or(default)
//!         .borrow_mut() += 1;
//!     assert_eq!(0x1, *tl2.get().unwrap().borrow());
//! })
//!     .join()
//!     .unwrap();
//!
//! *tl.get_or(default)
//!     .borrow_mut() += 2;
//! assert_eq!(0x2, *tl.get_or(default).borrow());
//! ```

#[cfg(not(feature = "loom"))]
mod loom;

#[cfg(feature = "loom")]
use loom;

mod thread;
mod page;

use std::ptr::NonNull;
use page::Storage;


/// Per-object thread-local storage
///
/// ## Capacity
///
/// `per-thread-object` has no max capacity limit,
/// each `ThreadLocal` instance will create its own memory space
/// instead of using global space.
///
/// this crate supports any number of threads,
/// but only the N threads are lock-free.
///
/// ## Panic when dropping
///
/// `ThreadLocal` will release object when calling `clean` or the end of thread.
/// If panic occurs during this process, it may cause a memory leak.
pub struct ThreadLocal<T: Send + 'static> {
    pool: Storage<T>
}

impl<T: Send + 'static> ThreadLocal<T> {
    pub fn new() -> ThreadLocal<T> {
        ThreadLocal {
            pool: Storage::new()
        }
    }

    #[inline]
    pub fn get(&self) -> Option<&T> {
        unsafe {
            self.pool.get(thread::get())
        }
    }

    #[inline]
    pub fn get_or<F: FnOnce() -> T>(&self, f: F) -> &T {
        use std::convert::Infallible;

        match self.get_or_try::<_, Infallible>(|| Ok(f())) {
            Ok(val) => val,
            Err(never) => match never {}
        }
    }

    #[inline]
    pub fn get_or_try<F, E>(&self, f: F) -> Result<&T, E>
    where
        F: FnOnce() -> Result<T, E>
    {
        let id = thread::get();
        let obj = unsafe { self.pool.get_or_new(id) };

        let val = match obj {
            Some(val) => val,
            None => {
                let ptr = NonNull::from(&*obj);
                let val = obj.get_or_insert(f()?);

                ThreadLocal::or_try(&self.pool, id, ptr);

                val
            }
        };

        Ok(val)
    }

    #[cold]
    fn or_try(pool: &Storage<T>, id: usize, ptr: NonNull<Option<T>>) {
        let thread_handle = unsafe {
            thread::push(pool.as_threads_ref(), ptr)
        };

        pool.insert_thread_handle(id, thread_handle);
    }
}

impl<T: Send + 'static> Default for ThreadLocal<T> {
    #[inline]
    fn default() -> ThreadLocal<T> {
        ThreadLocal::new()
    }
}

unsafe impl<T: Send> Send for ThreadLocal<T> {}
unsafe impl<T: Send> Sync for ThreadLocal<T> {}