1 // Copyright 2017 Amanieu d'Antras
2 //
3 // Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4 // http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5 // http://opensource.org/licenses/MIT>, at your option. This file may not be
6 // copied, modified, or distributed except according to those terms.
7
8 use crate::POINTER_WIDTH;
9 use once_cell::sync::Lazy;
10 use std::cmp::Reverse;
11 use std::collections::BinaryHeap;
12 use std::sync::Mutex;
13 use std::usize;
14
15 /// Thread ID manager which allocates thread IDs. It attempts to aggressively
16 /// reuse thread IDs where possible to avoid cases where a ThreadLocal grows
17 /// indefinitely when it is used by many short-lived threads.
18 struct ThreadIdManager {
19 free_from: usize,
20 free_list: BinaryHeap<Reverse<usize>>,
21 }
22 impl ThreadIdManager {
new() -> ThreadIdManager23 fn new() -> ThreadIdManager {
24 ThreadIdManager {
25 free_from: 0,
26 free_list: BinaryHeap::new(),
27 }
28 }
alloc(&mut self) -> usize29 fn alloc(&mut self) -> usize {
30 if let Some(id) = self.free_list.pop() {
31 id.0
32 } else {
33 let id = self.free_from;
34 self.free_from = self
35 .free_from
36 .checked_add(1)
37 .expect("Ran out of thread IDs");
38 id
39 }
40 }
free(&mut self, id: usize)41 fn free(&mut self, id: usize) {
42 self.free_list.push(Reverse(id));
43 }
44 }
45 static THREAD_ID_MANAGER: Lazy<Mutex<ThreadIdManager>> =
46 Lazy::new(|| Mutex::new(ThreadIdManager::new()));
47
48 /// Data which is unique to the current thread while it is running.
49 /// A thread ID may be reused after a thread exits.
50 #[derive(Clone, Copy)]
51 pub(crate) struct Thread {
52 /// The thread ID obtained from the thread ID manager.
53 pub(crate) id: usize,
54 /// The bucket this thread's local storage will be in.
55 pub(crate) bucket: usize,
56 /// The size of the bucket this thread's local storage will be in.
57 pub(crate) bucket_size: usize,
58 /// The index into the bucket this thread's local storage is in.
59 pub(crate) index: usize,
60 }
61 impl Thread {
new(id: usize) -> Thread62 fn new(id: usize) -> Thread {
63 let bucket = usize::from(POINTER_WIDTH) - id.leading_zeros() as usize;
64 let bucket_size = 1 << bucket.saturating_sub(1);
65 let index = if id != 0 { id ^ bucket_size } else { 0 };
66
67 Thread {
68 id,
69 bucket,
70 bucket_size,
71 index,
72 }
73 }
74 }
75
76 /// Wrapper around `Thread` that allocates and deallocates the ID.
77 struct ThreadHolder(Thread);
78 impl ThreadHolder {
new() -> ThreadHolder79 fn new() -> ThreadHolder {
80 ThreadHolder(Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc()))
81 }
82 }
83 impl Drop for ThreadHolder {
drop(&mut self)84 fn drop(&mut self) {
85 THREAD_ID_MANAGER.lock().unwrap().free(self.0.id);
86 }
87 }
88
89 thread_local!(static THREAD_HOLDER: ThreadHolder = ThreadHolder::new());
90
91 /// Get the current thread.
get() -> Thread92 pub(crate) fn get() -> Thread {
93 THREAD_HOLDER.with(|holder| holder.0)
94 }
95
96 #[test]
test_thread()97 fn test_thread() {
98 let thread = Thread::new(0);
99 assert_eq!(thread.id, 0);
100 assert_eq!(thread.bucket, 0);
101 assert_eq!(thread.bucket_size, 1);
102 assert_eq!(thread.index, 0);
103
104 let thread = Thread::new(1);
105 assert_eq!(thread.id, 1);
106 assert_eq!(thread.bucket, 1);
107 assert_eq!(thread.bucket_size, 1);
108 assert_eq!(thread.index, 0);
109
110 let thread = Thread::new(2);
111 assert_eq!(thread.id, 2);
112 assert_eq!(thread.bucket, 2);
113 assert_eq!(thread.bucket_size, 2);
114 assert_eq!(thread.index, 0);
115
116 let thread = Thread::new(3);
117 assert_eq!(thread.id, 3);
118 assert_eq!(thread.bucket, 2);
119 assert_eq!(thread.bucket_size, 2);
120 assert_eq!(thread.index, 1);
121
122 let thread = Thread::new(19);
123 assert_eq!(thread.id, 19);
124 assert_eq!(thread.bucket, 5);
125 assert_eq!(thread.bucket_size, 16);
126 assert_eq!(thread.index, 3);
127 }
128