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 fmt;
  68 use marker;
  69 use ptr;
  70 use sync::atomic::{AtomicUsize, AtomicBool, Ordering};
  71 use thread::{self, Thread};
  72
  73 /// A synchronization primitive which can be used to run a one-time global
  74 /// initialization. Useful for one-time initialization for FFI or related
  75 /// functionality. This type can only be constructed with the [`ONCE_INIT`]
  76 /// value.
  77 ///
  78 /// [`ONCE_INIT`]: constant.ONCE_INIT.html
  79 ///
  80 /// # Examples
  81 ///
  82 /// ```
  83 /// use std::sync::{Once, ONCE_INIT};
  84 ///
  85 /// static START: Once = ONCE_INIT;
  86 ///
  87 /// START.call_once(|| {
  88 ///     // run initialization here
  89 /// });
  90 /// ```
  91 #[stable(feature = "rust1", since = "1.0.0")]
  92 pub struct Once {
  93     // This `state` word is actually an encoded version of just a pointer to a
  94     // `Waiter`, so we add the `PhantomData` appropriately.
  95     state: AtomicUsize,
  96     _marker: marker::PhantomData<*mut Waiter>,
  97 }
  98
  99 // The `PhantomData` of a raw pointer removes these two auto traits, but we
 100 // enforce both below in the implementation so this should be safe to add.
 101 #[stable(feature = "rust1", since = "1.0.0")]
 102 unsafe impl Sync for Once {}
 103 #[stable(feature = "rust1", since = "1.0.0")]
 104 unsafe impl Send for Once {}
 105
 106 /// State yielded to the [`call_once_force`] method which can be used to query
 107 /// whether the [`Once`] was previously poisoned or not.
 108 ///
 109 /// [`call_once_force`]: struct.Once.html#method.call_once_force
 110 /// [`Once`]: struct.Once.html
 111 #[unstable(feature = "once_poison", issue = "33577")]
 112 #[derive(Debug)]
 113 pub struct OnceState {
 114     poisoned: bool,
 115 }
 116
 117 /// Initialization value for static [`Once`] values.
 118 ///
 119 /// [`Once`]: struct.Once.html
 120 ///
 121 /// # Examples
 122 ///
 123 /// ```
 124 /// use std::sync::{Once, ONCE_INIT};
 125 ///
 126 /// static START: Once = ONCE_INIT;
 127 /// ```
 128 #[stable(feature = "rust1", since = "1.0.0")]
 129 pub const ONCE_INIT: Once = Once::new();
 130
 131 // Four states that a Once can be in, encoded into the lower bits of `state` in
 132 // the Once structure.
 133 const INCOMPLETE: usize = 0x0;
 134 const POISONED: usize = 0x1;
 135 const RUNNING: usize = 0x2;
 136 const COMPLETE: usize = 0x3;
 137
 138 // Mask to learn about the state. All other bits are the queue of waiters if
 139 // this is in the RUNNING state.
 140 const STATE_MASK: usize = 0x3;
 141
 142 // Representation of a node in the linked list of waiters in the RUNNING state.
 143 struct Waiter {
 144     thread: Option<Thread>,
 145     signaled: AtomicBool,
 146     next: *mut Waiter,
 147 }
 148
 149 // Helper struct used to clean up after a closure call with a `Drop`
 150 // implementation to also run on panic.
 151 struct Finish {
 152     panicked: bool,
 153     me: &'static Once,
 154 }
 155
 156 impl Once {
 157     /// Creates a new `Once` value.
 158     #[stable(feature = "once_new", since = "1.2.0")]
 159     pub const fn new() -> Once {
 160         Once {
 161             state: AtomicUsize::new(INCOMPLETE),
 162             _marker: marker::PhantomData,
 163         }
 164     }
 165
 166     /// Performs an initialization routine once and only once. The given closure
 167     /// will be executed if this is the first time `call_once` has been called,
 168     /// and otherwise the routine will *not* be invoked.
 169     ///
 170     /// This method will block the calling thread if another initialization
 171     /// routine is currently running.
 172     ///
 173     /// When this function returns, it is guaranteed that some initialization
 174     /// has run and completed (it may not be the closure specified). It is also
 175     /// guaranteed that any memory writes performed by the executed closure can
 176     /// be reliably observed by other threads at this point (there is a
 177     /// happens-before relation between the closure and code executing after the
 178     /// return).
 179     ///
 180     /// # Examples
 181     ///
 182     /// ```
 183     /// use std::sync::{Once, ONCE_INIT};
 184     ///
 185     /// static mut VAL: usize = 0;
 186     /// static INIT: Once = ONCE_INIT;
 187     ///
 188     /// // Accessing a `static mut` is unsafe much of the time, but if we do so
 189     /// // in a synchronized fashion (e.g. write once or read all) then we're
 190     /// // good to go!
 191     /// //
 192     /// // This function will only call `expensive_computation` once, and will
 193     /// // otherwise always return the value returned from the first invocation.
 194     /// fn get_cached_val() -> usize {
 195     ///     unsafe {
 196     ///         INIT.call_once(|| {
 197     ///             VAL = expensive_computation();
 198     ///         });
 199     ///         VAL
 200     ///     }
 201     /// }
 202     ///
 203     /// fn expensive_computation() -> usize {
 204     ///     // ...
 205     /// # 2
 206     /// }
 207     /// ```
 208     ///
 209     /// # Panics
 210     ///
 211     /// The closure `f` will only be executed once if this is called
 212     /// concurrently amongst many threads. If that closure panics, however, then
 213     /// it will *poison* this `Once` instance, causing all future invocations of
 214     /// `call_once` to also panic.
 215     ///
 216     /// This is similar to [poisoning with mutexes][poison].
 217     ///
 218     /// [poison]: struct.Mutex.html#poisoning
 219     #[stable(feature = "rust1", since = "1.0.0")]
 220     pub fn call_once<F>(&'static self, f: F) where F: FnOnce() {
 221         // Fast path, just see if we've completed initialization.
 222         if self.state.load(Ordering::SeqCst) == COMPLETE {
 223             return
 224         }
 225
 226         let mut f = Some(f);
 227         self.call_inner(false, &mut |_| f.take().unwrap()());
 228     }
 229
 230     /// Performs the same function as [`call_once`] except ignores poisoning.
 231     ///
 232     /// [`call_once`]: struct.Once.html#method.call_once
 233     ///
 234     /// If this `Once` has been poisoned (some initialization panicked) then
 235     /// this function will continue to attempt to call initialization functions
 236     /// until one of them doesn't panic.
 237     ///
 238     /// The closure `f` is yielded a [`OnceState`] structure which can be used to query the
 239     /// state of this `Once` (whether initialization has previously panicked or
 240     /// not).
 241     ///
 242     /// [`OnceState`]: struct.OnceState.html
 243     #[unstable(feature = "once_poison", issue = "33577")]
 244     pub fn call_once_force<F>(&'static self, f: F) where F: FnOnce(&OnceState) {
 245         // same as above, just with a different parameter to `call_inner`.
 246         if self.state.load(Ordering::SeqCst) == COMPLETE {
 247             return
 248         }
 249
 250         let mut f = Some(f);
 251         self.call_inner(true, &mut |p| {
 252             f.take().unwrap()(&OnceState { poisoned: p })
 253         });
 254     }
 255
 256     // This is a non-generic function to reduce the monomorphization cost of
 257     // using `call_once` (this isn't exactly a trivial or small implementation).
 258     //
 259     // Additionally, this is tagged with `#[cold]` as it should indeed be cold
 260     // and it helps let LLVM know that calls to this function should be off the
 261     // fast path. Essentially, this should help generate more straight line code
 262     // in LLVM.
 263     //
 264     // Finally, this takes an `FnMut` instead of a `FnOnce` because there's
 265     // currently no way to take an `FnOnce` and call it via virtual dispatch
 266     // without some allocation overhead.
 267     #[cold]
 268     fn call_inner(&'static self,
 269                   ignore_poisoning: bool,
 270                   init: &mut FnMut(bool)) {
 271         let mut state = self.state.load(Ordering::SeqCst);
 272
 273         'outer: loop {
 274             match state {
 275                 // If we're complete, then there's nothing to do, we just
 276                 // jettison out as we shouldn't run the closure.
 277                 COMPLETE => return,
 278
 279                 // If we're poisoned and we're not in a mode to ignore
 280                 // poisoning, then we panic here to propagate the poison.
 281                 POISONED if !ignore_poisoning => {
 282                     panic!("Once instance has previously been poisoned");
 283                 }
 284
 285                 // Otherwise if we see a poisoned or otherwise incomplete state
 286                 // we will attempt to move ourselves into the RUNNING state. If
 287                 // we succeed, then the queue of waiters starts at null (all 0
 288                 // bits).
 289                 POISONED |
 290                 INCOMPLETE => {
 291                     let old = self.state.compare_and_swap(state, RUNNING,
 292                                                           Ordering::SeqCst);
 293                     if old != state {
 294                         state = old;
 295                         continue
 296                     }
 297
 298                     // Run the initialization routine, letting it know if we're
 299                     // poisoned or not. The `Finish` struct is then dropped, and
 300                     // the `Drop` implementation here is responsible for waking
 301                     // up other waiters both in the normal return and panicking
 302                     // case.
 303                     let mut complete = Finish {
 304                         panicked: true,
 305                         me: self,
 306                     };
 307                     init(state == POISONED);
 308                     complete.panicked = false;
 309                     return
 310                 }
 311
 312                 // All other values we find should correspond to the RUNNING
 313                 // state with an encoded waiter list in the more significant
 314                 // bits. We attempt to enqueue ourselves by moving us to the
 315                 // head of the list and bail out if we ever see a state that's
 316                 // not RUNNING.
 317                 _ => {
 318                     assert!(state & STATE_MASK == RUNNING);
 319                     let mut node = Waiter {
 320                         thread: Some(thread::current()),
 321                         signaled: AtomicBool::new(false),
 322                         next: ptr::null_mut(),
 323                     };
 324                     let me = &mut node as *mut Waiter as usize;
 325                     assert!(me & STATE_MASK == 0);
 326
 327                     while state & STATE_MASK == RUNNING {
 328                         node.next = (state & !STATE_MASK) as *mut Waiter;
 329                         let old = self.state.compare_and_swap(state,
 330                                                               me | RUNNING,
 331                                                               Ordering::SeqCst);
 332                         if old != state {
 333                             state = old;
 334                             continue
 335                         }
 336
 337                         // Once we've enqueued ourselves, wait in a loop.
 338                         // Afterwards reload the state and continue with what we
 339                         // were doing from before.
 340                         while !node.signaled.load(Ordering::SeqCst) {
 341                             thread::park();
 342                         }
 343                         state = self.state.load(Ordering::SeqCst);
 344                         continue 'outer
 345                     }
 346                 }
 347             }
 348         }
 349     }
 350 }
 351
 352 #[stable(feature = "std_debug", since = "1.16.0")]
 353 impl fmt::Debug for Once {
 354     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 355         f.pad("Once { .. }")
 356     }
 357 }
 358
 359 impl Drop for Finish {
 360     fn drop(&mut self) {
 361         // Swap out our state with however we finished. We should only ever see
 362         // an old state which was RUNNING.
 363         let queue = if self.panicked {
 364             self.me.state.swap(POISONED, Ordering::SeqCst)
 365         } else {
 366             self.me.state.swap(COMPLETE, Ordering::SeqCst)
 367         };
 368         assert_eq!(queue & STATE_MASK, RUNNING);
 369
 370         // Decode the RUNNING to a list of waiters, then walk that entire list
 371         // and wake them up. Note that it is crucial that after we store `true`
 372         // in the node it can be free'd! As a result we load the `thread` to
 373         // signal ahead of time and then unpark it after the store.
 374         unsafe {
 375             let mut queue = (queue & !STATE_MASK) as *mut Waiter;
 376             while !queue.is_null() {
 377                 let next = (*queue).next;
 378                 let thread = (*queue).thread.take().unwrap();
 379                 (*queue).signaled.store(true, Ordering::SeqCst);
 380                 thread.unpark();
 381                 queue = next;
 382             }
 383         }
 384     }
 385 }
 386
 387 impl OnceState {
 388     /// Returns whether the associated [`Once`] has been poisoned.
 389     ///
 390     /// Once an initialization routine for a [`Once`] has panicked it will forever
 391     /// indicate to future forced initialization routines that it is poisoned.
 392     ///
 393     /// [`Once`]: struct.Once.html
 394     #[unstable(feature = "once_poison", issue = "33577")]
 395     pub fn poisoned(&self) -> bool {
 396         self.poisoned
 397     }
 398 }
 399
 400 #[cfg(all(test, not(target_os = "emscripten")))]
 401 mod tests {
 402     use panic;
 403     use sync::mpsc::channel;
 404     use thread;
 405     use super::Once;
 406
 407     #[test]
 408     fn smoke_once() {
 409         static O: Once = Once::new();
 410         let mut a = 0;
 411         O.call_once(|| a += 1);
 412         assert_eq!(a, 1);
 413         O.call_once(|| a += 1);
 414         assert_eq!(a, 1);
 415     }
 416
 417     #[test]
 418     fn stampede_once() {
 419         static O: Once = Once::new();
 420         static mut RUN: bool = false;
 421
 422         let (tx, rx) = channel();
 423         for _ in 0..10 {
 424             let tx = tx.clone();
 425             thread::spawn(move|| {
 426                 for _ in 0..4 { thread::yield_now() }
 427                 unsafe {
 428                     O.call_once(|| {
 429                         assert!(!RUN);
 430                         RUN = true;
 431                     });
 432                     assert!(RUN);
 433                 }
 434                 tx.send(()).unwrap();
 435             });
 436         }
 437
 438         unsafe {
 439             O.call_once(|| {
 440                 assert!(!RUN);
 441                 RUN = true;
 442             });
 443             assert!(RUN);
 444         }
 445
 446         for _ in 0..10 {
 447             rx.recv().unwrap();
 448         }
 449     }
 450
 451     #[test]
 452     fn poison_bad() {
 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         // poisoning propagates
 462         let t = panic::catch_unwind(|| {
 463             O.call_once(|| {});
 464         });
 465         assert!(t.is_err());
 466
 467         // we can subvert poisoning, however
 468         let mut called = false;
 469         O.call_once_force(|p| {
 470             called = true;
 471             assert!(p.poisoned())
 472         });
 473         assert!(called);
 474
 475         // once any success happens, we stop propagating the poison
 476         O.call_once(|| {});
 477     }
 478
 479     #[test]
 480     fn wait_for_force_to_finish() {
 481         static O: Once = Once::new();
 482
 483         // poison the once
 484         let t = panic::catch_unwind(|| {
 485             O.call_once(|| panic!());
 486         });
 487         assert!(t.is_err());
 488
 489         // make sure someone's waiting inside the once via a force
 490         let (tx1, rx1) = channel();
 491         let (tx2, rx2) = channel();
 492         let t1 = thread::spawn(move || {
 493             O.call_once_force(|p| {
 494                 assert!(p.poisoned());
 495                 tx1.send(()).unwrap();
 496                 rx2.recv().unwrap();
 497             });
 498         });
 499
 500         rx1.recv().unwrap();
 501
 502         // put another waiter on the once
 503         let t2 = thread::spawn(|| {
 504             let mut called = false;
 505             O.call_once(|| {
 506                 called = true;
 507             });
 508             assert!(!called);
 509         });
 510
 511         tx2.send(()).unwrap();
 512
 513         assert!(t1.join().is_ok());
 514         assert!(t2.join().is_ok());
 515
 516     }
 517 }