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