src/libstd/sync/once.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 //! A "once initialization" primitive
  12 //!
  13 //! This primitive is meant to be used to run one-time initialization. An
  14 //! example use case would be for initializing an FFI library.
  15
  16 // A "once" is a relatively simple primitive, and it's also typically provided
  17 // by the OS as well (see `pthread_once` or `InitOnceExecuteOnce`). The OS
  18 // primitives, however, tend to have surprising restrictions, such as the Unix
  19 // one doesn't allow an argument to be passed to the function.
  20 //
  21 // As a result, we end up implementing it ourselves in the standard library.
  22 // This also gives us the opportunity to optimize the implementation a bit which
  23 // should help the fast path on call sites. Consequently, let's explain how this
  24 // primitive works now!
  25 //
  26 // So to recap, the guarantees of a Once are that it will call the
  27 // initialization closure at most once, and it will never return until the one
  28 // that's running has finished running. This means that we need some form of
  29 // blocking here while the custom callback is running at the very least.
  30 // Additionally, we add on the restriction of **poisoning**. Whenever an
  31 // initialization closure panics, the Once enters a "poisoned" state which means
  32 // that all future calls will immediately panic as well.
  33 //
  34 // So to implement this, one might first reach for a `StaticMutex`, but those
  35 // unfortunately need to be deallocated (e.g. call `destroy()`) to free memory
  36 // on all OSes (some of the BSDs allocate memory for mutexes). It also gets a
  37 // lot harder with poisoning to figure out when the mutex needs to be
  38 // deallocated because it's not after the closure finishes, but after the first
  39 // successful closure finishes.
  40 //
  41 // All in all, this is instead implemented with atomics and lock-free
  42 // operations! Whee! Each `Once` has one word of atomic state, and this state is
  43 // CAS'd on to determine what to do. There are four possible state of a `Once`:
  44 //
  45 // * Incomplete - no initialization has run yet, and no thread is currently
  46 //                using the Once.
  47 // * Poisoned - some thread has previously attempted to initialize the Once, but
  48 //              it panicked, so the Once is now poisoned. There are no other
  49 //              threads currently accessing this Once.
  50 // * Running - some thread is currently attempting to run initialization. It may
  51 //             succeed, so all future threads need to wait for it to finish.
  52 //             Note that this state is accompanied with a payload, described
  53 //             below.
  54 // * Complete - initialization has completed and all future calls should finish
  55 //              immediately.
  56 //
  57 // With 4 states we need 2 bits to encode this, and we use the remaining bits
  58 // in the word we have allocated as a queue of threads waiting for the thread
  59 // responsible for entering the RUNNING state. This queue is just a linked list
  60 // of Waiter nodes which is monotonically increasing in size. Each node is
  61 // allocated on the stack, and whenever the running closure finishes it will
  62 // consume the entire queue and notify all waiters they should try again.
  63 //
  64 // You'll find a few more details in the implementation, but that's the gist of
  65 // it!
  66
  67 use marker;
  68 use ptr;
  69 use sync::atomic::{AtomicUsize, AtomicBool, Ordering};
  70 use thread::{self, Thread};
  71
  72 /// A synchronization primitive which can be used to run a one-time global
  73 /// initialization. Useful for one-time initialization for FFI or related
  74 /// functionality. This type can only be constructed with the `ONCE_INIT`
  75 /// value.
  76 ///
  77 /// # Examples
  78 ///
  79 /// ```
  80 /// use std::sync::{Once, ONCE_INIT};
  81 ///
  82 /// static START: Once = ONCE_INIT;
  83 ///
  84 /// START.call_once(|| {
  85 ///     // run initialization here
  86 /// });
  87 /// ```
  88 #[stable(feature = "rust1", since = "1.0.0")]
  89 pub struct Once {
  90     // This `state` word is actually an encoded version of just a pointer to a
  91     // `Waiter`, so we add the `PhantomData` appropriately.
  92     state: AtomicUsize,
  93     _marker: marker::PhantomData<*mut Waiter>,
  94 }
  95
  96 // The `PhantomData` of a raw pointer removes these two auto traits, but we
  97 // enforce both below in the implementation so this should be safe to add.
  98 #[stable(feature = "rust1", since = "1.0.0")]
  99 unsafe impl Sync for Once {}
 100 #[stable(feature = "rust1", since = "1.0.0")]
 101 unsafe impl Send for Once {}
 102
 103 /// State yielded to the `call_once_force` method which can be used to query
 104 /// whether the `Once` was previously poisoned or not.
 105 #[unstable(feature = "once_poison", issue = "33577")]
 106 pub struct OnceState {
 107     poisoned: bool,
 108 }
 109
 110 /// Initialization value for static `Once` values.
 111 #[stable(feature = "rust1", since = "1.0.0")]
 112 pub const ONCE_INIT: Once = Once::new();
 113
 114 // Four states that a Once can be in, encoded into the lower bits of `state` in
 115 // the Once structure.
 116 const INCOMPLETE: usize = 0x0;
 117 const POISONED: usize = 0x1;
 118 const RUNNING: usize = 0x2;
 119 const COMPLETE: usize = 0x3;
 120
 121 // Mask to learn about the state. All other bits are the queue of waiters if
 122 // this is in the RUNNING state.
 123 const STATE_MASK: usize = 0x3;
 124
 125 // Representation of a node in the linked list of waiters in the RUNNING state.
 126 struct Waiter {
 127     thread: Option<Thread>,
 128     signaled: AtomicBool,
 129     next: *mut Waiter,
 130 }
 131
 132 // Helper struct used to clean up after a closure call with a `Drop`
 133 // implementation to also run on panic.
 134 struct Finish {
 135     panicked: bool,
 136     me: &'static Once,
 137 }
 138
 139 impl Once {
 140     /// Creates a new `Once` value.
 141     #[stable(feature = "once_new", since = "1.2.0")]
 142     pub const fn new() -> Once {
 143         Once {
 144             state: AtomicUsize::new(INCOMPLETE),
 145             _marker: marker::PhantomData,
 146         }
 147     }
 148
 149     /// Performs an initialization routine once and only once. The given closure
 150     /// will be executed if this is the first time `call_once` has been called,
 151     /// and otherwise the routine will *not* be invoked.
 152     ///
 153     /// This method will block the calling thread if another initialization
 154     /// routine is currently running.
 155     ///
 156     /// When this function returns, it is guaranteed that some initialization
 157     /// has run and completed (it may not be the closure specified). It is also
 158     /// guaranteed that any memory writes performed by the executed closure can
 159     /// be reliably observed by other threads at this point (there is a
 160     /// happens-before relation between the closure and code executing after the
 161     /// return).
 162     ///
 163     /// # Examples
 164     ///
 165     /// ```
 166     /// use std::sync::{Once, ONCE_INIT};
 167     ///
 168     /// static mut VAL: usize = 0;
 169     /// static INIT: Once = ONCE_INIT;
 170     ///
 171     /// // Accessing a `static mut` is unsafe much of the time, but if we do so
 172     /// // in a synchronized fashion (e.g. write once or read all) then we're
 173     /// // good to go!
 174     /// //
 175     /// // This function will only call `expensive_computation` once, and will
 176     /// // otherwise always return the value returned from the first invocation.
 177     /// fn get_cached_val() -> usize {
 178     ///     unsafe {
 179     ///         INIT.call_once(|| {
 180     ///             VAL = expensive_computation();
 181     ///         });
 182     ///         VAL
 183     ///     }
 184     /// }
 185     ///
 186     /// fn expensive_computation() -> usize {
 187     ///     // ...
 188     /// # 2
 189     /// }
 190     /// ```
 191     ///
 192     /// # Panics
 193     ///
 194     /// The closure `f` will only be executed once if this is called
 195     /// concurrently amongst many threads. If that closure panics, however, then
 196     /// it will *poison* this `Once` instance, causing all future invocations of
 197     /// `call_once` to also panic.
 198     ///
 199     /// This is similar to [poisoning with mutexes][poison].
 200     ///
 201     /// [poison]: struct.Mutex.html#poisoning
 202     #[stable(feature = "rust1", since = "1.0.0")]
 203     pub fn call_once<F>(&'static self, f: F) where F: FnOnce() {
 204         // Fast path, just see if we've completed initialization.
 205         if self.state.load(Ordering::SeqCst) == COMPLETE {
 206             return
 207         }
 208
 209         let mut f = Some(f);
 210         self.call_inner(false, &mut |_| f.take().unwrap()());
 211     }
 212
 213     /// Performs the same function as `call_once` except ignores poisoning.
 214     ///
 215     /// If this `Once` has been poisoned (some initialization panicked) then
 216     /// this function will continue to attempt to call initialization functions
 217     /// until one of them doesn't panic.
 218     ///
 219     /// The closure `f` is yielded a structure which can be used to query the
 220     /// state of this `Once` (whether initialization has previously panicked or
 221     /// not).
 222     #[unstable(feature = "once_poison", issue = "33577")]
 223     pub fn call_once_force<F>(&'static self, f: F) where F: FnOnce(&OnceState) {
 224         // same as above, just with a different parameter to `call_inner`.
 225         if self.state.load(Ordering::SeqCst) == COMPLETE {
 226             return
 227         }
 228
 229         let mut f = Some(f);
 230         self.call_inner(true, &mut |p| {
 231             f.take().unwrap()(&OnceState { poisoned: p })
 232         });
 233     }
 234
 235     // This is a non-generic function to reduce the monomorphization cost of
 236     // using `call_once` (this isn't exactly a trivial or small implementation).
 237     //
 238     // Additionally, this is tagged with `#[cold]` as it should indeed be cold
 239     // and it helps let LLVM know that calls to this function should be off the
 240     // fast path. Essentially, this should help generate more straight line code
 241     // in LLVM.
 242     //
 243     // Finally, this takes an `FnMut` instead of a `FnOnce` because there's
 244     // currently no way to take an `FnOnce` and call it via virtual dispatch
 245     // without some allocation overhead.
 246     #[cold]
 247     fn call_inner(&'static self,
 248                   ignore_poisoning: bool,
 249                   mut init: &mut FnMut(bool)) {
 250         let mut state = self.state.load(Ordering::SeqCst);
 251
 252         'outer: loop {
 253             match state {
 254                 // If we're complete, then there's nothing to do, we just
 255                 // jettison out as we shouldn't run the closure.
 256                 COMPLETE => return,
 257
 258                 // If we're poisoned and we're not in a mode to ignore
 259                 // poisoning, then we panic here to propagate the poison.
 260                 POISONED if !ignore_poisoning => {
 261                     panic!("Once instance has previously been poisoned");
 262                 }
 263
 264                 // Otherwise if we see a poisoned or otherwise incomplete state
 265                 // we will attempt to move ourselves into the RUNNING state. If
 266                 // we succeed, then the queue of waiters starts at null (all 0
 267                 // bits).
 268                 POISONED |
 269                 INCOMPLETE => {
 270                     let old = self.state.compare_and_swap(state, RUNNING,
 271                                                           Ordering::SeqCst);
 272                     if old != state {
 273                         state = old;
 274                         continue
 275                     }
 276
 277                     // Run the initialization routine, letting it know if we're
 278                     // poisoned or not. The `Finish` struct is then dropped, and
 279                     // the `Drop` implementation here is responsible for waking
 280                     // up other waiters both in the normal return and panicking
 281                     // case.
 282                     let mut complete = Finish {
 283                         panicked: true,
 284                         me: self,
 285                     };
 286                     init(state == POISONED);
 287                     complete.panicked = false;
 288                     return
 289                 }
 290
 291                 // All other values we find should correspond to the RUNNING
 292                 // state with an encoded waiter list in the more significant
 293                 // bits. We attempt to enqueue ourselves by moving us to the
 294                 // head of the list and bail out if we ever see a state that's
 295                 // not RUNNING.
 296                 _ => {
 297                     assert!(state & STATE_MASK == RUNNING);
 298                     let mut node = Waiter {
 299                         thread: Some(thread::current()),
 300                         signaled: AtomicBool::new(false),
 301                         next: ptr::null_mut(),
 302                     };
 303                     let me = &mut node as *mut Waiter as usize;
 304                     assert!(me & STATE_MASK == 0);
 305
 306                     while state & STATE_MASK == RUNNING {
 307                         node.next = (state & !STATE_MASK) as *mut Waiter;
 308                         let old = self.state.compare_and_swap(state,
 309                                                               me | RUNNING,
 310                                                               Ordering::SeqCst);
 311                         if old != state {
 312                             state = old;
 313                             continue
 314                         }
 315
 316                         // Once we've enqueued ourselves, wait in a loop.
 317                         // Aftewards reload the state and continue with what we
 318                         // were doing from before.
 319                         while !node.signaled.load(Ordering::SeqCst) {
 320                             thread::park();
 321                         }
 322                         state = self.state.load(Ordering::SeqCst);
 323                         continue 'outer
 324                     }
 325                 }
 326             }
 327         }
 328     }
 329 }
 330
 331 impl Drop for Finish {
 332     fn drop(&mut self) {
 333         // Swap out our state with however we finished. We should only ever see
 334         // an old state which was RUNNING.
 335         let queue = if self.panicked {
 336             self.me.state.swap(POISONED, Ordering::SeqCst)
 337         } else {
 338             self.me.state.swap(COMPLETE, Ordering::SeqCst)
 339         };
 340         assert_eq!(queue & STATE_MASK, RUNNING);
 341
 342         // Decode the RUNNING to a list of waiters, then walk that entire list
 343         // and wake them up. Note that it is crucial that after we store `true`
 344         // in the node it can be free'd! As a result we load the `thread` to
 345         // signal ahead of time and then unpark it after the store.
 346         unsafe {
 347             let mut queue = (queue & !STATE_MASK) as *mut Waiter;
 348             while !queue.is_null() {
 349                 let next = (*queue).next;
 350                 let thread = (*queue).thread.take().unwrap();
 351                 (*queue).signaled.store(true, Ordering::SeqCst);
 352                 thread.unpark();
 353                 queue = next;
 354             }
 355         }
 356     }
 357 }
 358
 359 impl OnceState {
 360     /// Returns whether the associated `Once` has been poisoned.
 361     ///
 362     /// Once an initalization routine for a `Once` has panicked it will forever
 363     /// indicate to future forced initialization routines that it is poisoned.
 364     #[unstable(feature = "once_poison", issue = "33577")]
 365     pub fn poisoned(&self) -> bool {
 366         self.poisoned
 367     }
 368 }
 369
 370 #[cfg(test)]
 371 mod tests {
 372     use prelude::v1::*;
 373
 374     use panic;
 375     use sync::mpsc::channel;
 376     use thread;
 377     use super::Once;
 378
 379     #[test]
 380     fn smoke_once() {
 381         static O: Once = Once::new();
 382         let mut a = 0;
 383         O.call_once(|| a += 1);
 384         assert_eq!(a, 1);
 385         O.call_once(|| a += 1);
 386         assert_eq!(a, 1);
 387     }
 388
 389     #[test]
 390     fn stampede_once() {
 391         static O: Once = Once::new();
 392         static mut run: bool = false;
 393
 394         let (tx, rx) = channel();
 395         for _ in 0..10 {
 396             let tx = tx.clone();
 397             thread::spawn(move|| {
 398                 for _ in 0..4 { thread::yield_now() }
 399                 unsafe {
 400                     O.call_once(|| {
 401                         assert!(!run);
 402                         run = true;
 403                     });
 404                     assert!(run);
 405                 }
 406                 tx.send(()).unwrap();
 407             });
 408         }
 409
 410         unsafe {
 411             O.call_once(|| {
 412                 assert!(!run);
 413                 run = true;
 414             });
 415             assert!(run);
 416         }
 417
 418         for _ in 0..10 {
 419             rx.recv().unwrap();
 420         }
 421     }
 422
 423     #[test]
 424     fn poison_bad() {
 425         static O: Once = Once::new();
 426
 427         // poison the once
 428         let t = panic::catch_unwind(|| {
 429             O.call_once(|| panic!());
 430         });
 431         assert!(t.is_err());
 432
 433         // poisoning propagates
 434         let t = panic::catch_unwind(|| {
 435             O.call_once(|| {});
 436         });
 437         assert!(t.is_err());
 438
 439         // we can subvert poisoning, however
 440         let mut called = false;
 441         O.call_once_force(|p| {
 442             called = true;
 443             assert!(p.poisoned())
 444         });
 445         assert!(called);
 446
 447         // once any success happens, we stop propagating the poison
 448         O.call_once(|| {});
 449     }
 450
 451     #[test]
 452     fn wait_for_force_to_finish() {
 453         static O: Once = Once::new();
 454
 455         // poison the once
 456         let t = panic::catch_unwind(|| {
 457             O.call_once(|| panic!());
 458         });
 459         assert!(t.is_err());
 460
 461         // make sure someone's waiting inside the once via a force
 462         let (tx1, rx1) = channel();
 463         let (tx2, rx2) = channel();
 464         let t1 = thread::spawn(move || {
 465             O.call_once_force(|p| {
 466                 assert!(p.poisoned());
 467                 tx1.send(()).unwrap();
 468                 rx2.recv().unwrap();
 469             });
 470         });
 471
 472         rx1.recv().unwrap();
 473
 474         // put another waiter on the once
 475         let t2 = thread::spawn(|| {
 476             let mut called = false;
 477             O.call_once(|| {
 478                 called = true;
 479             });
 480             assert!(!called);
 481         });
 482
 483         tx2.send(()).unwrap();
 484
 485         assert!(t1.join().is_ok());
 486         assert!(t2.join().is_ok());
 487
 488     }
 489 }