1use std::ptr::NonNull;
2use std::sync::atomic::Ordering;
3
4use crate::loom::sync::{Mutex, MutexGuard};
5use std::sync::atomic::AtomicUsize;
6
7use super::linked_list::{Link, LinkedList};
8
9/// An intrusive linked list supporting highly concurrent updates.
10///
11/// It currently relies on `LinkedList`, so it is the caller's
12/// responsibility to ensure the list is empty before dropping it.
13///
14/// Note: Due to its inner sharded design, the order of nodes cannot be guaranteed.
15pub(crate) struct ShardedList<L, T> {
16 lists: Box<[Mutex<LinkedList<L, T>>]>,
17 count: AtomicUsize,
18 shard_mask: usize,
19}
20
21/// Determines which linked list an item should be stored in.
22///
23/// # Safety
24///
25/// Implementations must guarantee that the id of an item does not change from
26/// call to call.
27pub(crate) unsafe trait ShardedListItem: Link {
28 /// # Safety
29 /// The provided pointer must point at a valid list item.
30 unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize;
31}
32
33impl<L, T> ShardedList<L, T> {
34 /// Creates a new and empty sharded linked list with the specified size.
35 pub(crate) fn new(sharded_size: usize) -> Self {
36 assert!(sharded_size.is_power_of_two());
37
38 let shard_mask = sharded_size - 1;
39 let mut lists = Vec::with_capacity(sharded_size);
40 for _ in 0..sharded_size {
41 lists.push(Mutex::new(LinkedList::<L, T>::new()))
42 }
43 Self {
44 lists: lists.into_boxed_slice(),
45 count: AtomicUsize::new(0),
46 shard_mask,
47 }
48 }
49}
50
51/// Used to get the lock of shard.
52pub(crate) struct ShardGuard<'a, L, T> {
53 lock: MutexGuard<'a, LinkedList<L, T>>,
54 count: &'a AtomicUsize,
55 id: usize,
56}
57
58impl<L: ShardedListItem> ShardedList<L, L::Target> {
59 /// Removes the last element from a list specified by shard_id and returns it, or None if it is
60 /// empty.
61 pub(crate) fn pop_back(&self, shard_id: usize) -> Option<L::Handle> {
62 let mut lock = self.shard_inner(shard_id);
63 let node = lock.pop_back();
64 if node.is_some() {
65 self.count.fetch_sub(1, Ordering::Relaxed);
66 }
67 node
68 }
69
70 /// Removes the specified node from the list.
71 ///
72 /// # Safety
73 ///
74 /// The caller **must** ensure that exactly one of the following is true:
75 /// - `node` is currently contained by `self`,
76 /// - `node` is not contained by any list,
77 /// - `node` is currently contained by some other `GuardedLinkedList`.
78 pub(crate) unsafe fn remove(&self, node: NonNull<L::Target>) -> Option<L::Handle> {
79 let id = L::get_shard_id(node);
80 let mut lock = self.shard_inner(id);
81 // SAFETY: Since the shard id cannot change, it's not possible for this node
82 // to be in any other list of the same sharded list.
83 let node = unsafe { lock.remove(node) };
84 if node.is_some() {
85 self.count.fetch_sub(1, Ordering::Relaxed);
86 }
87 node
88 }
89
90 /// Gets the lock of ShardedList, makes us have the write permission.
91 pub(crate) fn lock_shard(&self, val: &L::Handle) -> ShardGuard<'_, L, L::Target> {
92 let id = unsafe { L::get_shard_id(L::as_raw(val)) };
93 ShardGuard {
94 lock: self.shard_inner(id),
95 count: &self.count,
96 id,
97 }
98 }
99
100 /// Gets the count of elements in this list.
101 pub(crate) fn len(&self) -> usize {
102 self.count.load(Ordering::Relaxed)
103 }
104
105 /// Returns whether the linked list does not contain any node.
106 pub(crate) fn is_empty(&self) -> bool {
107 self.len() == 0
108 }
109
110 /// Gets the shard size of this SharedList.
111 ///
112 /// Used to help us to decide the parameter `shard_id` of the `pop_back` method.
113 pub(crate) fn shard_size(&self) -> usize {
114 self.shard_mask + 1
115 }
116
117 #[inline]
118 fn shard_inner(&self, id: usize) -> MutexGuard<'_, LinkedList<L, <L as Link>::Target>> {
119 // Safety: This modulo operation ensures that the index is not out of bounds.
120 unsafe { self.lists.get_unchecked(id & self.shard_mask).lock() }
121 }
122}
123
124impl<'a, L: ShardedListItem> ShardGuard<'a, L, L::Target> {
125 /// Push a value to this shard.
126 pub(crate) fn push(mut self, val: L::Handle) {
127 let id = unsafe { L::get_shard_id(L::as_raw(&val)) };
128 assert_eq!(id, self.id);
129 self.lock.push_front(val);
130 self.count.fetch_add(1, Ordering::Relaxed);
131 }
132}
133
134cfg_taskdump! {
135 impl<L: ShardedListItem> ShardedList<L, L::Target> {
136 pub(crate) fn for_each<F>(&self, mut f: F)
137 where
138 F: FnMut(&L::Handle),
139 {
140 let mut guards = Vec::with_capacity(self.lists.len());
141 for list in self.lists.iter() {
142 guards.push(list.lock());
143 }
144 for g in &mut guards {
145 g.for_each(&mut f);
146 }
147 }
148 }
149}
150