| // Copyright 2017 Amanieu d'Antras |
| // |
| // Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or |
| // http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or |
| // http://opensource.org/licenses/MIT>, at your option. This file may not be |
| // copied, modified, or distributed except according to those terms. |
| |
| use crate::POINTER_WIDTH; |
| use once_cell::sync::Lazy; |
| use std::cmp::Reverse; |
| use std::collections::BinaryHeap; |
| use std::sync::Mutex; |
| use std::usize; |
| |
| /// Thread ID manager which allocates thread IDs. It attempts to aggressively |
| /// reuse thread IDs where possible to avoid cases where a ThreadLocal grows |
| /// indefinitely when it is used by many short-lived threads. |
| struct ThreadIdManager { |
| free_from: usize, |
| free_list: BinaryHeap<Reverse<usize>>, |
| } |
| impl ThreadIdManager { |
| fn new() -> ThreadIdManager { |
| ThreadIdManager { |
| free_from: 0, |
| free_list: BinaryHeap::new(), |
| } |
| } |
| fn alloc(&mut self) -> usize { |
| if let Some(id) = self.free_list.pop() { |
| id.0 |
| } else { |
| let id = self.free_from; |
| self.free_from = self |
| .free_from |
| .checked_add(1) |
| .expect("Ran out of thread IDs"); |
| id |
| } |
| } |
| fn free(&mut self, id: usize) { |
| self.free_list.push(Reverse(id)); |
| } |
| } |
| static THREAD_ID_MANAGER: Lazy<Mutex<ThreadIdManager>> = |
| Lazy::new(|| Mutex::new(ThreadIdManager::new())); |
| |
| /// Data which is unique to the current thread while it is running. |
| /// A thread ID may be reused after a thread exits. |
| #[derive(Clone, Copy)] |
| pub(crate) struct Thread { |
| /// The thread ID obtained from the thread ID manager. |
| pub(crate) id: usize, |
| /// The bucket this thread's local storage will be in. |
| pub(crate) bucket: usize, |
| /// The size of the bucket this thread's local storage will be in. |
| pub(crate) bucket_size: usize, |
| /// The index into the bucket this thread's local storage is in. |
| pub(crate) index: usize, |
| } |
| impl Thread { |
| fn new(id: usize) -> Thread { |
| let bucket = usize::from(POINTER_WIDTH) - id.leading_zeros() as usize; |
| let bucket_size = 1 << bucket.saturating_sub(1); |
| let index = if id != 0 { id ^ bucket_size } else { 0 }; |
| |
| Thread { |
| id, |
| bucket, |
| bucket_size, |
| index, |
| } |
| } |
| } |
| |
| /// Wrapper around `Thread` that allocates and deallocates the ID. |
| struct ThreadHolder(Thread); |
| impl ThreadHolder { |
| fn new() -> ThreadHolder { |
| ThreadHolder(Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc())) |
| } |
| } |
| impl Drop for ThreadHolder { |
| fn drop(&mut self) { |
| THREAD_ID_MANAGER.lock().unwrap().free(self.0.id); |
| } |
| } |
| |
| thread_local!(static THREAD_HOLDER: ThreadHolder = ThreadHolder::new()); |
| |
| /// Get the current thread. |
| pub(crate) fn get() -> Thread { |
| THREAD_HOLDER.with(|holder| holder.0) |
| } |
| |
| #[test] |
| fn test_thread() { |
| let thread = Thread::new(0); |
| assert_eq!(thread.id, 0); |
| assert_eq!(thread.bucket, 0); |
| assert_eq!(thread.bucket_size, 1); |
| assert_eq!(thread.index, 0); |
| |
| let thread = Thread::new(1); |
| assert_eq!(thread.id, 1); |
| assert_eq!(thread.bucket, 1); |
| assert_eq!(thread.bucket_size, 1); |
| assert_eq!(thread.index, 0); |
| |
| let thread = Thread::new(2); |
| assert_eq!(thread.id, 2); |
| assert_eq!(thread.bucket, 2); |
| assert_eq!(thread.bucket_size, 2); |
| assert_eq!(thread.index, 0); |
| |
| let thread = Thread::new(3); |
| assert_eq!(thread.id, 3); |
| assert_eq!(thread.bucket, 2); |
| assert_eq!(thread.bucket_size, 2); |
| assert_eq!(thread.index, 1); |
| |
| let thread = Thread::new(19); |
| assert_eq!(thread.id, 19); |
| assert_eq!(thread.bucket, 5); |
| assert_eq!(thread.bucket_size, 16); |
| assert_eq!(thread.index, 3); |
| } |