src/libstd/sync/mpsc/sync.rs

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