1 //! The implementation is based on Dmitry Vyukov's bounded MPMC queue. 2 //! 3 //! Source: 4 //! - <http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue> 5 6 use alloc::boxed::Box; 7 use core::cell::UnsafeCell; 8 use core::fmt; 9 use core::marker::PhantomData; 10 use core::mem::{self, MaybeUninit}; 11 use core::sync::atomic::{self, AtomicUsize, Ordering}; 12 13 use crossbeam_utils::{Backoff, CachePadded}; 14 15 /// A slot in a queue. 16 struct Slot<T> { 17 /// The current stamp. 18 /// 19 /// If the stamp equals the tail, this node will be next written to. If it equals head + 1, 20 /// this node will be next read from. 21 stamp: AtomicUsize, 22 23 /// The value in this slot. 24 value: UnsafeCell<MaybeUninit<T>>, 25 } 26 27 /// A bounded multi-producer multi-consumer queue. 28 /// 29 /// This queue allocates a fixed-capacity buffer on construction, which is used to store pushed 30 /// elements. The queue cannot hold more elements than the buffer allows. Attempting to push an 31 /// element into a full queue will fail. Having a buffer allocated upfront makes this queue a bit 32 /// faster than [`SegQueue`]. 33 /// 34 /// [`SegQueue`]: super::SegQueue 35 /// 36 /// # Examples 37 /// 38 /// ``` 39 /// use crossbeam_queue::ArrayQueue; 40 /// 41 /// let q = ArrayQueue::new(2); 42 /// 43 /// assert_eq!(q.push('a'), Ok(())); 44 /// assert_eq!(q.push('b'), Ok(())); 45 /// assert_eq!(q.push('c'), Err('c')); 46 /// assert_eq!(q.pop(), Some('a')); 47 /// ``` 48 pub struct ArrayQueue<T> { 49 /// The head of the queue. 50 /// 51 /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a 52 /// single `usize`. The lower bits represent the index, while the upper bits represent the lap. 53 /// 54 /// Elements are popped from the head of the queue. 55 head: CachePadded<AtomicUsize>, 56 57 /// The tail of the queue. 58 /// 59 /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a 60 /// single `usize`. The lower bits represent the index, while the upper bits represent the lap. 61 /// 62 /// Elements are pushed into the tail of the queue. 63 tail: CachePadded<AtomicUsize>, 64 65 /// The buffer holding slots. 66 buffer: *mut Slot<T>, 67 68 /// The queue capacity. 69 cap: usize, 70 71 /// A stamp with the value of `{ lap: 1, index: 0 }`. 72 one_lap: usize, 73 74 /// Indicates that dropping an `ArrayQueue<T>` may drop elements of type `T`. 75 _marker: PhantomData<T>, 76 } 77 78 unsafe impl<T: Send> Sync for ArrayQueue<T> {} 79 unsafe impl<T: Send> Send for ArrayQueue<T> {} 80 81 impl<T> ArrayQueue<T> { 82 /// Creates a new bounded queue with the given capacity. 83 /// 84 /// # Panics 85 /// 86 /// Panics if the capacity is zero. 87 /// 88 /// # Examples 89 /// 90 /// ``` 91 /// use crossbeam_queue::ArrayQueue; 92 /// 93 /// let q = ArrayQueue::<i32>::new(100); 94 /// ``` new(cap: usize) -> ArrayQueue<T>95 pub fn new(cap: usize) -> ArrayQueue<T> { 96 assert!(cap > 0, "capacity must be non-zero"); 97 98 // Head is initialized to `{ lap: 0, index: 0 }`. 99 // Tail is initialized to `{ lap: 0, index: 0 }`. 100 let head = 0; 101 let tail = 0; 102 103 // Allocate a buffer of `cap` slots initialized 104 // with stamps. 105 let buffer = { 106 let mut boxed: Box<[Slot<T>]> = (0..cap) 107 .map(|i| { 108 // Set the stamp to `{ lap: 0, index: i }`. 109 Slot { 110 stamp: AtomicUsize::new(i), 111 value: UnsafeCell::new(MaybeUninit::uninit()), 112 } 113 }) 114 .collect(); 115 let ptr = boxed.as_mut_ptr(); 116 mem::forget(boxed); 117 ptr 118 }; 119 120 // One lap is the smallest power of two greater than `cap`. 121 let one_lap = (cap + 1).next_power_of_two(); 122 123 ArrayQueue { 124 buffer, 125 cap, 126 one_lap, 127 head: CachePadded::new(AtomicUsize::new(head)), 128 tail: CachePadded::new(AtomicUsize::new(tail)), 129 _marker: PhantomData, 130 } 131 } 132 133 /// Attempts to push an element into the queue. 134 /// 135 /// If the queue is full, the element is returned back as an error. 136 /// 137 /// # Examples 138 /// 139 /// ``` 140 /// use crossbeam_queue::ArrayQueue; 141 /// 142 /// let q = ArrayQueue::new(1); 143 /// 144 /// assert_eq!(q.push(10), Ok(())); 145 /// assert_eq!(q.push(20), Err(20)); 146 /// ``` push(&self, value: T) -> Result<(), T>147 pub fn push(&self, value: T) -> Result<(), T> { 148 let backoff = Backoff::new(); 149 let mut tail = self.tail.load(Ordering::Relaxed); 150 151 loop { 152 // Deconstruct the tail. 153 let index = tail & (self.one_lap - 1); 154 let lap = tail & !(self.one_lap - 1); 155 156 // Inspect the corresponding slot. 157 let slot = unsafe { &*self.buffer.add(index) }; 158 let stamp = slot.stamp.load(Ordering::Acquire); 159 160 // If the tail and the stamp match, we may attempt to push. 161 if tail == stamp { 162 let new_tail = if index + 1 < self.cap { 163 // Same lap, incremented index. 164 // Set to `{ lap: lap, index: index + 1 }`. 165 tail + 1 166 } else { 167 // One lap forward, index wraps around to zero. 168 // Set to `{ lap: lap.wrapping_add(1), index: 0 }`. 169 lap.wrapping_add(self.one_lap) 170 }; 171 172 // Try moving the tail. 173 match self.tail.compare_exchange_weak( 174 tail, 175 new_tail, 176 Ordering::SeqCst, 177 Ordering::Relaxed, 178 ) { 179 Ok(_) => { 180 // Write the value into the slot and update the stamp. 181 unsafe { 182 slot.value.get().write(MaybeUninit::new(value)); 183 } 184 slot.stamp.store(tail + 1, Ordering::Release); 185 return Ok(()); 186 } 187 Err(t) => { 188 tail = t; 189 backoff.spin(); 190 } 191 } 192 } else if stamp.wrapping_add(self.one_lap) == tail + 1 { 193 atomic::fence(Ordering::SeqCst); 194 let head = self.head.load(Ordering::Relaxed); 195 196 // If the head lags one lap behind the tail as well... 197 if head.wrapping_add(self.one_lap) == tail { 198 // ...then the queue is full. 199 return Err(value); 200 } 201 202 backoff.spin(); 203 tail = self.tail.load(Ordering::Relaxed); 204 } else { 205 // Snooze because we need to wait for the stamp to get updated. 206 backoff.snooze(); 207 tail = self.tail.load(Ordering::Relaxed); 208 } 209 } 210 } 211 212 /// Attempts to pop an element from the queue. 213 /// 214 /// If the queue is empty, `None` is returned. 215 /// 216 /// # Examples 217 /// 218 /// ``` 219 /// use crossbeam_queue::ArrayQueue; 220 /// 221 /// let q = ArrayQueue::new(1); 222 /// assert_eq!(q.push(10), Ok(())); 223 /// 224 /// assert_eq!(q.pop(), Some(10)); 225 /// assert!(q.pop().is_none()); 226 /// ``` pop(&self) -> Option<T>227 pub fn pop(&self) -> Option<T> { 228 let backoff = Backoff::new(); 229 let mut head = self.head.load(Ordering::Relaxed); 230 231 loop { 232 // Deconstruct the head. 233 let index = head & (self.one_lap - 1); 234 let lap = head & !(self.one_lap - 1); 235 236 // Inspect the corresponding slot. 237 let slot = unsafe { &*self.buffer.add(index) }; 238 let stamp = slot.stamp.load(Ordering::Acquire); 239 240 // If the the stamp is ahead of the head by 1, we may attempt to pop. 241 if head + 1 == stamp { 242 let new = if index + 1 < self.cap { 243 // Same lap, incremented index. 244 // Set to `{ lap: lap, index: index + 1 }`. 245 head + 1 246 } else { 247 // One lap forward, index wraps around to zero. 248 // Set to `{ lap: lap.wrapping_add(1), index: 0 }`. 249 lap.wrapping_add(self.one_lap) 250 }; 251 252 // Try moving the head. 253 match self.head.compare_exchange_weak( 254 head, 255 new, 256 Ordering::SeqCst, 257 Ordering::Relaxed, 258 ) { 259 Ok(_) => { 260 // Read the value from the slot and update the stamp. 261 let msg = unsafe { slot.value.get().read().assume_init() }; 262 slot.stamp 263 .store(head.wrapping_add(self.one_lap), Ordering::Release); 264 return Some(msg); 265 } 266 Err(h) => { 267 head = h; 268 backoff.spin(); 269 } 270 } 271 } else if stamp == head { 272 atomic::fence(Ordering::SeqCst); 273 let tail = self.tail.load(Ordering::Relaxed); 274 275 // If the tail equals the head, that means the channel is empty. 276 if tail == head { 277 return None; 278 } 279 280 backoff.spin(); 281 head = self.head.load(Ordering::Relaxed); 282 } else { 283 // Snooze because we need to wait for the stamp to get updated. 284 backoff.snooze(); 285 head = self.head.load(Ordering::Relaxed); 286 } 287 } 288 } 289 290 /// Returns the capacity of the queue. 291 /// 292 /// # Examples 293 /// 294 /// ``` 295 /// use crossbeam_queue::ArrayQueue; 296 /// 297 /// let q = ArrayQueue::<i32>::new(100); 298 /// 299 /// assert_eq!(q.capacity(), 100); 300 /// ``` capacity(&self) -> usize301 pub fn capacity(&self) -> usize { 302 self.cap 303 } 304 305 /// Returns `true` if the queue is empty. 306 /// 307 /// # Examples 308 /// 309 /// ``` 310 /// use crossbeam_queue::ArrayQueue; 311 /// 312 /// let q = ArrayQueue::new(100); 313 /// 314 /// assert!(q.is_empty()); 315 /// q.push(1).unwrap(); 316 /// assert!(!q.is_empty()); 317 /// ``` is_empty(&self) -> bool318 pub fn is_empty(&self) -> bool { 319 let head = self.head.load(Ordering::SeqCst); 320 let tail = self.tail.load(Ordering::SeqCst); 321 322 // Is the tail lagging one lap behind head? 323 // Is the tail equal to the head? 324 // 325 // Note: If the head changes just before we load the tail, that means there was a moment 326 // when the channel was not empty, so it is safe to just return `false`. 327 tail == head 328 } 329 330 /// Returns `true` if the queue is full. 331 /// 332 /// # Examples 333 /// 334 /// ``` 335 /// use crossbeam_queue::ArrayQueue; 336 /// 337 /// let q = ArrayQueue::new(1); 338 /// 339 /// assert!(!q.is_full()); 340 /// q.push(1).unwrap(); 341 /// assert!(q.is_full()); 342 /// ``` is_full(&self) -> bool343 pub fn is_full(&self) -> bool { 344 let tail = self.tail.load(Ordering::SeqCst); 345 let head = self.head.load(Ordering::SeqCst); 346 347 // Is the head lagging one lap behind tail? 348 // 349 // Note: If the tail changes just before we load the head, that means there was a moment 350 // when the queue was not full, so it is safe to just return `false`. 351 head.wrapping_add(self.one_lap) == tail 352 } 353 354 /// Returns the number of elements in the queue. 355 /// 356 /// # Examples 357 /// 358 /// ``` 359 /// use crossbeam_queue::ArrayQueue; 360 /// 361 /// let q = ArrayQueue::new(100); 362 /// assert_eq!(q.len(), 0); 363 /// 364 /// q.push(10).unwrap(); 365 /// assert_eq!(q.len(), 1); 366 /// 367 /// q.push(20).unwrap(); 368 /// assert_eq!(q.len(), 2); 369 /// ``` len(&self) -> usize370 pub fn len(&self) -> usize { 371 loop { 372 // Load the tail, then load the head. 373 let tail = self.tail.load(Ordering::SeqCst); 374 let head = self.head.load(Ordering::SeqCst); 375 376 // If the tail didn't change, we've got consistent values to work with. 377 if self.tail.load(Ordering::SeqCst) == tail { 378 let hix = head & (self.one_lap - 1); 379 let tix = tail & (self.one_lap - 1); 380 381 return if hix < tix { 382 tix - hix 383 } else if hix > tix { 384 self.cap - hix + tix 385 } else if tail == head { 386 0 387 } else { 388 self.cap 389 }; 390 } 391 } 392 } 393 } 394 395 impl<T> Drop for ArrayQueue<T> { drop(&mut self)396 fn drop(&mut self) { 397 // Get the index of the head. 398 let hix = self.head.load(Ordering::Relaxed) & (self.one_lap - 1); 399 400 // Loop over all slots that hold a message and drop them. 401 for i in 0..self.len() { 402 // Compute the index of the next slot holding a message. 403 let index = if hix + i < self.cap { 404 hix + i 405 } else { 406 hix + i - self.cap 407 }; 408 409 unsafe { 410 let p = { 411 let slot = &mut *self.buffer.add(index); 412 let value = &mut *slot.value.get(); 413 value.as_mut_ptr() 414 }; 415 p.drop_in_place(); 416 } 417 } 418 419 // Finally, deallocate the buffer, but don't run any destructors. 420 unsafe { 421 // Create a slice from the buffer to make 422 // a fat pointer. Then, use Box::from_raw 423 // to deallocate it. 424 let ptr = core::slice::from_raw_parts_mut(self.buffer, self.cap) as *mut [Slot<T>]; 425 Box::from_raw(ptr); 426 } 427 } 428 } 429 430 impl<T> fmt::Debug for ArrayQueue<T> { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result431 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 432 f.pad("ArrayQueue { .. }") 433 } 434 } 435