src/libstd/sync/mpsc/sync.rs

   1 /// Synchronous channels/ports
   2 ///
   3 /// This channel implementation differs significantly from the asynchronous
   4 /// implementations found next to it (oneshot/stream/share). This is an
   5 /// implementation of a synchronous, bounded buffer channel.
   6 ///
   7 /// Each channel is created with some amount of backing buffer, and sends will
   8 /// *block* until buffer space becomes available. A buffer size of 0 is valid,
   9 /// which means that every successful send is paired with a successful recv.
  10 ///
  11 /// This flavor of channels defines a new `send_opt` method for channels which
  12 /// is the method by which a message is sent but the thread does not panic if it
  13 /// cannot be delivered.
  14 ///
  15 /// Another major difference is that send() will *always* return back the data
  16 /// if it couldn't be sent. This is because it is deterministically known when
  17 /// the data is received and when it is not received.
  18 ///
  19 /// Implementation-wise, it can all be summed up with "use a mutex plus some
  20 /// logic". The mutex used here is an OS native mutex, meaning that no user code
  21 /// is run inside of the mutex (to prevent context switching). This
  22 /// implementation shares almost all code for the buffered and unbuffered cases
  23 /// of a synchronous channel. There are a few branches for the unbuffered case,
  24 /// but they're mostly just relevant to blocking senders.
  25
  26 pub use self::Failure::*;
  27 use self::Blocker::*;
  28
  29 use core::intrinsics::abort;
  30 use core::isize;
  31 use core::mem;
  32 use core::ptr;
  33
  34 use crate::sync::atomic::{Ordering, AtomicUsize};
  35 use crate::sync::mpsc::blocking::{self, WaitToken, SignalToken};
  36 use crate::sync::{Mutex, MutexGuard};
  37 use crate::time::Instant;
  38
  39 const MAX_REFCOUNT: usize = (isize::MAX) as usize;
  40
  41 pub struct Packet<T> {
  42     /// Only field outside of the mutex. Just done for kicks, but mainly because
  43     /// the other shared channel already had the code implemented
  44     channels: AtomicUsize,
  45
  46     lock: Mutex<State<T>>,
  47 }
  48
  49 unsafe impl<T: Send> Send for Packet<T> { }
  50
  51 unsafe impl<T: Send> Sync for Packet<T> { }
  52
  53 struct State<T> {
  54     disconnected: bool, // Is the channel disconnected yet?
  55     queue: Queue,       // queue of senders waiting to send data
  56     blocker: Blocker,   // currently blocked thread on this channel
  57     buf: Buffer<T>,     // storage for buffered messages
  58     cap: usize,         // capacity of this channel
  59
  60     /// A curious flag used to indicate whether a sender failed or succeeded in
  61     /// blocking. This is used to transmit information back to the thread that it
  62     /// must dequeue its message from the buffer because it was not received.
  63     /// This is only relevant in the 0-buffer case. This obviously cannot be
  64     /// safely constructed, but it's guaranteed to always have a valid pointer
  65     /// value.
  66     canceled: Option<&'static mut bool>,
  67 }
  68
  69 unsafe impl<T: Send> Send for State<T> {}
  70
  71 /// Possible flavors of threads who can be blocked on this channel.
  72 enum Blocker {
  73     BlockedSender(SignalToken),
  74     BlockedReceiver(SignalToken),
  75     NoneBlocked
  76 }
  77
  78 /// Simple queue for threading threads together. Nodes are stack-allocated, so
  79 /// this structure is not safe at all
  80 struct Queue {
  81     head: *mut Node,
  82     tail: *mut Node,
  83 }
  84
  85 struct Node {
  86     token: Option<SignalToken>,
  87     next: *mut Node,
  88 }
  89
  90 unsafe impl Send for Node {}
  91
  92 /// A simple ring-buffer
  93 struct Buffer<T> {
  94     buf: Vec<Option<T>>,
  95     start: usize,
  96     size: usize,
  97 }
  98
  99 #[derive(Debug)]
 100 pub enum Failure {
 101     Empty,
 102     Disconnected,
 103 }
 104
 105 /// Atomically blocks the current thread, placing it into `slot`, unlocking `lock`
 106 /// in the meantime. This re-locks the mutex upon returning.
 107 fn wait<'a, 'b, T>(lock: &'a Mutex<State<T>>,
 108                    mut guard: MutexGuard<'b, State<T>>,
 109                    f: fn(SignalToken) -> Blocker)
 110                    -> MutexGuard<'a, State<T>>
 111 {
 112     let (wait_token, signal_token) = blocking::tokens();
 113     match mem::replace(&mut guard.blocker, f(signal_token)) {
 114         NoneBlocked => {}
 115         _ => unreachable!(),
 116     }
 117     drop(guard);         // unlock
 118     wait_token.wait();   // block
 119     lock.lock().unwrap() // relock
 120 }
 121
 122 /// Same as wait, but waiting at most until `deadline`.
 123 fn wait_timeout_receiver<'a, 'b, T>(lock: &'a Mutex<State<T>>,
 124                                     deadline: Instant,
 125                                     mut guard: MutexGuard<'b, State<T>>,
 126                                     success: &mut bool)
 127                                     -> MutexGuard<'a, State<T>>
 128 {
 129     let (wait_token, signal_token) = blocking::tokens();
 130     match mem::replace(&mut guard.blocker, BlockedReceiver(signal_token)) {
 131         NoneBlocked => {}
 132         _ => unreachable!(),
 133     }
 134     drop(guard);         // unlock
 135     *success = wait_token.wait_max_until(deadline);   // block
 136     let mut new_guard = lock.lock().unwrap(); // relock
 137     if !*success {
 138         abort_selection(&mut new_guard);
 139     }
 140     new_guard
 141 }
 142
 143 fn abort_selection<'a, T>(guard: &mut MutexGuard<'a , State<T>>) -> bool {
 144     match mem::replace(&mut guard.blocker, NoneBlocked) {
 145         NoneBlocked => true,
 146         BlockedSender(token) => {
 147             guard.blocker = BlockedSender(token);
 148             true
 149         }
 150         BlockedReceiver(token) => { drop(token); false }
 151     }
 152 }
 153
 154 /// Wakes up a thread, dropping the lock at the correct time
 155 fn wakeup<T>(token: SignalToken, guard: MutexGuard<'_, State<T>>) {
 156     // We need to be careful to wake up the waiting thread *outside* of the mutex
 157     // in case it incurs a context switch.
 158     drop(guard);
 159     token.signal();
 160 }
 161
 162 impl<T> Packet<T> {
 163     pub fn new(cap: usize) -> Packet<T> {
 164         Packet {
 165             channels: AtomicUsize::new(1),
 166             lock: Mutex::new(State {
 167                 disconnected: false,
 168                 blocker: NoneBlocked,
 169                 cap,
 170                 canceled: None,
 171                 queue: Queue {
 172                     head: ptr::null_mut(),
 173                     tail: ptr::null_mut(),
 174                 },
 175                 buf: Buffer {
 176                     buf: (0..cap + if cap == 0 {1} else {0}).map(|_| None).collect(),
 177                     start: 0,
 178                     size: 0,
 179                 },
 180             }),
 181         }
 182     }
 183
 184     // wait until a send slot is available, returning locked access to
 185     // the channel state.
 186     fn acquire_send_slot(&self) -> MutexGuard<'_, State<T>> {
 187         let mut node = Node { token: None, next: ptr::null_mut() };
 188         loop {
 189             let mut guard = self.lock.lock().unwrap();
 190             // are we ready to go?
 191             if guard.disconnected || guard.buf.size() < guard.buf.cap() {
 192                 return guard;
 193             }
 194             // no room; actually block
 195             let wait_token = guard.queue.enqueue(&mut node);
 196             drop(guard);
 197             wait_token.wait();
 198         }
 199     }
 200
 201     pub fn send(&self, t: T) -> Result<(), T> {
 202         let mut guard = self.acquire_send_slot();
 203         if guard.disconnected { return Err(t) }
 204         guard.buf.enqueue(t);
 205
 206         match mem::replace(&mut guard.blocker, NoneBlocked) {
 207             // if our capacity is 0, then we need to wait for a receiver to be
 208             // available to take our data. After waiting, we check again to make
 209             // sure the port didn't go away in the meantime. If it did, we need
 210             // to hand back our data.
 211             NoneBlocked if guard.cap == 0 => {
 212                 let mut canceled = false;
 213                 assert!(guard.canceled.is_none());
 214                 guard.canceled = Some(unsafe { mem::transmute(&mut canceled) });
 215                 let mut guard = wait(&self.lock, guard, BlockedSender);
 216                 if canceled {Err(guard.buf.dequeue())} else {Ok(())}
 217             }
 218
 219             // success, we buffered some data
 220             NoneBlocked => Ok(()),
 221
 222             // success, someone's about to receive our buffered data.
 223             BlockedReceiver(token) => { wakeup(token, guard); Ok(()) }
 224
 225             BlockedSender(..) => panic!("lolwut"),
 226         }
 227     }
 228
 229     pub fn try_send(&self, t: T) -> Result<(), super::TrySendError<T>> {
 230         let mut guard = self.lock.lock().unwrap();
 231         if guard.disconnected {
 232             Err(super::TrySendError::Disconnected(t))
 233         } else if guard.buf.size() == guard.buf.cap() {
 234             Err(super::TrySendError::Full(t))
 235         } else if guard.cap == 0 {
 236             // With capacity 0, even though we have buffer space we can't
 237             // transfer the data unless there's a receiver waiting.
 238             match mem::replace(&mut guard.blocker, NoneBlocked) {
 239                 NoneBlocked => Err(super::TrySendError::Full(t)),
 240                 BlockedSender(..) => unreachable!(),
 241                 BlockedReceiver(token) => {
 242                     guard.buf.enqueue(t);
 243                     wakeup(token, guard);
 244                     Ok(())
 245                 }
 246             }
 247         } else {
 248             // If the buffer has some space and the capacity isn't 0, then we
 249             // just enqueue the data for later retrieval, ensuring to wake up
 250             // any blocked receiver if there is one.
 251             assert!(guard.buf.size() < guard.buf.cap());
 252             guard.buf.enqueue(t);
 253             match mem::replace(&mut guard.blocker, NoneBlocked) {
 254                 BlockedReceiver(token) => wakeup(token, guard),
 255                 NoneBlocked => {}
 256                 BlockedSender(..) => unreachable!(),
 257             }
 258             Ok(())
 259         }
 260     }
 261
 262     // Receives a message from this channel
 263     //
 264     // When reading this, remember that there can only ever be one receiver at
 265     // time.
 266     pub fn recv(&self, deadline: Option<Instant>) -> Result<T, Failure> {
 267         let mut guard = self.lock.lock().unwrap();
 268
 269         let mut woke_up_after_waiting = false;
 270         // Wait for the buffer to have something in it. No need for a
 271         // while loop because we're the only receiver.
 272         if !guard.disconnected && guard.buf.size() == 0 {
 273             if let Some(deadline) = deadline {
 274                 guard = wait_timeout_receiver(&self.lock,
 275                                               deadline,
 276                                               guard,
 277                                               &mut woke_up_after_waiting);
 278             } else {
 279                 guard = wait(&self.lock, guard, BlockedReceiver);
 280                 woke_up_after_waiting = true;
 281             }
 282         }
 283
 284         // N.B., channel could be disconnected while waiting, so the order of
 285         // these conditionals is important.
 286         if guard.disconnected && guard.buf.size() == 0 {
 287             return Err(Disconnected);
 288         }
 289
 290         // Pick up the data, wake up our neighbors, and carry on
 291         assert!(guard.buf.size() > 0 || (deadline.is_some() && !woke_up_after_waiting));
 292
 293         if guard.buf.size() == 0 { return Err(Empty); }
 294
 295         let ret = guard.buf.dequeue();
 296         self.wakeup_senders(woke_up_after_waiting, guard);
 297         Ok(ret)
 298     }
 299
 300     pub fn try_recv(&self) -> Result<T, Failure> {
 301         let mut guard = self.lock.lock().unwrap();
 302
 303         // Easy cases first
 304         if guard.disconnected && guard.buf.size() == 0 { return Err(Disconnected) }
 305         if guard.buf.size() == 0 { return Err(Empty) }
 306
 307         // Be sure to wake up neighbors
 308         let ret = Ok(guard.buf.dequeue());
 309         self.wakeup_senders(false, guard);
 310         ret
 311     }
 312
 313     // Wake up pending senders after some data has been received
 314     //
 315     // * `waited` - flag if the receiver blocked to receive some data, or if it
 316     //              just picked up some data on the way out
 317     // * `guard` - the lock guard that is held over this channel's lock
 318     fn wakeup_senders(&self, waited: bool, mut guard: MutexGuard<'_, State<T>>) {
 319         let pending_sender1: Option<SignalToken> = guard.queue.dequeue();
 320
 321         // If this is a no-buffer channel (cap == 0), then if we didn't wait we
 322         // need to ACK the sender. If we waited, then the sender waking us up
 323         // was already the ACK.
 324         let pending_sender2 = if guard.cap == 0 && !waited {
 325             match mem::replace(&mut guard.blocker, NoneBlocked) {
 326                 NoneBlocked => None,
 327                 BlockedReceiver(..) => unreachable!(),
 328                 BlockedSender(token) => {
 329                     guard.canceled.take();
 330                     Some(token)
 331                 }
 332             }
 333         } else {
 334             None
 335         };
 336         mem::drop(guard);
 337
 338         // only outside of the lock do we wake up the pending threads
 339         pending_sender1.map(|t| t.signal());
 340         pending_sender2.map(|t| t.signal());
 341     }
 342
 343     // Prepares this shared packet for a channel clone, essentially just bumping
 344     // a refcount.
 345     pub fn clone_chan(&self) {
 346         let old_count = self.channels.fetch_add(1, Ordering::SeqCst);
 347
 348         // See comments on Arc::clone() on why we do this (for `mem::forget`).
 349         if old_count > MAX_REFCOUNT {
 350             unsafe {
 351                 abort();
 352             }
 353         }
 354     }
 355
 356     pub fn drop_chan(&self) {
 357         // Only flag the channel as disconnected if we're the last channel
 358         match self.channels.fetch_sub(1, Ordering::SeqCst) {
 359             1 => {}
 360             _ => return
 361         }
 362
 363         // Not much to do other than wake up a receiver if one's there
 364         let mut guard = self.lock.lock().unwrap();
 365         if guard.disconnected { return }
 366         guard.disconnected = true;
 367         match mem::replace(&mut guard.blocker, NoneBlocked) {
 368             NoneBlocked => {}
 369             BlockedSender(..) => unreachable!(),
 370             BlockedReceiver(token) => wakeup(token, guard),
 371         }
 372     }
 373
 374     pub fn drop_port(&self) {
 375         let mut guard = self.lock.lock().unwrap();
 376
 377         if guard.disconnected { return }
 378         guard.disconnected = true;
 379
 380         // If the capacity is 0, then the sender may want its data back after
 381         // we're disconnected. Otherwise it's now our responsibility to destroy
 382         // the buffered data. As with many other portions of this code, this
 383         // needs to be careful to destroy the data *outside* of the lock to
 384         // prevent deadlock.
 385         let _data = if guard.cap != 0 {
 386             mem::replace(&mut guard.buf.buf, Vec::new())
 387         } else {
 388             Vec::new()
 389         };
 390         let mut queue = mem::replace(&mut guard.queue, Queue {
 391             head: ptr::null_mut(),
 392             tail: ptr::null_mut(),
 393         });
 394
 395         let waiter = match mem::replace(&mut guard.blocker, NoneBlocked) {
 396             NoneBlocked => None,
 397             BlockedSender(token) => {
 398                 *guard.canceled.take().unwrap() = true;
 399                 Some(token)
 400             }
 401             BlockedReceiver(..) => unreachable!(),
 402         };
 403         mem::drop(guard);
 404
 405         while let Some(token) = queue.dequeue() { token.signal(); }
 406         waiter.map(|t| t.signal());
 407     }
 408 }
 409
 410 impl<T> Drop for Packet<T> {
 411     fn drop(&mut self) {
 412         assert_eq!(self.channels.load(Ordering::SeqCst), 0);
 413         let mut guard = self.lock.lock().unwrap();
 414         assert!(guard.queue.dequeue().is_none());
 415         assert!(guard.canceled.is_none());
 416     }
 417 }
 418
 419
 420 ////////////////////////////////////////////////////////////////////////////////
 421 // Buffer, a simple ring buffer backed by Vec<T>
 422 ////////////////////////////////////////////////////////////////////////////////
 423
 424 impl<T> Buffer<T> {
 425     fn enqueue(&mut self, t: T) {
 426         let pos = (self.start + self.size) % self.buf.len();
 427         self.size += 1;
 428         let prev = mem::replace(&mut self.buf[pos], Some(t));
 429         assert!(prev.is_none());
 430     }
 431
 432     fn dequeue(&mut self) -> T {
 433         let start = self.start;
 434         self.size -= 1;
 435         self.start = (self.start + 1) % self.buf.len();
 436         let result = &mut self.buf[start];
 437         result.take().unwrap()
 438     }
 439
 440     fn size(&self) -> usize { self.size }
 441     fn cap(&self) -> usize { self.buf.len() }
 442 }
 443
 444 ////////////////////////////////////////////////////////////////////////////////
 445 // Queue, a simple queue to enqueue threads with (stack-allocated nodes)
 446 ////////////////////////////////////////////////////////////////////////////////
 447
 448 impl Queue {
 449     fn enqueue(&mut self, node: &mut Node) -> WaitToken {
 450         let (wait_token, signal_token) = blocking::tokens();
 451         node.token = Some(signal_token);
 452         node.next = ptr::null_mut();
 453
 454         if self.tail.is_null() {
 455             self.head = node as *mut Node;
 456             self.tail = node as *mut Node;
 457         } else {
 458             unsafe {
 459                 (*self.tail).next = node as *mut Node;
 460                 self.tail = node as *mut Node;
 461             }
 462         }
 463
 464         wait_token
 465     }
 466
 467     fn dequeue(&mut self) -> Option<SignalToken> {
 468         if self.head.is_null() {
 469             return None
 470         }
 471         let node = self.head;
 472         self.head = unsafe { (*node).next };
 473         if self.head.is_null() {
 474             self.tail = ptr::null_mut();
 475         }
 476         unsafe {
 477             (*node).next = ptr::null_mut();
 478             Some((*node).token.take().unwrap())
 479         }
 480     }
 481 }