1use std::cell::UnsafeCell;
2use std::ptr;
3use std::sync::atomic::{fence, AtomicPtr, AtomicUsize, Ordering};
4
5#[allow(dead_code)]
6pub struct LockFreeQueue<T> {
7 head: AtomicPtr<Node<T>>,
8 tail: AtomicPtr<Node<T>>,
9 cache_line_pad: [u8; 64],
10}
11
12struct Node<T> {
13 value: UnsafeCell<Option<T>>,
14 next: AtomicPtr<Node<T>>,
15}
16
17impl<T> LockFreeQueue<T> {
18 pub fn new() -> Self {
19 let sentinel = Box::into_raw(Box::new(Node {
20 value: UnsafeCell::new(None),
21 next: AtomicPtr::new(ptr::null_mut()),
22 }));
23
24 LockFreeQueue {
25 head: AtomicPtr::new(sentinel),
26 tail: AtomicPtr::new(sentinel),
27 cache_line_pad: [0; 64],
28 }
29 }
30
31 #[inline(always)]
32 pub fn enqueue(&self, value: T) {
33 let new_node = Box::into_raw(Box::new(Node {
34 value: UnsafeCell::new(Some(value)),
35 next: AtomicPtr::new(ptr::null_mut()),
36 }));
37
38 loop {
39 let tail = self.tail.load(Ordering::Acquire);
40 let tail_next = unsafe { (*tail).next.load(Ordering::Acquire) };
41
42 if tail == self.tail.load(Ordering::Acquire) {
43 if tail_next.is_null() {
44 if unsafe {
45 (*tail).next.compare_exchange_weak(
46 tail_next,
47 new_node,
48 Ordering::Release,
49 Ordering::Relaxed,
50 )
51 }
52 .is_ok()
53 {
54 let _ = self.tail.compare_exchange_weak(
55 tail,
56 new_node,
57 Ordering::Release,
58 Ordering::Relaxed,
59 );
60 return;
61 }
62 } else {
63 let _ = self.tail.compare_exchange_weak(
64 tail,
65 tail_next,
66 Ordering::Release,
67 Ordering::Relaxed,
68 );
69 }
70 }
71 core::hint::spin_loop();
72 }
73 }
74
75 #[inline(always)]
76 pub fn dequeue(&self) -> Option<T> {
77 loop {
78 let head = self.head.load(Ordering::Acquire);
79 let tail = self.tail.load(Ordering::Acquire);
80 let head_next = unsafe { (*head).next.load(Ordering::Acquire) };
81
82 if head == self.head.load(Ordering::Acquire) {
83 if head == tail {
84 if head_next.is_null() {
85 return None;
86 }
87 let _ = self.tail.compare_exchange_weak(
88 tail,
89 head_next,
90 Ordering::Release,
91 Ordering::Relaxed,
92 );
93 } else {
94 if self
95 .head
96 .compare_exchange_weak(
97 head,
98 head_next,
99 Ordering::Release,
100 Ordering::Relaxed,
101 )
102 .is_ok()
103 {
104 let value = unsafe { ptr::read(&(*head_next).value) }.into_inner();
105
106 unsafe { drop(Box::from_raw(head)) };
107
108 return value;
109 }
110 }
111 }
112 core::hint::spin_loop();
113 }
114 }
115
116 #[inline(always)]
117 pub fn is_empty(&self) -> bool {
118 let head = self.head.load(Ordering::Acquire);
119 let head_next = unsafe { (*head).next.load(Ordering::Acquire) };
120 head_next.is_null()
121 }
122}
123
124impl<T> Drop for LockFreeQueue<T> {
125 fn drop(&mut self) {
126 while self.dequeue().is_some() {}
127
128 let sentinel = self.head.load(Ordering::Relaxed);
129 unsafe { drop(Box::from_raw(sentinel)) };
130 }
131}
132
133#[allow(dead_code)]
134pub struct BoundedLockFreeQueue<T> {
135 buffer: *mut Node<T>,
136 capacity: usize,
137 head: AtomicUsize,
138 tail: AtomicUsize,
139 cache_line_pad: [u8; 64],
140}
141
142impl<T> BoundedLockFreeQueue<T> {
143 pub fn new(capacity: usize) -> Self {
144 let capacity = capacity.next_power_of_two();
145
146 let mut buffer = Vec::with_capacity(capacity);
147 for _ in 0..capacity {
148 buffer.push(Node {
149 value: UnsafeCell::new(None),
150 next: AtomicPtr::new(ptr::null_mut()),
151 });
152 }
153
154 let buffer = buffer.leak().as_mut_ptr();
155
156 BoundedLockFreeQueue {
157 buffer,
158 capacity,
159 head: AtomicUsize::new(0),
160 tail: AtomicUsize::new(0),
161 cache_line_pad: [0; 64],
162 }
163 }
164
165 #[inline(always)]
166 pub fn enqueue(&self, value: T) -> Result<(), T> {
167 let mask = self.capacity - 1;
168 let mut tail = self.tail.load(Ordering::Relaxed);
169
170 loop {
171 let head = self.head.load(Ordering::Acquire);
172
173 if (tail - head) >= self.capacity {
174 return Err(value);
175 }
176
177 if self
178 .tail
179 .compare_exchange_weak(tail, tail + 1, Ordering::Relaxed, Ordering::Relaxed)
180 .is_ok()
181 {
182 let index = tail & mask;
183 let node = unsafe { &*self.buffer.add(index) };
184
185 unsafe { *node.value.get() = Some(value) };
186
187 fence(Ordering::Release);
188
189 return Ok(());
190 }
191
192 tail = self.tail.load(Ordering::Relaxed);
193 core::hint::spin_loop();
194 }
195 }
196
197 #[inline(always)]
198 pub fn dequeue(&self) -> Option<T> {
199 let mask = self.capacity - 1;
200 let mut head = self.head.load(Ordering::Relaxed);
201
202 loop {
203 let tail = self.tail.load(Ordering::Acquire);
204
205 if head >= tail {
206 return None;
207 }
208
209 if self
210 .head
211 .compare_exchange_weak(head, head + 1, Ordering::Relaxed, Ordering::Relaxed)
212 .is_ok()
213 {
214 let index = head & mask;
215 let node = unsafe { &*self.buffer.add(index) };
216
217 fence(Ordering::Acquire);
218
219 let value = unsafe { (*node.value.get()).take() };
220
221 return value;
222 }
223
224 head = self.head.load(Ordering::Relaxed);
225 core::hint::spin_loop();
226 }
227 }
228
229 #[inline(always)]
230 pub fn is_empty(&self) -> bool {
231 self.head.load(Ordering::Acquire) >= self.tail.load(Ordering::Acquire)
232 }
233
234 #[inline(always)]
235 pub fn is_full(&self) -> bool {
236 let head = self.head.load(Ordering::Acquire);
237 let tail = self.tail.load(Ordering::Acquire);
238 (tail - head) >= self.capacity
239 }
240
241 #[inline(always)]
242 pub fn len(&self) -> usize {
243 let head = self.head.load(Ordering::Acquire);
244 let tail = self.tail.load(Ordering::Acquire);
245 tail.saturating_sub(head)
246 }
247
248 #[inline(always)]
249 pub fn capacity(&self) -> usize {
250 self.capacity
251 }
252}
253
254unsafe impl<T: Send> Send for BoundedLockFreeQueue<T> {}
255unsafe impl<T: Send> Sync for BoundedLockFreeQueue<T> {}
256
257impl<T> Drop for BoundedLockFreeQueue<T> {
258 fn drop(&mut self) {
259 while self.dequeue().is_some() {}
260
261 unsafe {
262 Vec::from_raw_parts(self.buffer, self.capacity, self.capacity);
263 }
264 }
265}
266
267unsafe impl<T: Send> Send for LockFreeQueue<T> {}
268unsafe impl<T: Send> Sync for LockFreeQueue<T> {}